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.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/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 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/examples/v1/z-image-turbo.py b/examples/v1/z-image-turbo.py new file mode 100644 index 00000000..87f8e51f --- /dev/null +++ b/examples/v1/z-image-turbo.py @@ -0,0 +1,29 @@ +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 + 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" + ) + + 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=8, # This actually results in 8 DiT forwards + guidance_scale=0.0, # Guidance should be 0 for the Turbo models + generator=torch.Generator().manual_seed(12345), + ).images[0] + + image.save(f"z-image-turbo-{precision}_r{rank}.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..c89ecf63 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. @@ -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. @@ -108,6 +110,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 +139,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 +158,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..2818baa0 --- /dev/null +++ b/nunchaku/models/attention_processors/zimage.py @@ -0,0 +1,88 @@ +""" +Attention processor implementations for :class:`~nunchaku.models.transformers.transformer_zimage.NunchakuZImageAttention`. +""" + +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): + """ + Nunchaku attention processor for Z-Image-Turbo. + Adapted from diffusers.models.transformers.transformer_z_image.ZSingleStreamAttnProcessor. + + """ + + def __init__(self): + super().__init__() + + # Adapted 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: + """ + 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) + + 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_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 new file mode 100644 index 00000000..1b89c90b --- /dev/null +++ b/nunchaku/models/transformers/transformer_zimage.py @@ -0,0 +1,269 @@ +""" +This module provides Nunchaku ZImageTransformer2DModel 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, patch_scale_key + + +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"). + + - ``"flashattn2"``: Standard FlashAttention-2. See :class:`~nunchaku.models.attention_processors.zimage.NunchakuZSingleStreamAttnProcessor`. + + 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. + + 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 + 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 ZImageTransformer2DModel") + + 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) + + 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..beabc75a 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,28 @@ 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): + """ + 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: + 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) 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/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", diff --git a/tests/v1/z_image/__init__.py b/tests/v1/z_image/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/v1/z_image/test_z_image_turbo.py b/tests/v1/z_image/test_z_image_turbo.py new file mode 100644 index 00000000..6a68770c --- /dev/null +++ b/tests/v1/z_image/test_z_image_turbo.py @@ -0,0 +1,112 @@ +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" + +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", + [ + (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( + 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.15