support Z-Image-Turbo model (#687)

Author: xinyiqin

assets/inputs/imgs/sketch-mountains-input.jpg

@@ -0,0 +1,40 @@
+{
+    "aspect_ratios": {
+        "1:1": [
+            1328,
+            1328
+        ],
+        "16:9": [
+            1664,
+            928
+        ],
+        "9:16": [
+            928,
+            1664
+        ],
+        "4:3": [
+            1472,
+            1140
+        ],
+        "3:4": [
+            768,
+            1024
+        ]
+    },
+    "aspect_ratio": "16:9",
+    "num_channels_latents": 16,
+    "batchsize": 1,
+    "vae_scale_factor": 8,
+    "infer_steps": 9,
+    "num_layers": 30,
+    "attention_out_dim": 3840,
+    "attention_dim_head": 128,
+    "attn_type": "flash_attn3",
+    "enable_cfg": true,
+    "sample_guide_scale": 0.0,
+    "n_refiner_layers": 2,
+    "patch_size": 2,
+    "strength": 0.6,
+    "transformer_in_channels": 64,
+    "_auto_resize": true
+}

configs/z_image/z_image_turbo_i2i.json

@@ -0,0 +1,37 @@
+{
+    "aspect_ratios": {
+        "1:1": [
+            1328,
+            1328
+        ],
+        "16:9": [
+            1664,
+            928
+        ],
+        "9:16": [
+            928,
+            1664
+        ],
+        "4:3": [
+            1472,
+            1140
+        ],
+        "3:4": [
+            768,
+            1024
+        ]
+    },
+    "aspect_ratio": "16:9",
+    "num_channels_latents": 16,
+    "batchsize": 1,
+    "vae_scale_factor": 8,
+    "infer_steps": 9,
+    "num_layers": 30,
+    "attention_out_dim": 3840,
+    "attention_dim_head": 128,
+    "attn_type": "flash_attn3",
+    "enable_cfg": false,
+    "sample_guide_scale": 0.0,
+    "n_refiner_layers": 2,
+    "patch_size": 2
+}

configs/z_image/z_image_turbo_t2i.json

@@ -15,6 +15,7 @@
 from lightx2v.models.runners.wan.wan_runner import Wan22MoeRunner, WanRunner  # noqa: F401
 from lightx2v.models.runners.wan.wan_sf_runner import WanSFRunner  # noqa: F401
 from lightx2v.models.runners.wan.wan_vace_runner import WanVaceRunner  # noqa: F401
+from lightx2v.models.runners.z_image.z_image_runner import ZImageRunner  # noqa: F401
 from lightx2v.utils.envs import *
 from lightx2v.utils.input_info import set_input_info
 from lightx2v.utils.profiler import *
@@ -56,6 +57,7 @@ def main():
             "wan2.2_animate",
             "hunyuan_video_1.5",
             "hunyuan_video_1.5_distill",
+            "z_image",
         ],
         default="wan2.1",
     )
@@ -118,6 +120,22 @@ def main():
     )
     parser.add_argument("--save_result_path", type=str, default=None, help="The path to save video path/file")
     parser.add_argument("--return_result_tensor", action="store_true", help="Whether to return result tensor. (Useful for comfyui)")
+
+    # Aspect ratio and custom shape for image tasks (t2i, i2i)
+    parser.add_argument(
+        "--aspect_ratio",
+        type=str,
+        default="16:9",
+        choices=["16:9", "9:16", "1:1", "4:3", "3:4"],
+        help="Aspect ratio for image generation. Only used for t2i and i2i tasks.",
+    )
+    parser.add_argument(
+        "--custom_shape",
+        type=str,
+        default=None,
+        help="Custom shape for image generation in format 'height,width' (e.g., '928,1664'). Only used for t2i and i2i tasks. Takes precedence over aspect_ratio.",
+    )
+    parser.add_argument("--strength", type=float, default=0.6, help="The strength for image-to-image generation")
     args = parser.parse_args()
     validate_task_arguments(args)
 

lightx2v/infer.py

@@ -0,0 +1,94 @@
+import gc
+import os
+
+import torch
+from PIL import Image
+
+try:
+    from transformers import Qwen2Tokenizer, Qwen3Model
+except ImportError:
+    Qwen2Tokenizer = None
+    Qwen3Model = None
+
+from lightx2v_platform.base.global_var import AI_DEVICE
+
+torch_device_module = getattr(torch, AI_DEVICE)
+
+try:
+    from diffusers.image_processor import VaeImageProcessor
+except ImportError:
+    VaeImageProcessor = None
+
+
+class Qwen3Model_TextEncoder:
+    def __init__(self, config):
+        self.config = config
+        self.tokenizer_max_length = 512
+        self.cpu_offload = config.get("qwen3_cpu_offload", config.get("cpu_offload", False))
+        self.dtype = torch.bfloat16
+        self.load()
+
+    def load(self):
+        self.text_encoder = Qwen3Model.from_pretrained(os.path.join(self.config["model_path"], "text_encoder"), torch_dtype=torch.bfloat16)
+        if not self.cpu_offload:
+            self.text_encoder = self.text_encoder.to(AI_DEVICE)
+
+        self.tokenizer = Qwen2Tokenizer.from_pretrained(os.path.join(self.config["model_path"], "tokenizer"))
+
+        if self.config["task"] == "i2i":
+            self.image_processor = VaeImageProcessor(vae_scale_factor=self.config.get("vae_scale_factor", 8) * 2)
+
+    def preprocess_image(self, image):
+        if isinstance(image, Image.Image):
+            preprocessed_image = self.image_processor.preprocess(image)
+        elif isinstance(image, torch.Tensor):
+            if image.dim() == 3:
+                image = image.unsqueeze(0)
+            preprocessed_image = image
+        else:
+            raise ValueError(f"Unsupported image type: {type(image)}")
+
+        return preprocessed_image
+
+    @torch.no_grad()
+    def infer(self, prompt, image_list=None):
+        if self.cpu_offload:
+            self.text_encoder.to(AI_DEVICE)
+
+        if isinstance(prompt, str):
+            prompt = [prompt]
+
+        for i, prompt_item in enumerate(prompt):
+            messages = [{"role": "user", "content": prompt_item}]
+            prompt_tokens = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
+            prompt[i] = prompt_tokens
+
+        text_inputs = self.tokenizer(prompt, max_length=self.tokenizer_max_length, padding="max_length", truncation=True, return_tensors="pt").to(AI_DEVICE)
+        prompt_masks = text_inputs.attention_mask.bool().to(AI_DEVICE)
+
+        prompt_embeds = self.text_encoder(
+            input_ids=text_inputs.input_ids,
+            attention_mask=prompt_masks,
+            output_hidden_states=True,
+        ).hidden_states[-2]
+        embedding_list = []
+        for i in range(len(prompt_embeds)):
+            extracted = prompt_embeds[i][prompt_masks[i]]
+            embedding_list.append(extracted)
+        image_info = {}
+        if self.config["task"] == "i2i" and image_list is not None:
+            vae_image_list = []
+            for image in image_list:
+                preprocessed_image = self.preprocess_image(image)
+                vae_image_list.append(preprocessed_image)
+
+            image_info = {
+                "vae_image_list": vae_image_list,
+            }
+
+        if self.cpu_offload:
+            self.text_encoder.to(torch.device("cpu"))
+            torch_device_module.empty_cache()
+            gc.collect()
+
+        return embedding_list, image_info

lightx2v/models/input_encoders/hf/z_image/qwen3_model.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class ZPreInferModuleOutput:
+    hidden_states: torch.Tensor
+    encoder_hidden_states: torch.Tensor
+    temb_img_silu: torch.Tensor
+    temb_txt_silu: torch.Tensor
+    x_freqs_cis: torch.Tensor
+    cap_freqs_cis: torch.Tensor
+    image_tokens_len: int
+    x_item_seqlens: list
+    cap_item_seqlens: list
+
+    @property
+    def adaln_input(self) -> torch.Tensor:
+        return self.temb_img_silu
+
+    @property
+    def image_rotary_emb(self) -> torch.Tensor:
+        return self.x_freqs_cis
+
+    @property
+    def freqs_cis(self) -> torch.Tensor:
+        return self.x_freqs_cis

lightx2v/models/networks/z_image/infer/module_io.py

@@ -0,0 +1,51 @@
+import torch
+
+from lightx2v.common.offload.manager import WeightAsyncStreamManager
+from lightx2v.models.networks.z_image.infer.transformer_infer import ZImageTransformerInfer
+from lightx2v_platform.base.global_var import AI_DEVICE
+
+torch_device_module = getattr(torch, AI_DEVICE)
+
+
+class ZImageOffloadTransformerInfer(ZImageTransformerInfer):
+    def __init__(self, config):
+        super().__init__(config)
+        self.phases_num = 3
+        self.num_blocks = config["num_layers"]
+        if self.config.get("cpu_offload", False):
+            if "offload_ratio" in self.config:
+                self.offload_ratio = self.config["offload_ratio"]
+            else:
+                self.offload_ratio = 1
+            offload_granularity = self.config.get("offload_granularity", "block")
+            if offload_granularity == "block":
+                if not self.config.get("lazy_load", False):
+                    self.infer_func = self.infer_with_blocks_offload
+                else:
+                    assert NotImplementedError
+
+            if offload_granularity != "model":
+                self.offload_manager = WeightAsyncStreamManager(offload_granularity=offload_granularity)
+            else:
+                assert NotImplementedError
+
+    def infer_with_blocks_offload(self, block_weights, hidden_states, encoder_hidden_states, temb, image_rotary_emb, modulate_index):
+        for block_idx in range(self.num_blocks):
+            self.block_idx = block_idx
+            if self.offload_manager.need_init_first_buffer:
+                self.offload_manager.init_first_buffer(block_weights.blocks)
+
+            self.offload_manager.prefetch_weights((block_idx + 1) % self.num_blocks, block_weights.blocks)
+            with torch_device_module.stream(self.offload_manager.compute_stream):
+                encoder_hidden_states, hidden_states = self.infer_block(
+                    block_weight=self.offload_manager.cuda_buffers[0],
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    modulate_index=modulate_index,
+                )
+
+            self.offload_manager.swap_blocks()
+
+        return encoder_hidden_states, hidden_states

lightx2v/models/networks/z_image/infer/offload/transformer_infer.py

@@ -0,0 +1,58 @@
+import torch.nn.functional as F
+
+
+class ZImagePostInfer:
+    def __init__(self, config):
+        self.config = config
+        self.cpu_offload = config.get("cpu_offload", False)
+        self.zero_cond_t = config.get("zero_cond_t", False)
+
+    def set_scheduler(self, scheduler):
+        self.scheduler = scheduler
+
+    def infer(self, weights, hidden_states, temb_img_silu, image_tokens_len=None):
+        temb_silu = F.silu(temb_img_silu)
+        temb1 = weights.norm_out_linear.apply(temb_silu)
+
+        scale = 1.0 + temb1
+        normed = weights.norm_out.apply(hidden_states)
+        scaled_norm = normed * scale.unsqueeze(1)
+        B, T, D = scaled_norm.shape
+        hidden_states_2d = scaled_norm.reshape(B * T, D)
+
+        output_2d = weights.proj_out_linear.apply(hidden_states_2d)
+        out_dim = output_2d.shape[-1]
+        output = output_2d.reshape(B, T, out_dim)
+
+        if image_tokens_len is not None:
+            output = output[:, :image_tokens_len, :]
+
+        patch_size = self.config.get("patch_size", 2)
+        f_patch_size = 1
+        transformer_out_channels = out_dim // (patch_size * patch_size * f_patch_size)
+        expected_out_dim = patch_size * patch_size * f_patch_size * transformer_out_channels
+
+        if out_dim != expected_out_dim:
+            raise ValueError(f"out_dim mismatch: {out_dim} != {expected_out_dim} (transformer_out_channels={transformer_out_channels})")
+
+        out_channels = transformer_out_channels
+        target_shape = self.scheduler.input_info.target_shape
+
+        _, _, height, width = target_shape
+        num_frames = 1
+        pH = pW = patch_size
+        pF = f_patch_size
+        F_tokens = num_frames // pF
+        H_tokens = height // pH
+        W_tokens = width // pW
+
+        expected_T = F_tokens * H_tokens * W_tokens
+        if output.shape[1] != expected_T:
+            raise ValueError(f"Token count mismatch: output.shape[1]={output.shape[1]} != expected_T={expected_T} (from target_shape={target_shape})")
+
+        output_reshaped = output.view(B, F_tokens, H_tokens, W_tokens, pF, pH, pW, out_channels)
+        output_permuted = output_reshaped.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        output_4d = output_permuted.reshape(B, out_channels, num_frames, height, width)
+        output_4d = output_4d.squeeze(2)
+
+        return output_4d