diff --git a/backend/loader.py b/backend/loader.py index 63c2fa65..4e8463f4 100644 --- a/backend/loader.py +++ b/backend/loader.py @@ -231,6 +231,8 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p assert isinstance(state_dict, dict) and len(state_dict) > 16, "You do not have model state dict!" model_loader = None + _nz = False # Nunchaku Z-Image + if cls_name == "UNet2DConditionModel": model_loader = lambda c: IntegratedUNet2DConditionModel.from_config(c) elif cls_name == "FluxTransformer2DModel": @@ -261,13 +263,16 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p model_loader = lambda c: QwenImageTransformer2DModel(**c) elif cls_name in ("Lumina2Transformer2DModel", "ZImageTransformer2DModel"): if guess.nunchaku: - from backend.nn.svdq import NunchakuZImageModel + from backend.nn.svdq import patch_nunchaku_zimage - model_loader = lambda c: NunchakuZImageModel(**c) - else: - from backend.nn.lumina import NextDiT + guess.unet_config.pop("filename") + precision = guess.unet_config.pop("precision") + rank = guess.unet_config.pop("rank") + _nz = True - model_loader = lambda c: NextDiT(**c) + from backend.nn.lumina import NextDiT + + model_loader = lambda c: NextDiT(**c) unet_config = guess.unet_config.copy() state_dict_parameters = memory_management.state_dict_parameters(state_dict) @@ -300,9 +305,11 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p need_manual_cast = storage_dtype != computation_dtype to_args = dict(device=initial_device, dtype=storage_dtype) - with using_forge_operations(**to_args, manual_cast_enabled=need_manual_cast): + with using_forge_operations(operations=False if _nz else None, **to_args, manual_cast_enabled=need_manual_cast): model = model_loader(unet_config).to(**to_args) + if _nz: + model = patch_nunchaku_zimage(model, precision, rank) load_state_dict(model, state_dict) if hasattr(model, "_internal_dict"): diff --git a/backend/nn/svdq.py b/backend/nn/svdq.py index 7adfa20d..fa3b7beb 100644 --- a/backend/nn/svdq.py +++ b/backend/nn/svdq.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Callable, Optional if TYPE_CHECKING: from transformers import T5EncoderModel @@ -15,6 +15,7 @@ from nunchaku.caching.diffusers_adapters.flux import apply_cache_on_transformer from nunchaku.caching.fbcache import cache_context, create_cache_context from nunchaku.lora.flux.compose import compose_lora from nunchaku.models.linear import AWQW4A16Linear, SVDQW4A4Linear +from nunchaku.models.transformers.utils import patch_scale_key from nunchaku.models.utils import CPUOffloadManager from nunchaku.ops.fused import fused_gelu_mlp from nunchaku.utils import load_state_dict_in_safetensors @@ -779,6 +780,8 @@ class NunchakuQwenImageTransformer2DModel(NunchakuModelMixin, QwenImageTransform # region: Z-Image +from functools import wraps + from backend.nn.lumina import JointAttention, NextDiT, clamp_fp16 @@ -891,39 +894,26 @@ def patch_z_image_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, t return patched_state_dict -class NunchakuZImageModel(NextDiT): - """https://github.com/nunchaku-tech/ComfyUI-nunchaku/blob/dev/models/zimage.py#L101""" +def patch_nunchaku_zimage(model: NextDiT, precision: str, rank: int): + kwargs = {"precision": precision, "rank": rank} - def __init__(self, filename: str, precision: str, rank: int, *args, **kwargs): - super().__init__(*args, **kwargs) + def patch_transformer_block(block_list: list[torch.nn.Module]): + for _, block in enumerate(block_list): + block.attention = NunchakuZImageAttention(block.attention, **kwargs) + block.feed_forward = NunchakuZImageFeedForward(block.feed_forward, **kwargs) - _kwargs = {"precision": precision, "rank": rank} + patch_transformer_block(model.layers) + patch_transformer_block(model.noise_refiner) + patch_transformer_block(model.context_refiner) - def _patch_transformer_block(block_list: list[torch.nn.Module]): - for _, block in enumerate(block_list): - block.attention = NunchakuZImageAttention(block.attention, **_kwargs) - block.feed_forward = NunchakuZImageFeedForward(block.feed_forward, **_kwargs) + _load_state_dict: Callable = model.load_state_dict - _patch_transformer_block(self.layers) - _patch_transformer_block(self.noise_refiner) - _patch_transformer_block(self.context_refiner) - - def load_state_dict(self, sd, *args, **kwargs): + @wraps(_load_state_dict) + def load_state_dict(sd, *args, **kwargs): sd = patch_z_image_state_dict(sd) - return self._load_state_dict(sd, *args, **kwargs) + patch_scale_key(model, sd) + return _load_state_dict(sd, *args, **kwargs) - def _load_state_dict(self, sd, *args, **kwargs): - state_dict = self.state_dict() - for k in state_dict.keys(): - if k not in sd: - if "dummy" in k: - continue - if ".wcscales" not in k: - raise ValueError(f"Key {k} not found in state_dict") - sd[k] = torch.ones_like(state_dict[k]) - for n, m in self.named_modules(): - if isinstance(m, SVDQW4A4Linear): - if m.wtscale is not None: - m.wtscale = sd.pop(f"{n}.wtscale", 1.0) + model.load_state_dict = load_state_dict - return super().load_state_dict(sd, *args, **kwargs) + return model diff --git a/backend/operations.py b/backend/operations.py index 36f86638..0d629a5a 100644 --- a/backend/operations.py +++ b/backend/operations.py @@ -513,6 +513,20 @@ def using_forge_operations(operations=None, device=None, dtype=None, manual_cast current_device, current_dtype, current_manual_cast_enabled, current_bnb_dtype = device, dtype, manual_cast_enabled, bnb_dtype + if operations is False: + _dev = torch.get_default_device() + _dtype = torch.get_default_dtype() + + torch.set_default_device(current_device) + torch.set_default_dtype(current_dtype) + + yield + + torch.set_default_device(_dev) + torch.set_default_dtype(_dtype) + + return + if operations is None: if bnb_dtype in ["gguf"]: operations = ForgeOperationsGGUF