Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Fix using VAE in quantization mode #1090

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix using VAE in quantization mode #1090

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 9 additions & 9 deletions onediff_comfy_nodes/modules/booster_cache.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,6 @@
from comfy.model_patcher import ModelPatcher
from comfy.sd import VAE
from onediff.torch_utils.module_operations import get_sub_module
from onediff.utils.import_utils import is_oneflow_available

from .._config import is_disable_oneflow_backend


@singledispatch
Expand Down Expand Up @@ -54,14 +51,17 @@ def _(model: ModelPatcher):

@get_cached_model.register
def _(model: VAE):
if is_oneflow_available() and not is_disable_oneflow_backend():
from .oneflow.utils.booster_utils import is_using_oneflow_backend
# from onediff.utils.import_utils import is_oneflow_available
# from .._config import is_disable_oneflow_backend
# if is_oneflow_available() and not is_disable_oneflow_backend():
# from .oneflow.utils.booster_utils import is_using_oneflow_backend

if is_using_oneflow_backend(model):
return None
# if is_using_oneflow_backend(model):
# return None

# TODO(TEST) if support cache
return model.first_stage_model
# # TODO(TEST) if support cache
# # return model.first_stage_model
return None


class BoosterCacheService:
Expand Down
13 changes: 13 additions & 0 deletions onediff_comfy_nodes/modules/oneflow/booster_quantization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os
import warnings
from dataclasses import dataclass
from functools import partial, singledispatchmethod
from typing import Any, Dict, Optional, Union
Expand All @@ -7,6 +8,7 @@
import torch.nn as nn
from comfy.controlnet import ControlNet
from comfy.model_patcher import ModelPatcher
from comfy.sd import VAE
from onediff.infer_compiler import oneflow_compile
from onediff.infer_compiler.backends.oneflow import (
OneflowDeployableModule as DeployableModule,
Expand Down Expand Up @@ -125,6 +127,17 @@ def execute(self, model, ckpt_name=None, **kwargs):
def extract_torch_module(self, model):
raise NotImplementedError(f"{type(model)}")

@execute.register(VAE)
def _(self, model: VAE, **kwargs):
# TODO: VAE does not support quantization and patch compatibility
from .booster_basic import BasicOneFlowBoosterExecutor

warnings.warn(
"TODO: VAE does not support quantization and patch compatibility",
UserWarning,
)
return BasicOneFlowBoosterExecutor().execute(model, **kwargs)
Comment on lines +130 to +139
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add stacklevel to the warning for better traceability.

The warning about VAE not supporting quantization lacks a stacklevel argument, which could help identify where the warning originates.

Apply this diff to improve the warning:

 warnings.warn(
     "TODO: VAE does not support quantization and patch compatibility",
     UserWarning,
+    stacklevel=2
 )
Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
@execute.register(VAE)
def _(self, model: VAE, **kwargs):
# TODO: VAE does not support quantization and patch compatibility
from .booster_basic import BasicOneFlowBoosterExecutor
warnings.warn(
"TODO: VAE does not support quantization and patch compatibility",
UserWarning,
)
return BasicOneFlowBoosterExecutor().execute(model, **kwargs)
@execute.register(VAE)
def _(self, model: VAE, **kwargs):
# TODO: VAE does not support quantization and patch compatibility
from .booster_basic import BasicOneFlowBoosterExecutor
warnings.warn(
"TODO: VAE does not support quantization and patch compatibility",
UserWarning,
stacklevel=2
)
return BasicOneFlowBoosterExecutor().execute(model, **kwargs)
Tools
Ruff

135-135: No explicit stacklevel keyword argument found

(B028)


@execute.register(ModelPatcher)
@execute.register(ControlNet)
def _(
Expand Down
Loading