nvfp4...

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"):

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

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