merge data process to training script

Author: Artiprocher

diffsynth/trainers/utils.py

@@ -520,14 +520,26 @@ def launch_training_task(
     dataset: torch.utils.data.Dataset,
     model: DiffusionTrainingModule,
     model_logger: ModelLogger,
-    optimizer: torch.optim.Optimizer,
-    scheduler: torch.optim.lr_scheduler.LRScheduler,
+    learning_rate: float = 1e-5,
+    weight_decay: float = 1e-2,
     num_workers: int = 8,
     save_steps: int = None,
     num_epochs: int = 1,
     gradient_accumulation_steps: int = 1,
     find_unused_parameters: bool = False,
+    args = None,
 ):
+    if args is not None:
+        learning_rate = args.learning_rate
+        weight_decay = args.weight_decay
+        num_workers = args.dataset_num_workers
+        save_steps = args.save_steps
+        num_epochs = args.num_epochs
+        gradient_accumulation_steps = args.gradient_accumulation_steps
+        find_unused_parameters = args.find_unused_parameters
+    
+    optimizer = torch.optim.AdamW(model.trainable_modules(), lr=learning_rate, weight_decay=weight_decay)
+    scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
     dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=num_workers)
     accelerator = Accelerator(
         gradient_accumulation_steps=gradient_accumulation_steps,
@@ -557,8 +569,12 @@ def launch_data_process_task(
     model: DiffusionTrainingModule,
     model_logger: ModelLogger,
     num_workers: int = 8,
+    args = None,
 ):
-    dataloader = torch.utils.data.DataLoader(dataset, shuffle=True, collate_fn=lambda x: x[0], num_workers=num_workers)
+    if args is not None:
+        num_workers = args.dataset_num_workers
+        
+    dataloader = torch.utils.data.DataLoader(dataset, shuffle=False, collate_fn=lambda x: x[0], num_workers=num_workers)
     accelerator = Accelerator()
     model, dataloader = accelerator.prepare(model, dataloader)
 
@@ -568,7 +584,7 @@ def launch_data_process_task(
                 folder = os.path.join(model_logger.output_path, str(accelerator.process_index))
                 os.makedirs(folder, exist_ok=True)
                 save_path = os.path.join(model_logger.output_path, str(accelerator.process_index), f"{data_id}.pth")
-                data = model(data)
+                data = model(data, return_inputs=True)
                 torch.save(data, save_path)
 
 
@@ -671,4 +687,5 @@ def qwen_image_parser():
     parser.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay.")
     parser.add_argument("--processor_path", type=str, default=None, help="Path to the processor. If provided, the processor will be used for image editing.")
     parser.add_argument("--enable_fp8_training", default=False, action="store_true", help="Whether to enable FP8 training. Only available for LoRA training on a single GPU.")
+    parser.add_argument("--task", type=str, default="sft", required=False, help="Task type.")
     return parser

examples/flux/model_training/train.py

@@ -25,7 +25,7 @@ def __init__(
 
         # Training mode
         self.switch_pipe_to_training_mode(
-            self, self.pipe, trainable_models,
+            self.pipe, trainable_models,
             lora_base_model, lora_target_modules, lora_rank, lora_checkpoint=lora_checkpoint,
             enable_fp8_training=False,
         )

examples/qwen_image/model_training/lora/Qwen-Image-Splited.sh

@@ -1,11 +1,12 @@
-accelerate launch examples/qwen_image/model_training/train_data_process.py \
+accelerate launch examples/qwen_image/model_training/train.py \
   --dataset_base_path data/example_image_dataset \
   --dataset_metadata_path data/example_image_dataset/metadata.csv \
   --max_pixels 1048576 \
   --model_id_with_origin_paths "Qwen/Qwen-Image:text_encoder/model*.safetensors,Qwen/Qwen-Image:vae/diffusion_pytorch_model.safetensors" \
   --output_path "./models/train/Qwen-Image_lora_cache" \
   --use_gradient_checkpointing \
-  --dataset_num_workers 8
+  --dataset_num_workers 8 \
+  --task data_process
 
 accelerate launch examples/qwen_image/model_training/train.py \
   --dataset_base_path models/train/Qwen-Image_lora_cache \

examples/qwen_image/model_training/train.py

@@ -2,7 +2,7 @@
 from diffsynth import load_state_dict
 from diffsynth.pipelines.qwen_image import QwenImagePipeline, ModelConfig
 from diffsynth.pipelines.flux_image_new import ControlNetInput
-from diffsynth.trainers.utils import DiffusionTrainingModule, ModelLogger, launch_training_task, qwen_image_parser
+from diffsynth.trainers.utils import DiffusionTrainingModule, ModelLogger, qwen_image_parser, launch_training_task, launch_data_process_task
 from diffsynth.trainers.unified_dataset import UnifiedDataset
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
@@ -29,7 +29,7 @@ def __init__(
 
         # Training mode
         self.switch_pipe_to_training_mode(
-            self, self.pipe, trainable_models,
+            self.pipe, trainable_models,
             lora_base_model, lora_target_modules, lora_rank, lora_checkpoint=lora_checkpoint,
             enable_fp8_training=enable_fp8_training,
         )
@@ -81,9 +81,10 @@ def forward_preprocess(self, data):
         return {**inputs_shared, **inputs_posi}
 
 
-    def forward(self, data, inputs=None):
+    def forward(self, data, inputs=None, return_inputs=False):
         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)
         return loss
@@ -123,13 +124,8 @@ def forward(self, data, inputs=None):
         enable_fp8_training=args.enable_fp8_training,
     )
     model_logger = ModelLogger(args.output_path, remove_prefix_in_ckpt=args.remove_prefix_in_ckpt)
-    optimizer = torch.optim.AdamW(model.trainable_modules(), lr=args.learning_rate, weight_decay=args.weight_decay)
-    scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
-    launch_training_task(
-        dataset, model, model_logger, optimizer, scheduler,
-        num_epochs=args.num_epochs,
-        gradient_accumulation_steps=args.gradient_accumulation_steps,
-        save_steps=args.save_steps,
-        find_unused_parameters=args.find_unused_parameters,
-        num_workers=args.dataset_num_workers,
-    )
+    launcher_map = {
+        "sft": launch_training_task,
+        "data_process": launch_data_process_task
+    }
+    launcher_map[args.task](dataset, model, model_logger, args=args)

examples/qwen_image/model_training/train_data_process.py

@@ -26,7 +26,7 @@ def __init__(
 
         # Training mode
         self.switch_pipe_to_training_mode(
-            self, self.pipe, trainable_models,
+            self.pipe, trainable_models,
             lora_base_model, lora_target_modules, lora_rank, lora_checkpoint=lora_checkpoint,
             enable_fp8_training=False,
         )