相信尝试过训练Z-Image LoRA的开发者会发现,直接基于 Z-Image Turbo 训练出的 LoRA 会失去加速能力,在加速配置(steps=8,cfg=1)下生成的图像变得模糊,在非加速配置(steps=30,cfg=2)下生成的图像正常。
为能够在 LoRA 中保留"Turbo"加速能力,ModelScope DiffSynth团队提出一种训练 Z-Image-Turbo LoRA的增强训练解决方案, 并开源了Z-Image Turbo 加速能力修复 LoRA的权重——Z-Image-Turbo-DistillPatch!
开源地址:
https://modelscope.cn/models/DiffSynth-Studio/Z-Image-Turbo-DistillPatch
另外,魔搭AIGC专区的 Z-Image 免费训练马上上线,敬请期待!
模型介绍
Tongyi-MAI/Z-Image-Turbo 已登顶多个主流开源模型社区(包括 Hugging Face 和 ModelScope)的热门模型榜单。该模型最受社区称赞的特性之一是,作为一个蒸馏版本,它能够在少量步数内生成高质量图像。然而,这也意味着该模型的训练较为复杂,特别是当我们希望在 LoRA 中保留"Turbo"能力以实现快速图像生成时。
为促进 Z-Image-Turbo 正确进行 LoRA 训练,我们探索和比较了不同的训练方案,并提出了一种增强训练解决方案,使我们能够依赖即插即用的标准化 SFT 流程,同时不影响推理时的 Turbo 加速能力。
训练方案
Tongyi-MAI/Z-Image-Turbo 是一个基于蒸馏技术的加速生成模型,其核心优势是支持低步数推理。
训练注意事项:直接更新模型权重(如完全微调或标准 LoRA)往往会破坏模型的预训练加速轨迹,导致以下现象:
- 使用默认"加速配置"(num_inference_steps=8, cfg_scale=1)推理时,生成质量显著下降。
- 使用"非加速配置"(num_inference_steps=30, cfg_scale=2)推理时,生成质量实际上有所提升,表明模型已退化为非 Turbo 版本。
为解决这一问题,DiffSynth-Studio 提供了四种训练和推理组合策略。您可以根据对推理速度和训练成本的需求选择最合适的方案。
通用实验设置:
- 数据集:
- 训练步数:5 个 epoch * 50 次重复 = 250 步
- 验证提示词:"a dog"
方案1:标准 SFT 训练 + 非加速配置推理
这是最通用的微调方法。如果您不依赖 Turbo 模型的快速推理能力,仅关注微调后的生成质量,可以直接使用标准 SFT 脚本进行训练。
- 适用场景:对推理速度不敏感;寻求简单的训练工作流。
- 训练方法:使用标准 SFT 训练。
accelerate launch examples/z_image/model_training/train.py \ --dataset_base_path data/example_image_dataset \ --dataset_metadata_path data/example_image_dataset/metadata.csv \ --max_pixels 1048576 \ --dataset_repeat 50 \ --model_id_with_origin_paths "Tongyi-MAI/Z-Image-Turbo:transformer/*.safetensors,Tongyi-MAI/Z-Image-Turbo:text_encoder/*.safetensors,Tongyi-MAI/Z-Image-Turbo:vae/diffusion_pytorch_model.safetensors" \ --learning_rate 1e-4 \ --num_epochs 5 \ --remove_prefix_in_ckpt "pipe.dit." \ --output_path "./models/train/Z-Image-Turbo_lora" \ --lora_base_model "dit" \ --lora_target_modules "to_q,to_k,to_v,to_out.0,w1,w2,w3" \ --lora_rank 32 \ --use_gradient_checkpointing \ --dataset_num_workers 8
- 推理配置:必须放弃加速配置。请调整为 num_inference_steps=30 和 cfg_scale=2。
每个 epoch 后的结果(8步,cfg=1):
最终结果(30步,cfg=2):
方案2:差分 LoRA 训练 + 加速配置推理
如果您希望微调后的模型保留 8 步生成加速能力,推荐使用差分 LoRA 训练。该方法通过引入预设 LoRA 来锁定加速轨迹。
- 适用场景:需要保持 8 步快速推理且显存占用低。
- 训练方法:通过加载预设 LoRA(例如 ostris/zimage_turbo_training_adapter)执行差分 LoRA 训练。
- 推理配置:保持加速配置,即 num_inference_steps=8 和 cfg_scale=1。
最终结果(8步,cfg=1):
方案3:标准 SFT 训练 + 轨迹模仿蒸馏训练 + 加速配置推理
这是一种两阶段"先微调、后加速"的训练方案,旨在让模型先学习内容,然后恢复速度。
- 适用场景:需要标准 SFT 训练并恢复加速能力。
- 训练方法:首先,执行方案1的标准 SFT 训练(此时加速能力将丢失);随后,基于 SFT 模型执行轨迹模仿蒸馏训练。
- 推理配置:恢复加速配置,即 num_inference_steps=8 和 cfg_scale=1。
最终结果(8步,cfg=1):
方案4:标准 SFT 训练 + 推理时加载蒸馏加速 LoRA + 加速配置推理
该方案使用标准 SFT 进行训练,并在推理时使用外部模块(https://www.modelscope.cn/models/DiffSynth-Studio/Z-Image-Turbo-DistillPatch)来恢复加速能力。
- 适用场景:希望使用标准 SFT 工作流,或已拥有训练好的 SFT 模型并希望在不重新训练的情况下恢复其加速特性。
- 训练方法:执行方案1的标准 SFT 训练。
- 推理方法:额外加载蒸馏加速 LoRA,并使用 num_inference_steps=8 和 cfg_scale=1 的加速配置。
最终结果(8步,cfg=1):
04结论:推荐使用方案4
方案4 提供了最佳权衡:您可以保持标准 SFT 的简洁性和强大功能,同时通过在推理时加载官方 Z-Image-Turbo-DistillPatch LoRA 轻松恢复 Turbo 加速。这种即插即用的方法无需重新训练,支持现有模型,并能提供高质量的 8 步生成,使其成为最实用且可扩展的选择。
模型推理
环境安装
git clone https://github.com/modelscope/DiffSynth-Studio.git cd DiffSynth-Studio pip install -e .
推理脚本
FP8 精度量化会导致明显的图像质量劣化,因此不建议在 Z-Image Turbo 模型上开启任何量化,仅建议开启 CPU Offload,最低 8G 显存即可运行。
from diffsynth.pipelines.z_image import ZImagePipeline, ModelConfig import torch pipe = ZImagePipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="transformer/*.safetensors"), ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors"), ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), ], tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"), ) pipe.load_lora(pipe.dit, "path/to/your/lora.safetensors") pipe.load_lora(pipe.dit, ModelConfig(model_id="DiffSynth-Studio/Z-Image-Turbo-DistillPatch", origin_file_pattern="model.safetensors")) image = pipe(prompt="a dog", seed=42, rand_device="cuda") image.save("image.jpg")
https://modelscope.cn/models/DiffSynth-Studio/Z-Image-Turbo-DistillPatch