support direct distill

examples/qwen_image/model_training/lora/Qwen-Image-Distill-LoRA.sh

@@ -0,0 +1,24 @@
+accelerate launch examples/qwen_image/model_training/train.py \
+  --dataset_base_path data/example_image_dataset \
+  --dataset_metadata_path data/example_image_dataset/metadata_distill_qwen_image.csv \
+  --data_file_keys "image" \
+  --extra_inputs "seed,rand_device,num_inference_steps,cfg_scale" \
+  --height 1328 \
+  --width 1328 \
+  --dataset_repeat 50 \
+  --model_id_with_origin_paths "Qwen/Qwen-Image:transformer/diffusion_pytorch_model*.safetensors,Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
+  --learning_rate 1e-4 \
+  --num_epochs 5 \
+  --remove_prefix_in_ckpt "pipe.dit." \
+  --output_path "./models/train/Qwen-Image-Distill-LoRA_lora" \
+  --lora_base_model "dit" \
+  --lora_target_modules "to_q,to_k,to_v,add_q_proj,add_k_proj,add_v_proj,to_out.0,to_add_out,img_mlp.net.2,img_mod.1,txt_mlp.net.2,txt_mod.1" \
+  --lora_rank 32 \
+  --use_gradient_checkpointing \
+  --dataset_num_workers 8 \
+  --find_unused_parameters \
+  --task direct_distill
+
+# This is an experimental training feature designed to directly distill the model, enabling generation results with fewer steps to approximate those achieved with more steps.
+# The model (https://www.modelscope.cn/models/DiffSynth-Studio/Qwen-Image-Distill-LoRA) is trained using this script.
+# The sample dataset is provided solely to demonstrate the dataset format. For actual usage, please construct a larger dataset using the base model.

examples/qwen_image/model_training/train.py

@@ -19,6 +19,7 @@ class QwenImageTrainingModule(DiffusionTrainingModule):
         use_gradient_checkpointing_offload=False,
         extra_inputs=None,
         enable_fp8_training=False,
+        task="sft",
     ):
         super().__init__()
         # Load models
@@ -38,6 +39,7 @@ class QwenImageTrainingModule(DiffusionTrainingModule):
         self.use_gradient_checkpointing = use_gradient_checkpointing
         self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
         self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else []
+        self.task = task
 
     
     def forward_preprocess(self, data):
@@ -82,11 +84,21 @@ class QwenImageTrainingModule(DiffusionTrainingModule):
     
     
     def forward(self, data, inputs=None, return_inputs=False):
+        # Inputs
         if inputs is None: inputs = self.forward_preprocess(data)
         else: inputs = self.transfer_data_to_device(inputs, self.pipe.device)
         if return_inputs: return inputs
-        models = {name: getattr(self.pipe, name) for name in self.pipe.in_iteration_models}
-        loss = self.pipe.training_loss(**models, **inputs)
+        
+        # Loss
+        if self.task == "sft":
+            models = {name: getattr(self.pipe, name) for name in self.pipe.in_iteration_models}
+            loss = self.pipe.training_loss(**models, **inputs)
+        elif self.task == "data_process":
+            loss = inputs
+        elif self.task == "direct_distill":
+            loss = self.pipe.direct_distill_loss(**inputs)
+        else:
+            raise NotImplementedError(f"Unsupported task: {self.task}.")
         return loss
 
 
@@ -122,10 +134,12 @@ if __name__ == "__main__":
         use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
         extra_inputs=args.extra_inputs,
         enable_fp8_training=args.enable_fp8_training,
+        task=args.task,
     )
     model_logger = ModelLogger(args.output_path, remove_prefix_in_ckpt=args.remove_prefix_in_ckpt)
     launcher_map = {
         "sft": launch_training_task,
-        "data_process": launch_data_process_task
+        "data_process": launch_data_process_task,
+        "direct_distill": launch_training_task,
     }
     launcher_map[args.task](dataset, model, model_logger, args=args)

examples/qwen_image/model_training/validate_lora/Qwen-Image-Distill-LoRA.py

@@ -0,0 +1,23 @@
+from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig
+import torch
+
+
+pipe = QwenImagePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="transformer/diffusion_pytorch_model*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="text_encoder/model*.safetensors"),
+        ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
+    ],
+    tokenizer_config=ModelConfig(model_id="Qwen/Qwen-Image", origin_file_pattern="tokenizer/"),
+)
+pipe.load_lora(pipe.dit, "models/train/Qwen-Image-Distill-LoRA_lora/epoch-4.safetensors")
+prompt = "精致肖像，水下少女，蓝裙飘逸，发丝轻扬，光影透澈，气泡环绕，面容恬静，细节精致，梦幻唯美。"
+image = pipe(
+    prompt,
+    seed=0,
+    num_inference_steps=4,
+    cfg_scale=1,
+)
+image.save("image.jpg")