DOCSEARCHbyAlgolia
Submit Docs

Z-Image Image Generation Model

repository·main·Indexed 11 days ago

https://github.com/tongyi-mai/z-image

An efficient 6B parameter image generation foundation model family utilizing a Scalable Single-Stream Diffusion Transformer (S3-DiT) architecture. It includes variants such as Z-Image-Turbo for sub-second latency, Z-Image for high-quality aesthetics, Z-Image-Omni-Base for generation and editing, and Z-Image-Edit for natural language image-to-image tasks. Supports PyTorch native inference and integration with the diffusers library via ZImagePipeline.

Tokens
9.9K
Snippets
35
Records
47
Agent score
94%

What's inside Z-Image

  1. Overview of Z-Image Model Variants

    main

    Z-Image is a 6B parameter image generation model family using a Scalable Single-Stream DiT (S3-DiT) architecture. There are four primary variants:

    • Z-Image-Turbo: A distilled version optimized for speed. It requires only 8 NFEs (Number of Function Evaluations), offering sub-second latency on H800 GPUs and fitting in 16G VRAM consumer devices. Best for photorealism and bilingual (English/Chinese) text.
    • Z-Image: The foundation model focused on high-quality aesthetics, diversity, and controllability. Ideal for creative generation and fine-tuning.
    • Z-Image-Omni-Base: A versatile foundation model designed for both generation and editing tasks.
    • Z-Image-Edit: A variant fine-tuned specifically for image editing via natural language instructions (image-to-image).

  2. How Z-Image Model Manifests work

    main

    Model manifest files list all required files for a specific Z-Image model variant. They can optionally include MD5 checksums to ensure file integrity during loading.

    Naming Conventions:

    • z-image-turbo.txt: Used specifically for the Z-Image Turbo model.
    • {model-name}.txt: Used for any custom model variants.

    Verification Modes:

    • verify=False (Default): Only checks if the files exist. This is the faster option.
    • verify=True: Checks both file existence and verifies the MD5 checksums against the manifest. This is the thorough option for ensuring data integrity.

  3. Accelerate Z-Image inference with community tools

    main

    Several third-party tools provide acceleration and specialized inference for Z-Image:

    • Cache-DiT: Provides inference acceleration via DBCache, Context Parallelism, and Tensor Parallelism (up to ~4x speedup on 4 GPUs).
    • stable-diffusion.cpp: A pure C++ engine for fast, memory-efficient inference. Supports low-VRAM environments (as little as 4GB VRAM) across CUDA, Vulkan, etc.
    • LeMiCa: A training-free, timestep-level acceleration method.
    • MeanCache: A training-free acceleration method for Flow Matching models, offering up to 3.7x speedup with plug-and-play integration.
    • SGLang-Diffusion: Accelerates image and video generation using SGLang's performance optimizations.
    • vllm-omni: Supports fast inference and serving for Z-Image.
    • Candle: A minimalist ML framework in Rust that supports Z-Image.

  4. Optimize Z-Image performance with bfloat16 and Flash Attention

    main

    To achieve optimal performance when running Z-Image on supported GPUs, use bfloat16 precision. Additionally, you can improve efficiency by switching the attention backend from the default SDPA to Flash Attention-2 using the set_attention_backend method on the transformer component of the pipeline.

    # Use bfloat16 for optimal performance on supported GPUs
# [Optional] Switch to Flash Attention for better efficiency
pip.transformer.set_attention_backend("flash")    # Enable Flash-Attention-2

  5. Quick Start: PyTorch Native Inference

    main

    To run inference using the native PyTorch implementation, install the project in editable mode and run the provided inference script.

    # Install dependencies
pip install -e .

# Run inference
python inference.py

  6. Quick Start: Diffusers Inference (Z-Image Foundation)

    main

    To use the standard Z-Image foundation model with diffusers, install diffusers from source.

    Recommended Parameters for Z-Image:

    • Resolution: 512×512 to 2048×2048.
    • Guidance scale: 3.0 – 5.0.
    • Inference steps: 28 – 50.
    • Negative prompts: Highly recommended for better control.
    • CFG normalization: Set to False for general stylism, or True for realism.
    import torch
from diffusers import ZImagePipeline

# Load the pipeline
pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
pipe.to("cuda")

# Generate image
image = pipe(
    prompt="Your prompt here",
    negative_prompt="",
    height=1280,
    width=720,
    cfg_normalization=False,
    num_inference_steps=50,
    guidance_scale=4,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]

image.save("example.png")

  7. Generate a Model Manifest

    main

    You can generate manifest files for your model directories using the src.tools.generate_manifest module. The standard format includes MD5 checksums and is recommended for integrity verification.

    # Generate with MD5 checksums (auto-saves to this directory)
python -m src.tools.generate_manifest ckpts/Z-Image-Turbo

# Generate without checksums (faster, not recommended)
python -m src.tools.generate_manifest ckpts/Z-Image-Turbo --no-checksums

# With verbose output
python -m src.tools.generate_manifest ckpts/Z-Image-Turbo --verbose

# Custom output path
python -m src.tools.generate_manifest ckpts/Z-Image-Turbo --output custom.txt

  8. Quick Start: Diffusers Inference (Z-Image-Turbo)

    main

    To use Z-Image-Turbo with the diffusers library, you must install diffusers from source to ensure support for the Z-Image architecture.

    Key Configuration for Turbo:

    • dtype: Use torch.bfloat16 for optimal performance.
    • guidance_scale: Must be set to 0.0 for Turbo models.
    • num_inference_steps: Typically set to 9 (which results in 8 DiT forwards).
    import torch
from diffusers import ZImagePipeline

# 1. Load the pipeline
pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
pipe.to("cuda")

# [Optional] Enable Flash-Attention-2
# pipe.transformer.set_attention_backend("flash")

prompt = "Your prompt here"

# 2. Generate Image
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=9,
    guidance_scale=0.0,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]

image.save("example.png")

  9. Use Z-Image in specialized environments

    main

    Z-Image is supported by various specialized frameworks and interfaces:

    • ComfyUI: Use ComfyUI ZImageLatent for official Z-Image resolution latents.
    • DiffSynth-Studio: Supports LoRA training, full training, distillation training, and low-VRAM inference.
    • Rust: Supported via the Candle ML framework.

  10. How Z-Image Transformer layers work (Refinement vs Unified)

    main

    The architecture uses three distinct stages of processing:

    1. Noise Refiner (noise_refiner): Uses ZImageTransformerBlock with modulation=True. It uses Adaptive Layer Norm (AdaLN) to modulate features based on the timestep t.
    2. Context Refiner (context_refiner): Uses ZImageTransformerBlock with modulation=False. It processes caption features without timestep modulation.
    3. Unified Layers (layers): Combines image and caption features into a single sequence. These layers use modulation=True (AdaLN) to integrate timestep information into the unified sequence.

    This multi-stage approach allows the model to refine noise and context independently before performing joint reasoning.

  12. Load model weights using ensure_model_weights

    main

    Use the ensure_model_weights function to load model weights. The function can automatically detect the correct manifest based on the directory name, or you can specify a manifest explicitly.

    Auto-detection: If the model directory is named Z-Image-Turbo, the system automatically looks for z-image-turbo.txt.

    Explicit Manifest: If using a custom model or non-standard naming, pass the manifest_name argument.

    # Auto-detects manifest from "Z-Image-Turbo" -> uses z-image-turbo.txt
model_path = ensure_model_weights("ckpts/Z-Image-Turbo")

# Explicit manifest
model_path = ensure_model_weights("ckpts/Z-Image-Turbo", manifest_name="z-image-turbo.txt")