Skip to content
  • Stefy Lanza (nextime / spora )'s avatar
    1f586669
    loras: native Z-Image (DiT) LoRA training (_train_dit) · 1f586669
    Stefy Lanza (nextime / spora ) authored
    coderai's trainer only targeted SDXL/SD1.x U-Nets, producing unet.-prefixed LoRAs
    that silently no-op on Z-Image's DiT (ZImageTransformer2DModel, transformer. keys,
    Qwen3 text encoder) — "No LoRA keys associated to ZImageTransformer2DModel found
    with prefix='transformer'". The lora_train_base_model=SDXL workaround can't make a
    LoRA that loads on Z-Image.
    
    Add _train_dit: a native flow-matching DiT LoRA trainer for Z-Image, reusing the
    existing job/progress/checkpoint/queue infra (modeled on the Wan video DiT trainer,
    which is video-only). All deps (ZImageTransformer2DModel, ZImagePipeline, PEFT,
    Qwen3) are already in the main venv — no ai-toolkit, no separate venv.
    
    Reverse-engineered from diffusers pipeline_z_image.py so training matches inference:
    chat-template + Qwen3 hidden_states[-2] masked per-sample embed list; AutoencoderKL
    latents scaled (lat-shift)*scale; list-based transformer I/O with normalized
    timestep (1000-t)/1000; RAW target = x0-...
    1f586669
    loras: native Z-Image (DiT) LoRA training (_train_dit)
    Stefy Lanza (nextime / spora ) authored
    coderai's trainer only targeted SDXL/SD1.x U-Nets, producing unet.-prefixed LoRAs
    that silently no-op on Z-Image's DiT (ZImageTransformer2DModel, transformer. keys,
    Qwen3 text encoder) — "No LoRA keys associated to ZImageTransformer2DModel found
    with prefix='transformer'". The lora_train_base_model=SDXL workaround can't make a
    LoRA that loads on Z-Image.
    
    Add _train_dit: a native flow-matching DiT LoRA trainer for Z-Image, reusing the
    existing job/progress/checkpoint/queue infra (modeled on the Wan video DiT trainer,
    which is video-only). All deps (ZImageTransformer2DModel, ZImagePipeline, PEFT,
    Qwen3) are already in the main venv — no ai-toolkit, no separate venv.
    
    Reverse-engineered from diffusers pipeline_z_image.py so training matches inference:
    chat-template + Qwen3 hidden_states[-2] masked per-sample embed list; AutoencoderKL
    latents scaled (lat-shift)*scale; list-based transformer I/O with normalized
    timestep (1000-t)/1000; RAW target = x0-...
Loading