Z-Image-Turbo_fp32-fp16-bf16_full_and_ema-only

✨ Z-Image-Turbo FP32 / FP16 / BF16 EMA-ONLY & FULL

Multiple versions of Z-Image-Turbo model in various precisions and configurations, prepared directly from the original Tongyi-MAI/Z-Image-Turbo repository.

📦 Available Variants

EMA vs Full - Which to Choose?

• EMA-only: Contains only Exponential Moving Average parameters - averaged weights from training process. Provides more stable and better results during image generation, smaller file size. Use this for inference.

• Full: Contains all parameters (training + EMA). Only needed if you want to continue training the model.

🔧 Preparation Process

Models were processed using:

1. merge-safetensors - merging split transformer parts into single *.safetensors file (placed in transformer directory)

2. PyTorch-Precision-Converter - converting precision from FP32 to FP16/BF16 and creating EMA-only variants

💡 For Diffusers Users

⚠️ This is NOT compatible with ComfyUI - models are prepared for diffusers library.

Required File Names

To use with ZImagePipeline without specifying full paths, rename model files in appropriate folders:

text_encoder/
  └── model.safetensors              # Text encoder

transformer/
  └── diffusion_pytorch_model.safetensors   # Transformer

vae/
  └── diffusion_pytorch_model.safetensors   # VAE

Example Usage (based on one from original repo)

pip install git+https://github.com/huggingface/diffusers

Example /w repo+models downloader

import argparse
import os
from huggingface_hub import snapshot_download, hf_hub_download


REPO_ID = "tsqn/Z-Image-Turbo_fp32-fp16-bf16_full_and_ema-only"

def main(local_dir):
    TRANSFORMER_DIR = f"{local_dir}\\transformer"
    TEXT_ENCODER_DIR = f"{local_dir}\\text_encoder"
    VAE_DIR = f"{local_dir}\\vae"

    def _download_model_files():
        snapshot_download(repo_id=REPO_ID, ignore_patterns="*.safetensors", local_dir=local_dir)
        hf_hub_download(repo_id=REPO_ID, filename="diffusion_pytorch_model-ema-only-fp32.safetensors", local_dir=local_dir)
        hf_hub_download(repo_id=REPO_ID, subfolder="text_encoder", filename="qwen_3_4b_bf16.safetensors", local_dir=local_dir)
        hf_hub_download(repo_id=REPO_ID, subfolder="vae", filename="ae_bf16.safetensors", local_dir=local_dir)

    def _rename_model_files():
        os.rename(f"{TRANSFORMER_DIR}\\diffusion_pytorch_model-ema-only-fp32.safetensors", f"{TRANSFORMER_DIR}\\diffusion_pytorch_model.safetensors")
        os.rename(f"{TEXT_ENCODER_DIR}\\qwen_3_4b_bf16.safetensors", f"{TEXT_ENCODER_DIR}\\model.safetensors")
        os.rename(f"{VAE_DIR}\\ae_bf16.safetensors", f"{VAE_DIR}\\diffusion_pytorch_model.safetensors")

    try:
        _download_model_files()
    except:
        print("ERROR when downloading model files!")
        raise Exception
    
    try:
        _rename_model_files()
    except:
        print("ERROR when renaming model files!")
        raise Exception

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--local_dir", default=None, type=str, required=True, help="Whether to save repository with model files"
    )

    args = parser.parse_args()
    main(args.local_dir)

import argparse
import torch
from diffusers import ZImagePipeline, ZImageTransformer2DModel, AutoencoderKL, FlowMatchEulerDiscreteScheduler
from transformers import Qwen3Model, Qwen2Tokenizer


def setup_zimage_pipeline(model_path, model_cpu_offload=True):
    def _setup_pipeline_components():
        vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae", torch_dtype=torch.bfloat16)
        text_encoder = Qwen3Model.from_pretrained(model_path, subfolder="text_encoder", torch_dtype=torch.bfloat16)
        tokenizer = Qwen2Tokenizer.from_pretrained(model_path, subfolder="tokenizer")
        transformer = ZImageTransformer2DModel.from_pretrained(model_path, subfolder="transformer", torch_dtype=torch.float32)
        return {
            "vae": vae,
            "text_encoder": text_encoder,
            "tokenizer": tokenizer,
            "transformer": transformer
        }
    pipeline = ZImagePipeline.from_pretrained(
        model_path,
        torch_dtype=torch.float32,
        low_cpu_mem_usage=False,
        **_setup_pipeline_components()
    )
    pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipeline.scheduler.config)
    
    if model_cpu_offload:
        pipeline.enable_model_cpu_offload()
    return pipeline

def generate_image(pipe, prompt, height, width, num_inference_steps, guidance_scale, seed, output_save_path):
    with torch.inference_mode():
        image = pipe(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=torch.Generator("cuda").manual_seed(seed),
        ).images[0]

        if output_save_path:
            image.save(f"{output_save_path}\\example_{seed}.png")
        else:
            image.save(f"example_{seed}.png")
        torch.cuda.empty_cache()

def main(args):
    pipeline = setup_zimage_pipeline(args.local_dir)
    generate_image(
        pipe=pipeline,
        prompt=args.prompt,
        width=args.width,
        height=args.height,
        num_inference_steps=args.num_inference_steps,
        guidance_scale=args.guidance_scale,
        seed=args.seed,
        output_save_path=args.output_save_path
    )

if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--local_dir", default=None, type=str, required=True, help="Path to the zimage diffusers local repository"
    )
    parser.add_argument(
        "--prompt", default="Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.", type=str, required=False, help="Prompt used to generate image."
    )
    parser.add_argument(
        "--width", default=1024, type=int, required=False, help="Width of the generated image."
    )
    parser.add_argument(
        "--height", default=1024, type=int, required=False, help="Height of the generated image."
    )
    parser.add_argument(
        "--num_inference_steps", default=9, type=int, required=False, help="Number of the inference steps."
    )
    parser.add_argument(
        "--guidance_scale", default=0.0, type=float, required=False, help="Guidance scale setting."
    )
    parser.add_argument(
        "--seed", default=42, type=int, required=False, help="Random seed value."
    )
    parser.add_argument(
        "--output_save_path",
        default=None,
        type=str,
        required=False,
        help="Path to the directory for generated image.",
    )

    args = parser.parse_args()
    main(args)

python download_repo.py --local_dir "C:\\zimage"
python generate_image.py --local_dir "C:\\zimage"

with custom parameters:

python generate_image.py --local_dir "C:\\zimage" --prompt "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights." --width 1024 --height 1042 --num_inference_steps 9 --guidance_scale 0.0 --seed 42 --output_save_path "C:\\zimage_generations"

Example 2

import torch
from diffusers import ZImagePipeline

# 1. Load the pipeline
# Use bfloat16 for optimal performance on supported GPUs
pipe = ZImagePipeline.from_pretrained(
    "path/to/model_files_main_dir",
    torch_dtype=torch.float32, # or torch.bfloat16 / torch.float16
    low_cpu_mem_usage=False,
)
pipe.to("cuda")

# [Optional] Attention Backend
# Diffusers uses SDPA by default. Switch to Flash Attention for better efficiency if supported:
# pipe.transformer.set_attention_backend("flash")    # Enable Flash-Attention-2
# pipe.transformer.set_attention_backend("_flash_3") # Enable Flash-Attention-3

# [Optional] Model Compilation
# Compiling the DiT model accelerates inference, but the first run will take longer to compile.
# pipe.transformer.compile()

# [Optional] CPU Offloading
# Enable CPU offloading for memory-constrained devices.
# pipe.enable_model_cpu_offload()

prompt = "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."

# 2. Generate Image
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(42),
).images[0]

image.save("example.png")

Example 3

import torch
from diffusers import ZImagePipeline, ZImageTransformer2DModel, AutoencoderKL, FlowMatchEulerDiscreteScheduler
from transformers import Qwen3Model, Qwen2Tokenizer


MODEL_PATH = "tsqn/Z-Image-Turbo_fp32-fp16-bf16_full_and_ema-only"

vae = AutoencoderKL.from_pretrained(MODEL_PATH, subfolder="vae", torch_dtype=torch.bfloat16)
text_encoder = Qwen3Model.from_pretrained(MODEL_PATH, subfolder="text_encoder", torch_dtype=torch.bfloat16)
tokenizer = Qwen2Tokenizer.from_pretrained(MODEL_PATH, subfolder="tokenizer")
transformer = ZImageTransformer2DModel.from_pretrained(MODEL_PATH, subfolder="transformer", torch_dtype=torch.float32)

pipe = ZImagePipeline.from_pretrained(
    MODEL_PATH,
    vae=vae,
    text_encoder=text_encoder,
    tokenizer=tokenizer,
    transformer=transformer,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=False,
)
pipe.scheduler = FlowMatchEulerDiscreteScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()

prompt = "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."

with torch.inference_mode():
    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")
    torch.cuda.empty_cache()

🎯 Recommendations

• RTX 3060 and similar: Use BF16 or FP16 for optimal performance
• Less than 12GB VRAM: FP16 EMA-only
• 12GB+ VRAM: BF16 EMA-only (better numerical stability)
• Training: FP32 Full

📝 License

Same as original Z-Image-Turbo model.

README was generated with a help of AI

@article{team2025zimage,
  title={Z-Image: An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer},
  author={Z-Image Team},
  journal={arXiv preprint arXiv:2511.22699},
  year={2025}
}

@article{liu2025decoupled,
  title={Decoupled DMD: CFG Augmentation as the Spear, Distribution Matching as the Shield},
  author={Dongyang Liu and Peng Gao and David Liu and Ruoyi Du and Zhen Li and Qilong Wu and Xin Jin and Sihan Cao and Shifeng Zhang and Hongsheng Li and Steven Hoi},
  journal={arXiv preprint arXiv:2511.22677},
  year={2025}
}

@article{jiang2025distribution,
  title={Distribution Matching Distillation Meets Reinforcement Learning},
  author={Jiang, Dengyang and Liu, Dongyang and Wang, Zanyi and Wu, Qilong and Jin, Xin and Liu, David and Li, Zhen and Wang, Mengmeng and Gao, Peng and Yang, Harry},
  journal={arXiv preprint arXiv:2511.13649},
  year={2025}
}