• Stefy Lanza (nextime / spora )'s avatar
    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-noise (Z-Image negates the model output);
    timesteps sampled from the discrete turbo schedule (set_timesteps mu-shift) to keep
    the distilled model sharp; save via ZImagePipeline.save_lora_weights (transformer.
    keys) so _apply_loras' load_lora_weights applies it.
    
    _train_lora_sync routes a Z-Image DiT base to _train_dit (detected via
    model_index/transformer config _class_name); other DiTs (Flux/SD3) still raise with
    guidance. v1 — needs one validation train; recipe + knobs in
    docs/zimage-lora-training.md.
    Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
    Claude-Session: https://claude.ai/code/session_01RdMufYvtTbtGDWsiZVoXce
    1f586669
zimage-lora-training.md 4.22 KB