feat: add Turing (20~ Series) GPU compatibility for Z Image Turbo

examples/v1/z-image-turbo.py

@@ -2,18 +2,21 @@
 from diffusers.pipelines.z_image.pipeline_z_image import ZImagePipeline
 
 from nunchaku import NunchakuZImageTransformer2DModel
-from nunchaku.utils import get_precision
+from nunchaku.utils import get_precision, is_turing
 
 if __name__ == "__main__":
     precision = get_precision()  # auto-detect your precision is 'int4' or 'fp4' based on your GPU
     rank = 128  # Use 32 for faster sampling; 256 (INT4 only) for best quality
+    dtype = torch.float16 if is_turing() else torch.bfloat16  # Use float16 when Turing (20- series) GPU is used.
     transformer = NunchakuZImageTransformer2DModel.from_pretrained(
-        f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}-z-image-turbo.safetensors"
+        f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}-z-image-turbo.safetensors", torch_dtype=dtype
     )
 
     pipe = ZImagePipeline.from_pretrained(
-        "Tongyi-MAI/Z-Image-Turbo", transformer=transformer, torch_dtype=torch.bfloat16, low_cpu_mem_usage=False
-    ).to("cuda")
+        "Tongyi-MAI/Z-Image-Turbo", transformer=transformer, torch_dtype=dtype, low_cpu_mem_usage=False
+    )
+    pipe.enable_sequential_cpu_offload()  # enable sequential CPU offload for low vram
+    # pipe = pipe.to("cuda") # or else comment the line above and uncomment this line to put all components to GPU
 
     prompt = "a young military male cooking in the kitchen for therapy"
 
@@ -26,4 +29,4 @@
         generator=torch.Generator().manual_seed(12345),
     ).images[0]
 
-    image.save(f"z-image-turbo-{precision}_r{rank}.png")
+    image.save(f"z-image-turbo-{precision}_r{rank}_{str(dtype)}.png")

nunchaku/models/attention_processors/zimage.py

@@ -8,8 +8,6 @@
 from diffusers.models.attention_dispatch import dispatch_attention_fn
 from diffusers.models.transformers.transformer_z_image import ZSingleStreamAttnProcessor
 
-from ...ops.fused import fused_qkv_norm_rottary
-
 
 class NunchakuZSingleStreamAttnProcessor(ZSingleStreamAttnProcessor):
     """
@@ -33,13 +31,8 @@ def __call__(
         """
         Forward pass of the attention module. Adapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor#__call__.
         """
-        qkv = fused_qkv_norm_rottary(
-            hidden_states,
-            attn.to_qkv,
-            attn.norm_q,
-            attn.norm_k,
-            freqs_cis,
-        )
+        qkv = attn.fused_module(hidden_states, freqs_cis)
+
         query, key, value = qkv.chunk(3, dim=-1)
         query = query.unflatten(-1, (attn.heads, -1))
         key = key.unflatten(-1, (attn.heads, -1))

nunchaku/models/linear.py

@@ -149,11 +149,12 @@ def from_linear(cls, linear: nn.Linear, **kwargs):
         SVDQW4A4Linear
         """
         in_features = kwargs.pop("in_features", linear.in_features)
+        torch_dtype = kwargs.pop("torch_dtype", linear.weight.dtype)
         return cls(
             in_features=in_features,
             out_features=linear.out_features,
             bias=linear.bias is not None,
-            torch_dtype=linear.weight.dtype,
+            torch_dtype=torch_dtype,
             device=linear.weight.device,
             **kwargs,
         )

nunchaku/models/transformers/transformer_zimage.py

@@ -11,19 +11,22 @@
 import torch.nn as nn
 from diffusers.models.attention import FeedForward
 from diffusers.models.attention_processor import Attention
+from diffusers.models.normalization import RMSNorm
 from diffusers.models.transformers.transformer_z_image import FeedForward as ZImageFeedForward
 from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel, ZImageTransformerBlock
 from huggingface_hub import utils
 
 from nunchaku.models.unets.unet_sdxl import NunchakuSDXLFeedForward
 
+from ...ops.gemm import svdq_gemm_w4a4_cuda
+from ...ops.quantize import svdq_quantize_w4a4_act_fuse_lora_cuda
 from ...utils import get_precision, pad_tensor
 from ..attention import NunchakuBaseAttention
 from ..attention_processors.zimage import NunchakuZSingleStreamAttnProcessor
 from ..embeddings import pack_rotemb
 from ..linear import SVDQW4A4Linear
 from ..utils import fuse_linears
-from .utils import NunchakuModelLoaderMixin, patch_scale_key
+from .utils import NunchakuModelLoaderMixin, convert_fp16, patch_scale_key
 
 
 class NunchakuZImageRopeHook:
@@ -50,6 +53,77 @@ def __call__(self, module: nn.Module, input_args: tuple, input_kwargs: dict):
         return input_args, new_input_kwargs
 
 
+class NunchakuZImageFusedModule(nn.Module):
+    """
+    Fused module for quantized QKV projection, RMS normalization, and rotary embedding for ZImage attention.
+
+    Parameters
+    ----------
+    qkv : SVDQW4A4Linear
+        Quantized QKV projection layer.
+    norm_q : RMSNorm
+        RMSNorm for query.
+    norm_k : RMSNorm
+        RMSNorm for key.
+    """
+
+    def __init__(self, qkv: SVDQW4A4Linear, norm_q: RMSNorm, norm_k: RMSNorm):
+        super().__init__()
+        for name, param in qkv.named_parameters(prefix="qkv_"):
+            setattr(self, name.replace(".", ""), param)
+        self.qkv_precision = qkv.precision
+        self.qkv_out_features = qkv.out_features
+        for name, param in norm_q.named_parameters(prefix="norm_q_"):
+            setattr(self, name.replace(".", ""), param)
+        for name, param in norm_k.named_parameters(prefix="norm_k_"):
+            setattr(self, name.replace(".", ""), param)
+
+    def forward(self, x: torch.Tensor, freqs_cis: Optional[torch.Tensor] = None):
+        """
+        Fuse QKV projection, RMS normalizaion and rotary embedding.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The hidden states tensor
+        freqs_cis : torch.Tensor, optional
+            The rotary embedding tensor
+
+        Returns
+        -------
+        The projection results of q, k, v. q result and k result are RMS-normalized and applied RoPE.
+        """
+        batch_size, seq_len, channels = x.shape
+        x = x.view(batch_size * seq_len, channels)
+        quantized_x, ascales, lora_act_out = svdq_quantize_w4a4_act_fuse_lora_cuda(
+            x,
+            lora_down=self.qkv_proj_down,
+            smooth=self.qkv_smooth_factor,
+            fp4=self.qkv_precision == "nvfp4",
+            pad_size=256,
+        )
+        output = torch.empty(batch_size * seq_len, self.qkv_out_features, dtype=x.dtype, device=x.device)
+        svdq_gemm_w4a4_cuda(
+            act=quantized_x,
+            wgt=self.qkv_qweight,
+            out=output,
+            ascales=ascales,
+            wscales=self.qkv_wscales,
+            lora_act_in=lora_act_out,
+            lora_up=self.qkv_proj_up,
+            bias=getattr(self, "qkv_bias", None),
+            fp4=self.qkv_precision == "nvfp4",
+            alpha=1.0 if self.qkv_precision == "nvfp4" else None,
+            wcscales=self.qkv_wcscales if self.qkv_precision == "nvfp4" else None,
+            norm_q=self.norm_q_weight,
+            norm_k=self.norm_k_weight,
+            rotary_emb=freqs_cis,
+        )
+
+        output = output.view(batch_size, seq_len, -1)
+        return output
+
+
 class NunchakuZImageAttention(NunchakuBaseAttention):
     """
     Nunchaku-optimized Attention module for ZImage with quantized and fused QKV projections.
@@ -172,6 +246,14 @@ def _convert_z_image_ff(z_ff: ZImageFeedForward) -> FeedForward:
     return converted_ff
 
 
+def replace_fused_module(module, incompatible_keys):
+    assert isinstance(module, NunchakuZImageAttention)
+    module.fused_module = NunchakuZImageFusedModule(module.to_qkv, module.norm_q, module.norm_k)
+    del module.to_qkv
+    del module.norm_q
+    del module.norm_k
+
+
 class NunchakuZImageFeedForward(NunchakuSDXLFeedForward):
     """
     Quantized feed-forward block for :class:`NunchakuZImageTransformerBlock`.
@@ -218,6 +300,7 @@ def _patch_model(self, skip_refiners: bool = False, **kwargs):
         def _patch_transformer_block(block_list: List[ZImageTransformerBlock]):
             for _, block in enumerate(block_list):
                 block.attention = NunchakuZImageAttention(block.attention, **kwargs)
+                block.attention.register_load_state_dict_post_hook(replace_fused_module)
                 block.feed_forward = NunchakuZImageFeedForward(block.feed_forward, **kwargs)
 
         def _convert_feed_forward(block_list: List[ZImageTransformerBlock]):
@@ -323,10 +406,12 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str],
 
         print(f"quantization_config: {quantization_config}, rank={rank}, skip_refiners={skip_refiners}")
 
-        transformer._patch_model(skip_refiners=skip_refiners, precision=precision, rank=rank)
+        transformer._patch_model(skip_refiners=skip_refiners, precision=precision, rank=rank, **kwargs)
         transformer = transformer.to_empty(device=device)
 
         patch_scale_key(transformer, model_state_dict)
+        if torch_dtype == torch.float16:
+            convert_fp16(transformer, model_state_dict)
 
         transformer.load_state_dict(model_state_dict)
 

nunchaku/models/transformers/utils.py

@@ -164,10 +164,21 @@ def patch_scale_key(transformer_from_config: nn.Module, state_dict_from_checkpoi
         if k not in state_dict_from_checkpoint:
             assert ".wcscales" in k
             state_dict_from_checkpoint[k] = torch.ones_like(state_dict[k])
-        else:
-            assert state_dict[k].dtype == state_dict_from_checkpoint[k].dtype
 
     for n, m in transformer_from_config.named_modules():
         if isinstance(m, SVDQW4A4Linear):
             if m.wtscale is not None:
                 m.wtscale = state_dict_from_checkpoint.pop(f"{n}.wtscale", 1.0)
+
+
+def convert_fp16(transformer_from_config: nn.Module, state_dict_from_checkpoint: dict):
+    state_dict = transformer_from_config.state_dict()
+    for k in state_dict.keys():
+        if state_dict[k].dtype != state_dict_from_checkpoint[k].dtype:
+            assert (
+                state_dict[k].dtype == torch.float16 and state_dict_from_checkpoint[k].dtype == torch.bfloat16
+            ), f"Unexpected dtype difference for key: {k}, model dtype: {state_dict[k].dtype}, \
+                checkpoint dtype: {state_dict_from_checkpoint[k].dtype}"
+            state_dict_from_checkpoint[k] = torch.nan_to_num(
+                state_dict_from_checkpoint[k].to(torch.float16), nan=0.0, posinf=65504, neginf=-65504
+            )

tests/v1/z_image/test_z_image_turbo.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 import pytest
+import requests
 import torch
 from diffusers import ZImagePipeline
 
@@ -48,6 +49,7 @@
 ]
 
 
+@pytest.mark.skipif(is_turing(), reason="Turing GPUs. Skip tests.")
 @pytest.mark.parametrize(
     "rank,expected_lpips",
     [
@@ -110,3 +112,67 @@ def test_zimage_turbo(rank: int, expected_lpips: dict[str, float]):
     lpips = compute_lpips(save_dir_16bit, save_dir_nunchaku)
     print(f"lpips: {lpips}")
     assert lpips < expected_lpips[f"{precision}-{dtype_str}"] * 1.15
+
+
+def download_ref_images(local_save_dir, filenames):
+    for filename in filenames:
+        try:
+            url = f"https://huggingface.co/datasets/nunchaku-tech/test-data/resolve/main/inputs/test-ref-z-image-turbo-{filename}.png"
+            save_path = local_save_dir / f"{filename}.png"
+            response = requests.get(url, stream=True, timeout=10)
+            response.raise_for_status()
+            if not os.path.exists(local_save_dir):
+                os.makedirs(local_save_dir)
+            with open(save_path, "wb") as file:
+                for chunk in response.iter_content(chunk_size=2048):
+                    file.write(chunk)
+            print(f"ref image downloaded: url: {url}, save_path: {save_path}")
+        except Exception as e:
+            print(f"download ref image failed: {e}")
+
+
+@pytest.mark.parametrize(
+    "rank,expected_lpips",
+    [
+        (32, {"int4-fp16": 0.4}),
+        (128, {"int4-fp16": 0.38}),
+        (256, {"int4-fp16": 0.37}),
+    ],
+)
+def test_zimage_turbo_turing(rank: int, expected_lpips: dict[str, float]):
+    if f"{precision}-{dtype_str}" not in expected_lpips:
+        return
+
+    if not already_generate(save_dir_16bit, len(dataset)):
+        filenames = [d["filename"] for d in dataset]
+        download_ref_images(save_dir_16bit, filenames)
+
+    save_dir_nunchaku = (
+        Path("test_results") / "nunchaku" / model_name / f"{precision}_r{rank}-fp16" / f"{folder_name}-bs{batch_size}"
+    )
+    path = f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}-z-image-turbo.safetensors"
+    transformer = NunchakuZImageTransformer2DModel.from_pretrained(path, torch_dtype=torch_dtype)
+
+    pipe = ZImagePipeline.from_pretrained(repo_id, transformer=transformer, torch_dtype=torch_dtype)
+    pipe.enable_sequential_cpu_offload()
+
+    run_pipeline(
+        dataset=dataset,
+        batch_size=batch_size,
+        pipeline=pipe,
+        save_dir=save_dir_nunchaku,
+        forward_kwargs={
+            "width": width,
+            "height": height,
+            "num_inference_steps": num_inference_steps,
+            "guidance_scale": guidance_scale,
+        },
+    )
+    del transformer
+    del pipe
+    gc.collect()
+    torch.cuda.empty_cache()
+
+    lpips = compute_lpips(save_dir_16bit, save_dir_nunchaku)
+    print(f"lpips: {lpips}")
+    assert lpips < expected_lpips[f"{precision}-fp16"] * 1.15