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

Author: devgdovg

examples/v1/z-image-turbo.py

@@ -2,17 +2,18 @@
 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
+        "Tongyi-MAI/Z-Image-Turbo", transformer=transformer, torch_dtype=dtype, low_cpu_mem_usage=False
     ).to("cuda")
 
     prompt = "a young military male cooking in the kitchen for therapy"
@@ -26,4 +27,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/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

@@ -21,7 +21,7 @@
 from ..attention_processors.zimage import NunchakuZSingleStreamAttnProcessor
 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 NunchakuZImageAttention(NunchakuBaseAttention):
@@ -259,10 +259,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
+            )