From 8e7385f8b35e7396f643366292581619afc66467 Mon Sep 17 00:00:00 2001 From: dengdong Date: Fri, 5 Dec 2025 11:18:00 -0500 Subject: [PATCH 01/17] [FEAT] support for z-image --- examples/v1/z-image-turbo.py | 31 +++ nunchaku/__init__.py | 2 + nunchaku/merge_safetensors.py | 21 +- nunchaku/models/__init__.py | 2 + .../models/attention_processors/zimage.py | 76 ++++++ nunchaku/models/transformers/__init__.py | 2 + .../models/transformers/transformer_zimage.py | 258 ++++++++++++++++++ 7 files changed, 389 insertions(+), 3 deletions(-) create mode 100644 examples/v1/z-image-turbo.py create mode 100644 nunchaku/models/attention_processors/zimage.py create mode 100644 nunchaku/models/transformers/transformer_zimage.py diff --git a/examples/v1/z-image-turbo.py b/examples/v1/z-image-turbo.py new file mode 100644 index 00000000..80052cad --- /dev/null +++ b/examples/v1/z-image-turbo.py @@ -0,0 +1,31 @@ +import torch +from diffusers.pipelines.z_image.pipeline_z_image import ZImagePipeline + +from nunchaku import NunchakuZImageTransformer2DModel +from nunchaku.utils import get_precision + +if __name__ == "__main__": + precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU + transformer = NunchakuZImageTransformer2DModel.from_pretrained( + f"/PATH/TO/svdq-{precision}_r128-z-image-turbo.safetensors" + ) + + pipe = ZImagePipeline.from_pretrained( + "Tongyi-MAI/Z-Image-Turbo", + transformer=transformer, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=False, + ).to("cuda") + + prompt = "a young military male cooking in the kitchen for therapy" + + image = pipe( + prompt=prompt, + height=1024, + width=1024, + num_inference_steps=9, # This actually results in 8 DiT forwards + guidance_scale=0.0, # Guidance should be 0 for the Turbo models + generator=torch.Generator("cuda").manual_seed(12345), + ).images[0] + + image.save(f"tmp_imgs/z-image-turbo-{precision}.png") diff --git a/nunchaku/__init__.py b/nunchaku/__init__.py index 977619dd..cb5eb872 100644 --- a/nunchaku/__init__.py +++ b/nunchaku/__init__.py @@ -4,6 +4,7 @@ NunchakuQwenImageTransformer2DModel, NunchakuSanaTransformer2DModel, NunchakuT5EncoderModel, + NunchakuZImageTransformer2DModel, ) __all__ = [ @@ -12,4 +13,5 @@ "NunchakuT5EncoderModel", "NunchakuFluxTransformer2DModelV2", "NunchakuQwenImageTransformer2DModel", + "NunchakuZImageTransformer2DModel", ] diff --git a/nunchaku/merge_safetensors.py b/nunchaku/merge_safetensors.py index aa1def94..d48e6fb2 100644 --- a/nunchaku/merge_safetensors.py +++ b/nunchaku/merge_safetensors.py @@ -34,7 +34,7 @@ def merge_safetensors( - pretrained_model_name_or_path: str | os.PathLike[str], **kwargs + pretrained_model_name_or_path: str | os.PathLike[str], model_class: str, **kwargs ) -> tuple[dict[str, torch.Tensor], dict[str, str]]: """ Merge split safetensors model files into a single state dict and metadata. @@ -108,6 +108,9 @@ def merge_safetensors( state_dict.update(transformer_block_sd) rank = next((v.shape[1] for k, v in transformer_block_sd.items() if ".lora_down" in k), 32) + if "ZImage" in model_class: + rank = next((v.shape[1] for k, v in transformer_block_sd.items() if ".proj_down" in k), 32) + skip_refiners = not any(("refiner" in k and "attention.to_qkv" in k) for k in transformer_block_sd.keys()) precision = "int4" for v in state_dict.values(): @@ -134,10 +137,12 @@ def merge_safetensors( }, "rank": rank, } + if "ZImage" in model_class: + quantization_config["skip_refiners"] = skip_refiners return state_dict, { "config": Path(config_path).read_text(), "comfy_config": Path(comfy_config_path).read_text(), - "model_class": "NunchakuFluxTransformer2dModel", + "model_class": model_class, "quantization_config": json.dumps(quantization_config), } @@ -151,10 +156,20 @@ def merge_safetensors( required=True, help="Path to model directory. It can also be a huggingface repo.", ) + parser.add_argument( + "-m", + "--model-class", + type=str, + required=True, + help="Specify model class. E.g. NunchakuFluxTransformer2dModel or NunchakuZImageTransformer2DModel", + ) parser.add_argument("-o", "--output-path", type=Path, required=True, help="Path to output path") args = parser.parse_args() - state_dict, metadata = merge_safetensors(args.input_path) + state_dict, metadata = merge_safetensors(args.input_path, args.model_class) output_path = Path(args.output_path) + print(f" --input-path: {args.input_path}") + print(f" --model-class: {args.model_class}") + print(f" --output-path: {args.output_path}") dirpath = output_path.parent dirpath.mkdir(parents=True, exist_ok=True) save_file(state_dict, output_path, metadata=metadata) diff --git a/nunchaku/models/__init__.py b/nunchaku/models/__init__.py index 58474f33..70a25eeb 100644 --- a/nunchaku/models/__init__.py +++ b/nunchaku/models/__init__.py @@ -4,6 +4,7 @@ NunchakuFluxTransformer2DModelV2, NunchakuQwenImageTransformer2DModel, NunchakuSanaTransformer2DModel, + NunchakuZImageTransformer2DModel, ) __all__ = [ @@ -12,4 +13,5 @@ "NunchakuT5EncoderModel", "NunchakuFluxTransformer2DModelV2", "NunchakuQwenImageTransformer2DModel", + "NunchakuZImageTransformer2DModel", ] diff --git a/nunchaku/models/attention_processors/zimage.py b/nunchaku/models/attention_processors/zimage.py new file mode 100644 index 00000000..f8f1b847 --- /dev/null +++ b/nunchaku/models/attention_processors/zimage.py @@ -0,0 +1,76 @@ +from typing import Optional + +import torch +from diffusers.models.attention_dispatch import dispatch_attention_fn +from diffusers.models.transformers.transformer_z_image import ZSingleStreamAttnProcessor + + +class NunchakuZSingleStreamAttnProcessor(ZSingleStreamAttnProcessor): + + def __init__(self): + super().__init__() + + # Apapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor#__call__ + def __call__( + self, + attn, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + freqs_cis: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + + qkv = attn.to_qkv(hidden_states) + query, key, value = qkv.chunk(3, dim=-1) + + query = query.unflatten(-1, (attn.heads, -1)) + key = key.unflatten(-1, (attn.heads, -1)) + value = value.unflatten(-1, (attn.heads, -1)) + + # Apply Norms + if attn.norm_q is not None: + query = attn.norm_q(query) + if attn.norm_k is not None: + key = attn.norm_k(key) + + # Apply RoPE + def apply_rotary_emb(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + with torch.amp.autocast("cuda", enabled=False): + x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x * freqs_cis).flatten(3) + return x_out.type_as(x_in) # todo + + if freqs_cis is not None: + query = apply_rotary_emb(query, freqs_cis) + key = apply_rotary_emb(key, freqs_cis) + + # Cast to correct dtype + dtype = query.dtype + query, key = query.to(dtype), key.to(dtype) + + # From [batch, seq_len] to [batch, 1, 1, seq_len] -> broadcast to [batch, heads, seq_len, seq_len] + if attention_mask is not None and attention_mask.ndim == 2: + attention_mask = attention_mask[:, None, None, :] + + # Compute joint attention + hidden_states = dispatch_attention_fn( + query, + key, + value, + attn_mask=attention_mask, + dropout_p=0.0, + is_causal=False, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + # Reshape back + hidden_states = hidden_states.flatten(2, 3) + hidden_states = hidden_states.to(dtype) + + output = attn.to_out[0](hidden_states) + if len(attn.to_out) > 1: # dropout + output = attn.to_out[1](output) + + return output diff --git a/nunchaku/models/transformers/__init__.py b/nunchaku/models/transformers/__init__.py index eddc8d46..6233f23b 100644 --- a/nunchaku/models/transformers/__init__.py +++ b/nunchaku/models/transformers/__init__.py @@ -2,10 +2,12 @@ from .transformer_flux_v2 import NunchakuFluxTransformer2DModelV2 from .transformer_qwenimage import NunchakuQwenImageTransformer2DModel from .transformer_sana import NunchakuSanaTransformer2DModel +from .transformer_zimage import NunchakuZImageTransformer2DModel __all__ = [ "NunchakuFluxTransformer2dModel", "NunchakuSanaTransformer2DModel", "NunchakuFluxTransformer2DModelV2", "NunchakuQwenImageTransformer2DModel", + "NunchakuZImageTransformer2DModel", ] diff --git a/nunchaku/models/transformers/transformer_zimage.py b/nunchaku/models/transformers/transformer_zimage.py new file mode 100644 index 00000000..19266802 --- /dev/null +++ b/nunchaku/models/transformers/transformer_zimage.py @@ -0,0 +1,258 @@ +""" +This module provides Nunchaku FluxTransformer2DModel and its building blocks in Python. +""" + +import json +import os +from pathlib import Path +from typing import List, Optional + +import torch +from diffusers.models.attention import FeedForward +from diffusers.models.attention_processor import Attention +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 ...utils import get_precision +from ..attention import NunchakuBaseAttention +from ..attention_processors.zimage import NunchakuZSingleStreamAttnProcessor +from ..linear import SVDQW4A4Linear +from ..utils import fuse_linears +from .utils import NunchakuModelLoaderMixin + + +class NunchakuZImageAttention(NunchakuBaseAttention): + """ + Nunchaku-optimized Attention module for ZImage with quantized and fused QKV projections. + + Parameters + ---------- + other : Attention + The original Attention module in ZImage model. + processor : str, optional + The attention processor to use ("flashattn2" or "nunchaku-fp16"). + **kwargs + Additional arguments for quantization. + """ + + def __init__(self, orig_attn: Attention, processor: str = "flashattn2", **kwargs): + super(NunchakuZImageAttention, self).__init__(processor) + self.inner_dim = orig_attn.inner_dim + self.query_dim = orig_attn.query_dim + self.use_bias = orig_attn.use_bias + self.dropout = orig_attn.dropout + self.out_dim = orig_attn.out_dim + self.context_pre_only = orig_attn.context_pre_only + self.pre_only = orig_attn.pre_only + self.heads = orig_attn.heads + self.rescale_output_factor = orig_attn.rescale_output_factor + self.is_cross_attention = orig_attn.is_cross_attention + + # region sub-modules + self.norm_q = orig_attn.norm_q + self.norm_k = orig_attn.norm_k + with torch.device("meta"): + to_qkv = fuse_linears([orig_attn.to_q, orig_attn.to_k, orig_attn.to_v]) + self.to_qkv = SVDQW4A4Linear.from_linear(to_qkv, **kwargs) + self.to_out = orig_attn.to_out + self.to_out[0] = SVDQW4A4Linear.from_linear(self.to_out[0], **kwargs) + # end of region + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: Optional[torch.Tensor] = None, + attention_mask: Optional[torch.Tensor] = None, + **cross_attention_kwargs, + ) -> torch.Tensor: + """ + Forward pass for NunchakuZImageAttention. + + Parameters + ---------- + hidden_states : torch.Tensor + Input tensor. + encoder_hidden_states : torch.Tensor, optional + Encoder hidden states for cross-attention. + attention_mask : torch.Tensor, optional + Attention mask. + **cross_attention_kwargs + Additional arguments for cross attention. + + Returns + ------- + Output of the attention processor. + """ + return self.processor( + attn=self, + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=attention_mask, + **cross_attention_kwargs, + ) + + def set_processor(self, processor: str): + """ + Set the attention processor. + + Parameters + ---------- + processor : str + Name of the processor ("flashattn2" or "nunchaku-fp16"). + + - ``"flashattn2"``: Standard FlashAttention-2. See :class:`~nunchaku.models.attention_processors.zimage.NunchakuZImageFA2Processor`. + + Raises + ------ + ValueError + If the processor is not supported. + """ + if processor == "flashattn2": + self.processor = NunchakuZSingleStreamAttnProcessor() + else: + raise ValueError(f"Processor {processor} is not supported") + + +def _convert_z_image_ff(z_ff: ZImageFeedForward) -> FeedForward: + """Replace custom FeedForward module in `ZImageTransformerBlock`s with standard FeedForward in diffusers lib.""" + assert isinstance(z_ff, ZImageFeedForward) + assert z_ff.w1.in_features == z_ff.w3.in_features + assert z_ff.w1.out_features == z_ff.w3.out_features + assert z_ff.w1.out_features == z_ff.w2.in_features + converted_ff = FeedForward( + dim=z_ff.w1.in_features, + dim_out=z_ff.w2.out_features, + dropout=0.0, + activation_fn="swiglu", + inner_dim=z_ff.w2.in_features, + bias=False, + ).to(dtype=z_ff.w1.weight.dtype, device=z_ff.w1.weight.device) + return converted_ff + + +class NunchakuZImageFeedForward(NunchakuSDXLFeedForward): + """ + Quantized feed-forward block for :class:`NunchakuZImageTransformerBlock`. + + Replaces linear layers in a FeedForward block with :class:`~nunchaku.models.linear.SVDQW4A4Linear` for quantized inference. + + Parameters + ---------- + ff : FeedForward + Source ZImage FeedForward module to quantize. + **kwargs : + Additional arguments for SVDQW4A4Linear. + """ + + def __init__(self, ff: ZImageFeedForward, **kwargs): + converted_ff = _convert_z_image_ff(ff) + # forward pass are equivalent to NunchakuSDXLFeedForward + NunchakuSDXLFeedForward.__init__(self, converted_ff, **kwargs) + + +class NunchakuZImageTransformer2DModel(ZImageTransformer2DModel, NunchakuModelLoaderMixin): + """ + Nunchaku-optimized ZImageTransformer2DModel. + """ + + def _patch_model(self, skip_refiners: bool = False, **kwargs): + """ + Patch the model by replacing attention and feed_forward modules in the orginal ZImageTransformerBlock. + + Parameters + ---------- + skip_refiners: bool + Default to `False` + if `True`, transformer blocks of `noise_refiner` and `context_refiner` will NOT be replaced. + **kwargs + Additional arguments for quantization. + + Returns + ------- + self : NunchakuZImageTransformer2DModel + The patched model. + """ + + def _patch_transformer_block(block_list: List[ZImageTransformerBlock]): + for _, block in enumerate(block_list): + block.attention = NunchakuZImageAttention(block.attention, **kwargs) + block.feed_forward = NunchakuZImageFeedForward(block.feed_forward, **kwargs) + + def _convert_feed_forward(block_list: List[ZImageTransformerBlock]): + for _, block in enumerate(block_list): + block.feed_forward = _convert_z_image_ff(block.feed_forward) + + _patch_transformer_block(self.layers) + if skip_refiners: + _convert_feed_forward(self.noise_refiner) + _convert_feed_forward(self.context_refiner) + else: + _patch_transformer_block(self.noise_refiner) + _patch_transformer_block(self.context_refiner) + return self + + @classmethod + @utils.validate_hf_hub_args + def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], **kwargs): + """ + Load a pretrained NunchakuZImageTransformer2DModel from a safetensors file. + + Parameters + ---------- + pretrained_model_name_or_path : str or os.PathLike + Path to the safetensors file. It can be a local file or a remote HuggingFace path. + **kwargs + Additional arguments (e.g., device, torch_dtype). + + Returns + ------- + NunchakuZImageTransformer2DModel + The loaded and quantized model. + + Raises + ------ + NotImplementedError + If offload is requested. + AssertionError + If the file is not a safetensors file. + """ + device = kwargs.get("device", "cpu") + offload = kwargs.get("offload", False) + + if offload: + raise NotImplementedError("Offload is not supported for FluxTransformer2DModelV2") + + torch_dtype = kwargs.get("torch_dtype", torch.bfloat16) + + if isinstance(pretrained_model_name_or_path, str): + pretrained_model_name_or_path = Path(pretrained_model_name_or_path) + + assert pretrained_model_name_or_path.is_file() or pretrained_model_name_or_path.name.endswith( + (".safetensors", ".sft") + ), "Only safetensors are supported" + transformer, model_state_dict, metadata = cls._build_model(pretrained_model_name_or_path, **kwargs) + quantization_config = json.loads(metadata.get("quantization_config", "{}")) + + rank = quantization_config.get("rank", 32) + skip_refiners = quantization_config.get("skip_refiners", False) + transformer = transformer.to(torch_dtype) + + precision = get_precision() + if precision == "fp4": + precision = "nvfp4" + + print(f"quantization_config: {quantization_config}, rank={rank}, skip_refiners={skip_refiners}") + + transformer._patch_model(skip_refiners=skip_refiners, precision=precision, rank=rank) + + transformer = transformer.to_empty(device=device) + # converted_state_dict = convert_flux_state_dict(model_state_dict) + + transformer.load_state_dict(model_state_dict) + + print(f"transformer type: {type(transformer)}") + + return transformer From de1bb361ea90497306ded4aedd98b9c1f7e77b76 Mon Sep 17 00:00:00 2001 From: dengdong Date: Fri, 5 Dec 2025 11:28:59 -0500 Subject: [PATCH 02/17] fix wrong doc string --- nunchaku/models/transformers/transformer_zimage.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/nunchaku/models/transformers/transformer_zimage.py b/nunchaku/models/transformers/transformer_zimage.py index 19266802..a8cff72d 100644 --- a/nunchaku/models/transformers/transformer_zimage.py +++ b/nunchaku/models/transformers/transformer_zimage.py @@ -1,5 +1,5 @@ """ -This module provides Nunchaku FluxTransformer2DModel and its building blocks in Python. +This module provides Nunchaku ZImageTransformer2DModel and its building blocks in Python. """ import json @@ -223,7 +223,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], offload = kwargs.get("offload", False) if offload: - raise NotImplementedError("Offload is not supported for FluxTransformer2DModelV2") + raise NotImplementedError("Offload is not supported for ZImageTransformer2DModel") torch_dtype = kwargs.get("torch_dtype", torch.bfloat16) @@ -247,12 +247,7 @@ 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 = transformer.to_empty(device=device) - # converted_state_dict = convert_flux_state_dict(model_state_dict) - transformer.load_state_dict(model_state_dict) - print(f"transformer type: {type(transformer)}") - return transformer From 85667ace93fa4558b0c462dfc11699a5bac2d86e Mon Sep 17 00:00:00 2001 From: Dong Deng Date: Sat, 6 Dec 2025 01:46:30 -0500 Subject: [PATCH 03/17] patch scale key --- .../models/transformers/transformer_flux_v2.py | 17 ++--------------- .../transformers/transformer_qwenimage.py | 17 +++-------------- .../models/transformers/transformer_zimage.py | 5 ++++- nunchaku/models/transformers/utils.py | 16 ++++++++++++++++ 4 files changed, 25 insertions(+), 30 deletions(-) diff --git a/nunchaku/models/transformers/transformer_flux_v2.py b/nunchaku/models/transformers/transformer_flux_v2.py index f2b2d9d0..0ce7ae2c 100644 --- a/nunchaku/models/transformers/transformer_flux_v2.py +++ b/nunchaku/models/transformers/transformer_flux_v2.py @@ -26,7 +26,7 @@ from ..linear import SVDQW4A4Linear from ..normalization import NunchakuAdaLayerNormZero, NunchakuAdaLayerNormZeroSingle from ..utils import fuse_linears -from .utils import NunchakuModelLoaderMixin +from .utils import NunchakuModelLoaderMixin, patch_scale_key class NunchakuFluxAttention(NunchakuBaseAttention): @@ -421,20 +421,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], transformer = transformer.to_empty(device=device) converted_state_dict = convert_flux_state_dict(model_state_dict) - state_dict = transformer.state_dict() - - for k in state_dict.keys(): - if k not in converted_state_dict: - assert ".wcscales" in k - converted_state_dict[k] = torch.ones_like(state_dict[k]) - else: - assert state_dict[k].dtype == converted_state_dict[k].dtype - - # Load the wtscale from the converted state dict. - for n, m in transformer.named_modules(): - if isinstance(m, SVDQW4A4Linear): - if m.wtscale is not None: - m.wtscale = converted_state_dict.pop(f"{n}.wtscale", 1.0) + patch_scale_key(transformer, converted_state_dict) transformer.load_state_dict(converted_state_dict) diff --git a/nunchaku/models/transformers/transformer_qwenimage.py b/nunchaku/models/transformers/transformer_qwenimage.py index 06ab6970..975dd7d5 100644 --- a/nunchaku/models/transformers/transformer_qwenimage.py +++ b/nunchaku/models/transformers/transformer_qwenimage.py @@ -26,7 +26,7 @@ from ..attention_processors.qwenimage import NunchakuQwenImageNaiveFA2Processor from ..linear import AWQW4A16Linear, SVDQW4A4Linear from ..utils import CPUOffloadManager, fuse_linears -from .utils import NunchakuModelLoaderMixin +from .utils import NunchakuModelLoaderMixin, patch_scale_key logger = diffusers_logging.get_logger(__name__) @@ -405,19 +405,8 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], theta=10000, axes_dim=list(config.get("axes_dims_rope", [16, 56, 56])), scale_rope=True ) - state_dict = transformer.state_dict() - for k in state_dict.keys(): - if k not in model_state_dict: - assert ".wcscales" in k - model_state_dict[k] = torch.ones_like(state_dict[k]) - else: - assert state_dict[k].dtype == model_state_dict[k].dtype - - # load the wtscale from the state dict, as it is a float on CPU - for n, m in transformer.named_modules(): - if isinstance(m, SVDQW4A4Linear): - if m.wtscale is not None: - m.wtscale = model_state_dict.pop(f"{n}.wtscale", 1.0) + patch_scale_key(transformer, model_state_dict) + transformer.load_state_dict(model_state_dict) transformer.set_offload(offload) diff --git a/nunchaku/models/transformers/transformer_zimage.py b/nunchaku/models/transformers/transformer_zimage.py index a8cff72d..95465938 100644 --- a/nunchaku/models/transformers/transformer_zimage.py +++ b/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 +from .utils import NunchakuModelLoaderMixin, patch_scale_key class NunchakuZImageAttention(NunchakuBaseAttention): @@ -248,6 +248,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], transformer._patch_model(skip_refiners=skip_refiners, precision=precision, rank=rank) transformer = transformer.to_empty(device=device) + + patch_scale_key(transformer, model_state_dict) + transformer.load_state_dict(model_state_dict) return transformer diff --git a/nunchaku/models/transformers/utils.py b/nunchaku/models/transformers/utils.py index 179ddd4d..1194bc96 100644 --- a/nunchaku/models/transformers/utils.py +++ b/nunchaku/models/transformers/utils.py @@ -13,6 +13,7 @@ from torch import nn from ...utils import load_state_dict_in_safetensors +from ..linear import SVDQW4A4Linear # Get log level from environment variable (default to INFO) log_level = os.getenv("LOG_LEVEL", "INFO").upper() @@ -145,3 +146,18 @@ def _build_model_legacy( with torch.device("meta"): transformer = cls.from_config(config).to(kwargs.get("torch_dtype", torch.bfloat16)) return transformer, unquantized_part_path, transformer_block_path + + +def patch_scale_key(transformer_from_config: nn.Module, state_dict_from_checkpoint: dict): + state_dict = transformer_from_config.state_dict() + for k in state_dict.keys(): + 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) From 431be72ae0e9be535e1e768dfd959a33b282bfd6 Mon Sep 17 00:00:00 2001 From: Dong Deng Date: Sun, 7 Dec 2025 10:13:52 -0500 Subject: [PATCH 04/17] add pytest case for z-image-turbo --- tests/v1/zimage/__init__.py | 0 tests/v1/zimage/test_zimage_turbo.py | 107 +++++++++++++++++++++++++++ 2 files changed, 107 insertions(+) create mode 100644 tests/v1/zimage/__init__.py create mode 100644 tests/v1/zimage/test_zimage_turbo.py diff --git a/tests/v1/zimage/__init__.py b/tests/v1/zimage/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/v1/zimage/test_zimage_turbo.py b/tests/v1/zimage/test_zimage_turbo.py new file mode 100644 index 00000000..9772942c --- /dev/null +++ b/tests/v1/zimage/test_zimage_turbo.py @@ -0,0 +1,107 @@ +import gc +import os +from pathlib import Path + +import pytest +import torch +from diffusers import ZImagePipeline + +from nunchaku import NunchakuZImageTransformer2DModel +from nunchaku.utils import get_precision, is_turing + +from ...utils import already_generate, compute_lpips +from ..utils import run_pipeline + +precision = get_precision() +torch_dtype = torch.float16 if is_turing() else torch.bfloat16 +dtype_str = "fp16" if torch_dtype == torch.float16 else "bf16" + + +@pytest.mark.parametrize( + "rank,expected_lpips", + [ + (32, {"int4-bf16": 0.37, "fp4-bf16": 0.33}), + (128, {"int4-bf16": 0.34, "fp4-bf16": 0.3}), + ], +) +def test_zimage_turbo(rank: int, expected_lpips: dict[str, float]): + model_name = "z-image-turbo" + batch_size = 1 + width = 1024 + height = 1024 + num_inference_steps = 9 + guidance_scale = 0.0 + + ref_root = os.environ.get("NUNCHAKU_TEST_CACHE_ROOT", os.path.join("test_results", "ref")) + folder_name = f"w{width}h{height}t{num_inference_steps}g{guidance_scale}" + save_dir_16bit = Path(ref_root) / model_name / dtype_str / folder_name + + repo_id = "Tongyi-MAI/Z-Image-Turbo" + dataset = [ + { + "prompt": "Table Mountain, South Africa, covered in clouds on a hot, bright summers day. Use a Sony alpha 1 to capture a lot of details. use a 100mm lense. Use aperture F 1.2 to make the mountain standout. Photo taken from Blouberg Beach ", + "negative_prompt": " ", + "filename": "landscape", + }, + { + "prompt": "A futuristic tibetan god wearing ornate robes embroidered with an infinitely complex gold mandala, very old man, white beard, character concept full body, a weathered magical Gate with glowing runes carved into a granite cliff face, stairs lined with cherry blossom trees and jacaranda trees the entrance of goddess, ornate, beautiful, weapons, lush, nature, low angle, Protoctist style Zeng Chuanxing, widescreen, anamorphic 2 39, gold , intricate detail, hyper realistic, low angle Symmetrical, epic scale Cinematic, Color Grading, F 2. 8, 8K, Ultra HD, AMOLED, Ray Tracing Global Illumination, spiritual vibes, Transparent, Translucent, Iridescent, Ray Tracing Reflections, Harris Shutter, De Noise, VFX, SFX, anamorphic 2 39 ", + "negative_prompt": " ", + "filename": "art", + }, + { + "prompt": "年轻的中国女子，身着红色汉服，绣工细密。妆容精致无瑕，额间点着红色花钿。发髻高盘而华丽，簪着金色凤凰头饰、红花与串珠。右手持一柄圆形折扇，扇面绘有仕女、树木与鸟。左手微抬，掌上方悬着一盏霓虹闪电形灯（⚡️），散发明亮的黄色光辉。背景是柔和灯光下的户外夜景，层叠的宝塔（西安大雁塔）成剪影状隐现，远处彩光朦胧。", + "negative_prompt": " ", + "filename": "portrait_chinese_prompt", + }, + ] + + if not already_generate(save_dir_16bit, len(dataset)): + pipe = ZImagePipeline.from_pretrained(repo_id, torch_dtype=torch_dtype).to("cuda") + run_pipeline( + dataset=dataset, + batch_size=1, + pipeline=pipe, + save_dir=save_dir_16bit, + forward_kwargs={ + "width": width, + "height": height, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + }, + ) + del pipe + gc.collect() + torch.cuda.empty_cache() + + save_dir_nunchaku = ( + Path("test_results") + / "nunchaku" + / model_name + / f"{precision}_r{rank}-{dtype_str}" + / 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).to("cuda") + + 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}-{dtype_str}"] * 1.10 From 938e4188cdf24461b4687cad0f4ef79730a93e32 Mon Sep 17 00:00:00 2001 From: Dong Deng Date: Sun, 7 Dec 2025 10:46:25 -0500 Subject: [PATCH 05/17] fix test case --- tests/v1/zimage/test_zimage_turbo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/v1/zimage/test_zimage_turbo.py b/tests/v1/zimage/test_zimage_turbo.py index 9772942c..058e9dce 100644 --- a/tests/v1/zimage/test_zimage_turbo.py +++ b/tests/v1/zimage/test_zimage_turbo.py @@ -80,7 +80,7 @@ def test_zimage_turbo(rank: int, expected_lpips: dict[str, float]): / f"{precision}_r{rank}-{dtype_str}" / f"{folder_name}-bs{batch_size}" ) - path = (f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}-z-image-turbo.safetensors",) + 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).to("cuda") From 579826df791d1efa4da2a8554ef19846e534f7f2 Mon Sep 17 00:00:00 2001 From: dengdong Date: Sun, 7 Dec 2025 10:59:09 -0500 Subject: [PATCH 06/17] add test case for `skip_refiners` model --- tests/v1/zimage/test_zimage_turbo.py | 121 ++++++++++++++++++++------- 1 file changed, 91 insertions(+), 30 deletions(-) diff --git a/tests/v1/zimage/test_zimage_turbo.py b/tests/v1/zimage/test_zimage_turbo.py index 058e9dce..d990cdcb 100644 --- a/tests/v1/zimage/test_zimage_turbo.py +++ b/tests/v1/zimage/test_zimage_turbo.py @@ -16,6 +16,37 @@ torch_dtype = torch.float16 if is_turing() else torch.bfloat16 dtype_str = "fp16" if torch_dtype == torch.float16 else "bf16" +model_name = "z-image-turbo" +batch_size = 1 +width = 1024 +height = 1024 +num_inference_steps = 9 +guidance_scale = 0.0 + +ref_root = os.environ.get("NUNCHAKU_TEST_CACHE_ROOT", os.path.join("test_results", "ref")) +folder_name = f"w{width}h{height}t{num_inference_steps}g{guidance_scale}" +save_dir_16bit = Path(ref_root) / model_name / dtype_str / folder_name + +repo_id = "Tongyi-MAI/Z-Image-Turbo" + +dataset = [ + { + "prompt": "Table Mountain, South Africa, covered in clouds on a hot, bright summers day. Use a Sony alpha 1 to capture a lot of details. use a 100mm lense. Use aperture F 1.2 to make the mountain standout. Photo taken from Blouberg Beach ", + "negative_prompt": " ", + "filename": "landscape", + }, + { + "prompt": "A futuristic tibetan god wearing ornate robes embroidered with an infinitely complex gold mandala, very old man, white beard, character concept full body, a weathered magical Gate with glowing runes carved into a granite cliff face, stairs lined with cherry blossom trees and jacaranda trees the entrance of goddess, ornate, beautiful, weapons, lush, nature, low angle, Protoctist style Zeng Chuanxing, widescreen, anamorphic 2 39, gold , intricate detail, hyper realistic, low angle Symmetrical, epic scale Cinematic, Color Grading, F 2. 8, 8K, Ultra HD, AMOLED, Ray Tracing Global Illumination, spiritual vibes, Transparent, Translucent, Iridescent, Ray Tracing Reflections, Harris Shutter, De Noise, VFX, SFX, anamorphic 2 39 ", + "negative_prompt": " ", + "filename": "art", + }, + { + "prompt": "年轻的中国女子，身着红色汉服，绣工细密。妆容精致无瑕，额间点着红色花钿。发髻高盘而华丽，簪着金色凤凰头饰、红花与串珠。右手持一柄圆形折扇，扇面绘有仕女、树木与鸟。左手微抬，掌上方悬着一盏霓虹闪电形灯（⚡️），散发明亮的黄色光辉。背景是柔和灯光下的户外夜景，层叠的宝塔（西安大雁塔）成剪影状隐现，远处彩光朦胧。", + "negative_prompt": " ", + "filename": "portrait_chinese_prompt", + }, +] + @pytest.mark.parametrize( "rank,expected_lpips", @@ -25,36 +56,6 @@ ], ) def test_zimage_turbo(rank: int, expected_lpips: dict[str, float]): - model_name = "z-image-turbo" - batch_size = 1 - width = 1024 - height = 1024 - num_inference_steps = 9 - guidance_scale = 0.0 - - ref_root = os.environ.get("NUNCHAKU_TEST_CACHE_ROOT", os.path.join("test_results", "ref")) - folder_name = f"w{width}h{height}t{num_inference_steps}g{guidance_scale}" - save_dir_16bit = Path(ref_root) / model_name / dtype_str / folder_name - - repo_id = "Tongyi-MAI/Z-Image-Turbo" - dataset = [ - { - "prompt": "Table Mountain, South Africa, covered in clouds on a hot, bright summers day. Use a Sony alpha 1 to capture a lot of details. use a 100mm lense. Use aperture F 1.2 to make the mountain standout. Photo taken from Blouberg Beach ", - "negative_prompt": " ", - "filename": "landscape", - }, - { - "prompt": "A futuristic tibetan god wearing ornate robes embroidered with an infinitely complex gold mandala, very old man, white beard, character concept full body, a weathered magical Gate with glowing runes carved into a granite cliff face, stairs lined with cherry blossom trees and jacaranda trees the entrance of goddess, ornate, beautiful, weapons, lush, nature, low angle, Protoctist style Zeng Chuanxing, widescreen, anamorphic 2 39, gold , intricate detail, hyper realistic, low angle Symmetrical, epic scale Cinematic, Color Grading, F 2. 8, 8K, Ultra HD, AMOLED, Ray Tracing Global Illumination, spiritual vibes, Transparent, Translucent, Iridescent, Ray Tracing Reflections, Harris Shutter, De Noise, VFX, SFX, anamorphic 2 39 ", - "negative_prompt": " ", - "filename": "art", - }, - { - "prompt": "年轻的中国女子，身着红色汉服，绣工细密。妆容精致无瑕，额间点着红色花钿。发髻高盘而华丽，簪着金色凤凰头饰、红花与串珠。右手持一柄圆形折扇，扇面绘有仕女、树木与鸟。左手微抬，掌上方悬着一盏霓虹闪电形灯（⚡️），散发明亮的黄色光辉。背景是柔和灯光下的户外夜景，层叠的宝塔（西安大雁塔）成剪影状隐现，远处彩光朦胧。", - "negative_prompt": " ", - "filename": "portrait_chinese_prompt", - }, - ] - if not already_generate(save_dir_16bit, len(dataset)): pipe = ZImagePipeline.from_pretrained(repo_id, torch_dtype=torch_dtype).to("cuda") run_pipeline( @@ -105,3 +106,63 @@ 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.10 + + +@pytest.mark.skipif(get_precision() == "fp4", reason="The `skip_refiners` option is valid only for int4 precision.") +@pytest.mark.parametrize( + "rank,expected_lpips", + [ + (128, {"int4-bf16": 0.34}), + ], +) +def test_zimage_turbo_r128_sr(rank: int, expected_lpips: dict[str, float]): + if not already_generate(save_dir_16bit, len(dataset)): + pipe = ZImagePipeline.from_pretrained(repo_id, torch_dtype=torch_dtype).to("cuda") + run_pipeline( + dataset=dataset, + batch_size=1, + pipeline=pipe, + save_dir=save_dir_16bit, + forward_kwargs={ + "width": width, + "height": height, + "num_inference_steps": num_inference_steps, + "guidance_scale": guidance_scale, + }, + ) + del pipe + gc.collect() + torch.cuda.empty_cache() + + save_dir_nunchaku = ( + Path("test_results") + / "nunchaku" + / model_name + / f"{precision}_r{rank}_sr-{dtype_str}" + / f"{folder_name}-bs{batch_size}" + ) + path = f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}_sr-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).to("cuda") + + 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}-{dtype_str}"] * 1.10 From a054f75932d4ae3976287b6ec7de2448ffcbf5e6 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Sun, 7 Dec 2025 13:31:24 -0500 Subject: [PATCH 07/17] update example --- examples/v1/z-image-turbo.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/examples/v1/z-image-turbo.py b/examples/v1/z-image-turbo.py index 80052cad..dd13861b 100644 --- a/examples/v1/z-image-turbo.py +++ b/examples/v1/z-image-turbo.py @@ -7,14 +7,14 @@ if __name__ == "__main__": precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU transformer = NunchakuZImageTransformer2DModel.from_pretrained( - f"/PATH/TO/svdq-{precision}_r128-z-image-turbo.safetensors" + f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r128-z-image-turbo.safetensors" ) pipe = ZImagePipeline.from_pretrained( "Tongyi-MAI/Z-Image-Turbo", transformer=transformer, torch_dtype=torch.bfloat16, - low_cpu_mem_usage=False, + low_cpu_mem_usage=False ).to("cuda") prompt = "a young military male cooking in the kitchen for therapy" @@ -23,9 +23,9 @@ prompt=prompt, height=1024, width=1024, - num_inference_steps=9, # This actually results in 8 DiT forwards + num_inference_steps=8, # This actually results in 8 DiT forwards guidance_scale=0.0, # Guidance should be 0 for the Turbo models - generator=torch.Generator("cuda").manual_seed(12345), + generator=torch.Generator().manual_seed(12345), ).images[0] - image.save(f"tmp_imgs/z-image-turbo-{precision}.png") + image.save(f"z-image-turbo-{precision}.png") From 5a101a496273933417654bb27a0166469d55ad14 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Sun, 7 Dec 2025 14:32:46 -0500 Subject: [PATCH 08/17] update --- examples/v1/z-image-turbo.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/v1/z-image-turbo.py b/examples/v1/z-image-turbo.py index dd13861b..8fedfc7f 100644 --- a/examples/v1/z-image-turbo.py +++ b/examples/v1/z-image-turbo.py @@ -6,8 +6,9 @@ if __name__ == "__main__": precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU + rank = 32 # you can also use rank=32 model to improve the speed transformer = NunchakuZImageTransformer2DModel.from_pretrained( - f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r128-z-image-turbo.safetensors" + f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}-z-image-turbo.safetensors" ) pipe = ZImagePipeline.from_pretrained( @@ -28,4 +29,4 @@ generator=torch.Generator().manual_seed(12345), ).images[0] - image.save(f"z-image-turbo-{precision}.png") + image.save(f"z-image-turbo-{precision}_r{rank}.png") From 06f309111f61802cff3d33e564fff4b68617483b Mon Sep 17 00:00:00 2001 From: dengdong Date: Sun, 7 Dec 2025 23:09:35 -0500 Subject: [PATCH 09/17] update docs --- ...models.transformers.transformer_zimage.rst | 7 +++++++ docs/source/usage/zimage.rst | 16 ++++++++++++++++ nunchaku/merge_safetensors.py | 2 ++ .../models/transformers/transformer_zimage.py | 19 ++++++++++++++++--- nunchaku/models/transformers/utils.py | 10 ++++++++++ 5 files changed, 51 insertions(+), 3 deletions(-) create mode 100644 docs/source/python_api/nunchaku.models.transformers.transformer_zimage.rst create mode 100644 docs/source/usage/zimage.rst diff --git a/docs/source/python_api/nunchaku.models.transformers.transformer_zimage.rst b/docs/source/python_api/nunchaku.models.transformers.transformer_zimage.rst new file mode 100644 index 00000000..e7a06355 --- /dev/null +++ b/docs/source/python_api/nunchaku.models.transformers.transformer_zimage.rst @@ -0,0 +1,7 @@ +nunchaku.models.transformers.transformer\_zimage +================================================ + +.. automodule:: nunchaku.models.transformers.transformer_zimage + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/source/usage/zimage.rst b/docs/source/usage/zimage.rst new file mode 100644 index 00000000..2c825fc4 --- /dev/null +++ b/docs/source/usage/zimage.rst @@ -0,0 +1,16 @@ +Z-Image +======= + +The following is the example of running Nunchaku version of Z-Image text-to-image pipeline. + +.. tabs:: + + .. tab:: Z-Image-Turbo + + .. literalinclude:: ../../../examples/v1/z-image-turbo.py + :language: python + :caption: Running Z-Image-Turbo (`examples/v1/z-image-turbo.py `__) + :linenos: + + +For more details, see :class:`~nunchaku.models.transformers.transformer_zimage.NunchakuZImageTransformer2DModel`. diff --git a/nunchaku/merge_safetensors.py b/nunchaku/merge_safetensors.py index d48e6fb2..c89ecf63 100644 --- a/nunchaku/merge_safetensors.py +++ b/nunchaku/merge_safetensors.py @@ -47,6 +47,8 @@ def merge_safetensors( ---------- pretrained_model_name_or_path : str or os.PathLike Path to the model directory or HuggingFace repo. + model_class : str + Specify model class. E.g. NunchakuFluxTransformer2dModel or NunchakuZImageTransformer2DModel **kwargs Additional keyword arguments for subfolder, comfy_config_path, and HuggingFace download options. diff --git a/nunchaku/models/transformers/transformer_zimage.py b/nunchaku/models/transformers/transformer_zimage.py index 95465938..1b89c90b 100644 --- a/nunchaku/models/transformers/transformer_zimage.py +++ b/nunchaku/models/transformers/transformer_zimage.py @@ -101,9 +101,9 @@ def set_processor(self, processor: str): Parameters ---------- processor : str - Name of the processor ("flashattn2" or "nunchaku-fp16"). + Name of the processor ("flashattn2"). - - ``"flashattn2"``: Standard FlashAttention-2. See :class:`~nunchaku.models.attention_processors.zimage.NunchakuZImageFA2Processor`. + - ``"flashattn2"``: Standard FlashAttention-2. See :class:`~nunchaku.models.attention_processors.zimage.NunchakuZSingleStreamAttnProcessor`. Raises ------ @@ -117,7 +117,20 @@ def set_processor(self, processor: str): def _convert_z_image_ff(z_ff: ZImageFeedForward) -> FeedForward: - """Replace custom FeedForward module in `ZImageTransformerBlock`s with standard FeedForward in diffusers lib.""" + """ + Replace custom FeedForward module in `ZImageTransformerBlock`s with standard FeedForward in diffusers lib. + + Parameters + ---------- + z_ff : ZImageFeedForward + The feed forward sub-module in the ZImageTransformerBlock module + + Returns + ------- + FeedForward + A diffusers FeedForward module which is equivalent to the input `z_ff` + + """ assert isinstance(z_ff, ZImageFeedForward) assert z_ff.w1.in_features == z_ff.w3.in_features assert z_ff.w1.out_features == z_ff.w3.out_features diff --git a/nunchaku/models/transformers/utils.py b/nunchaku/models/transformers/utils.py index 1194bc96..beabc75a 100644 --- a/nunchaku/models/transformers/utils.py +++ b/nunchaku/models/transformers/utils.py @@ -149,6 +149,16 @@ def _build_model_legacy( def patch_scale_key(transformer_from_config: nn.Module, state_dict_from_checkpoint: dict): + """ + Modify scale parameters so that the state dict from the checkpoint file can be loaded to the transformer model created from the config. + + Parameters + ---------- + transformer_from_config : nn.Module + The transformer model created from the `config.json` + state_dict_from_checkpoint : dict + The state dict loaded from the checkpoint file (typically .safetensors) + """ state_dict = transformer_from_config.state_dict() for k in state_dict.keys(): if k not in state_dict_from_checkpoint: From 52a691e253402834aaa1cb5f321d1345d8ce12bd Mon Sep 17 00:00:00 2001 From: dengdong Date: Sun, 7 Dec 2025 23:21:30 -0500 Subject: [PATCH 10/17] update docs --- ...nunchaku.models.attention_processors.zimage.rst | 7 +++++++ nunchaku/models/attention_processors/zimage.py | 14 +++++++++++++- 2 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 docs/source/python_api/nunchaku.models.attention_processors.zimage.rst diff --git a/docs/source/python_api/nunchaku.models.attention_processors.zimage.rst b/docs/source/python_api/nunchaku.models.attention_processors.zimage.rst new file mode 100644 index 00000000..9d919490 --- /dev/null +++ b/docs/source/python_api/nunchaku.models.attention_processors.zimage.rst @@ -0,0 +1,7 @@ +nunchaku.models.attention_processors.zimage +=========================================== + +.. automodule:: nunchaku.models.attention_processors.zimage + :members: + :undoc-members: + :show-inheritance: diff --git a/nunchaku/models/attention_processors/zimage.py b/nunchaku/models/attention_processors/zimage.py index f8f1b847..2818baa0 100644 --- a/nunchaku/models/attention_processors/zimage.py +++ b/nunchaku/models/attention_processors/zimage.py @@ -1,3 +1,7 @@ +""" +Attention processor implementations for :class:`~nunchaku.models.transformers.transformer_zimage.NunchakuZImageAttention`. +""" + from typing import Optional import torch @@ -6,11 +10,16 @@ class NunchakuZSingleStreamAttnProcessor(ZSingleStreamAttnProcessor): + """ + Nunchaku attention processor for Z-Image-Turbo. + Adapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor. + + """ def __init__(self): super().__init__() - # Apapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor#__call__ + # Adapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor#__call__ def __call__( self, attn, @@ -19,6 +28,9 @@ def __call__( attention_mask: Optional[torch.Tensor] = None, freqs_cis: Optional[torch.Tensor] = None, ) -> torch.Tensor: + """ + Forward pass of the attention module. Adapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor#__call__. + """ qkv = attn.to_qkv(hidden_states) query, key, value = qkv.chunk(3, dim=-1) From 727fac4cf97d9689acbdfb62060637713b1ebd41 Mon Sep 17 00:00:00 2001 From: Dong Deng Date: Mon, 15 Dec 2025 04:31:56 -0500 Subject: [PATCH 11/17] update docs --- docs/source/index.rst | 1 + docs/source/python_api/nunchaku.models.attention_processors.rst | 1 + docs/source/python_api/nunchaku.models.transformers.rst | 1 + 3 files changed, 3 insertions(+) diff --git a/docs/source/index.rst b/docs/source/index.rst index 231512b0..dded1bca 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -28,6 +28,7 @@ Check out `DeepCompressor `_ for the quantization librar usage/fbcache.rst usage/pulid.rst usage/ip_adapter.rst + usage/zimage.rst .. toctree:: :maxdepth: 1 diff --git a/docs/source/python_api/nunchaku.models.attention_processors.rst b/docs/source/python_api/nunchaku.models.attention_processors.rst index a06eb40d..1f43b394 100644 --- a/docs/source/python_api/nunchaku.models.attention_processors.rst +++ b/docs/source/python_api/nunchaku.models.attention_processors.rst @@ -6,3 +6,4 @@ nunchaku.models.attention_processors nunchaku.models.attention_processors.flux nunchaku.models.attention_processors.qwenimage + nunchaku.models.attention_processors.zimage diff --git a/docs/source/python_api/nunchaku.models.transformers.rst b/docs/source/python_api/nunchaku.models.transformers.rst index 37c5b776..38b74a1f 100644 --- a/docs/source/python_api/nunchaku.models.transformers.rst +++ b/docs/source/python_api/nunchaku.models.transformers.rst @@ -7,5 +7,6 @@ nunchaku.models.transformers nunchaku.models.transformers.transformer_flux nunchaku.models.transformers.transformer_flux_v2 nunchaku.models.transformers.transformer_qwenimage + nunchaku.models.transformers.transformer_zimage nunchaku.models.transformers.transformer_sana nunchaku.models.transformers.utils From 8bbdc1ca8bff0719761e3c251cb693de758f22be Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Fri, 19 Dec 2025 22:27:56 -0500 Subject: [PATCH 12/17] update --- examples/v1/z-image-turbo.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/examples/v1/z-image-turbo.py b/examples/v1/z-image-turbo.py index 8fedfc7f..e09e9b95 100644 --- a/examples/v1/z-image-turbo.py +++ b/examples/v1/z-image-turbo.py @@ -6,16 +6,13 @@ if __name__ == "__main__": precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU - rank = 32 # you can also use rank=32 model to improve the speed + rank = 32 # you can also use rank=32 model to improve the speed transformer = NunchakuZImageTransformer2DModel.from_pretrained( f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}-z-image-turbo.safetensors" ) 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=torch.bfloat16, low_cpu_mem_usage=False ).to("cuda") prompt = "a young military male cooking in the kitchen for therapy" From 1d1c58167ea022f56f05ebe423f95ec43e2e4934 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Fri, 19 Dec 2025 22:43:07 -0500 Subject: [PATCH 13/17] runnable --- examples/v1/z-image-turbo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/v1/z-image-turbo.py b/examples/v1/z-image-turbo.py index e09e9b95..87f8e51f 100644 --- a/examples/v1/z-image-turbo.py +++ b/examples/v1/z-image-turbo.py @@ -6,7 +6,7 @@ if __name__ == "__main__": precision = get_precision() # auto-detect your precision is 'int4' or 'fp4' based on your GPU - rank = 32 # you can also use rank=32 model to improve the speed + rank = 128 # Use 32 for faster sampling; 256 (INT4 only) for best quality transformer = NunchakuZImageTransformer2DModel.from_pretrained( f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}-z-image-turbo.safetensors" ) From c7d9e7a9f2b5aa5c8850d82304bf4dfa1a24a3f3 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Fri, 19 Dec 2025 22:52:29 -0500 Subject: [PATCH 14/17] rename the test --- tests/v1/{zimage => z_image}/__init__.py | 0 .../test_zimage_turbo.py => z_image/test_z_image_turbo.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename tests/v1/{zimage => z_image}/__init__.py (100%) rename tests/v1/{zimage/test_zimage_turbo.py => z_image/test_z_image_turbo.py} (100%) diff --git a/tests/v1/zimage/__init__.py b/tests/v1/z_image/__init__.py similarity index 100% rename from tests/v1/zimage/__init__.py rename to tests/v1/z_image/__init__.py diff --git a/tests/v1/zimage/test_zimage_turbo.py b/tests/v1/z_image/test_z_image_turbo.py similarity index 100% rename from tests/v1/zimage/test_zimage_turbo.py rename to tests/v1/z_image/test_z_image_turbo.py From da61478da309bf6b1799b94e02b403a278fa7824 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Fri, 19 Dec 2025 23:11:40 -0500 Subject: [PATCH 15/17] skip the fp4_r256 test --- tests/v1/z_image/test_z_image_turbo.py | 70 +++----------------------- 1 file changed, 7 insertions(+), 63 deletions(-) diff --git a/tests/v1/z_image/test_z_image_turbo.py b/tests/v1/z_image/test_z_image_turbo.py index d990cdcb..e2dacba1 100644 --- a/tests/v1/z_image/test_z_image_turbo.py +++ b/tests/v1/z_image/test_z_image_turbo.py @@ -51,11 +51,15 @@ @pytest.mark.parametrize( "rank,expected_lpips", [ - (32, {"int4-bf16": 0.37, "fp4-bf16": 0.33}), - (128, {"int4-bf16": 0.34, "fp4-bf16": 0.3}), + (32, {"int4-bf16": 0.4, "fp4-bf16": 0.33}), + (128, {"int4-bf16": 0.38, "fp4-bf16": 0.3}), + (256, {"int4-bf16": 0.37}), ], ) def test_zimage_turbo(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)): pipe = ZImagePipeline.from_pretrained(repo_id, torch_dtype=torch_dtype).to("cuda") run_pipeline( @@ -105,64 +109,4 @@ 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.10 - - -@pytest.mark.skipif(get_precision() == "fp4", reason="The `skip_refiners` option is valid only for int4 precision.") -@pytest.mark.parametrize( - "rank,expected_lpips", - [ - (128, {"int4-bf16": 0.34}), - ], -) -def test_zimage_turbo_r128_sr(rank: int, expected_lpips: dict[str, float]): - if not already_generate(save_dir_16bit, len(dataset)): - pipe = ZImagePipeline.from_pretrained(repo_id, torch_dtype=torch_dtype).to("cuda") - run_pipeline( - dataset=dataset, - batch_size=1, - pipeline=pipe, - save_dir=save_dir_16bit, - forward_kwargs={ - "width": width, - "height": height, - "num_inference_steps": num_inference_steps, - "guidance_scale": guidance_scale, - }, - ) - del pipe - gc.collect() - torch.cuda.empty_cache() - - save_dir_nunchaku = ( - Path("test_results") - / "nunchaku" - / model_name - / f"{precision}_r{rank}_sr-{dtype_str}" - / f"{folder_name}-bs{batch_size}" - ) - path = f"nunchaku-tech/nunchaku-z-image-turbo/svdq-{precision}_r{rank}_sr-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).to("cuda") - - 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}-{dtype_str}"] * 1.10 + assert lpips < expected_lpips[f"{precision}-{dtype_str}"] * 1.15 From 14dd5f96b50ad1320ed54944b35d17bd0d28c82e Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Fri, 19 Dec 2025 23:51:49 -0500 Subject: [PATCH 16/17] update the dependencies --- pyproject.toml | 2 +- tests/v1/z_image/test_z_image_turbo.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3e377ecb..faf486a6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,7 +36,7 @@ optional-dependencies.ci = [ "accelerate==1.9", "controlnet-aux==0.0.10", "datasets==3.6", - "diffusers @ git+https://github.com/huggingface/diffusers@041501a", + "diffusers @ git+https://github.com/huggingface/diffusers@262ce19", "facexlib==0.3", "huggingface-hub==0.34", "image-gen-aux @ git+https://github.com/asomoza/image_gen_aux.git@e08ae3d", diff --git a/tests/v1/z_image/test_z_image_turbo.py b/tests/v1/z_image/test_z_image_turbo.py index e2dacba1..6a68770c 100644 --- a/tests/v1/z_image/test_z_image_turbo.py +++ b/tests/v1/z_image/test_z_image_turbo.py @@ -59,7 +59,7 @@ def test_zimage_turbo(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)): pipe = ZImagePipeline.from_pretrained(repo_id, torch_dtype=torch_dtype).to("cuda") run_pipeline( From ef6efbb318b65f76dd4fcf8e409e692d90ce9e90 Mon Sep 17 00:00:00 2001 From: Muyang Li Date: Sat, 20 Dec 2025 10:31:53 -0500 Subject: [PATCH 17/17] update image link --- tests/flux/test_flux_dev_pulid.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/flux/test_flux_dev_pulid.py b/tests/flux/test_flux_dev_pulid.py index 5f40113a..c2418149 100644 --- a/tests/flux/test_flux_dev_pulid.py +++ b/tests/flux/test_flux_dev_pulid.py @@ -31,7 +31,7 @@ def test_flux_dev_pulid(): pipeline.transformer.forward = MethodType(pulid_forward, pipeline.transformer) - id_image = load_image("https://github.com/ToTheBeginning/PuLID/blob/main/example_inputs/liuyifei.png?raw=true") + id_image = load_image("https://raw.githubusercontent.com/ToTheBeginning/PuLID/main/example_inputs/liuyifei.png") image = pipeline( "A woman holding a sign that says hello world",