theluyuan/DiffSynth-Studio
1
0
Fork 0
mirror of https://github.com/modelscope/DiffSynth-Studio.git synced 2026-04-15 14:48:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

update template framework

Browse Source
This commit is contained in:
Artiprocher
2026-04-15 14:07:51 +08:00
parent 9f8c352a15
commit 59b4bbb62c
7 changed files with 85 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
examples/flux2/model_training/train.py
View File

@@ -19,6 +19,7 @@ class Flux2ImageTrainingModule(DiffusionTrainingModule):
fp8_models=None,
offload_models=None,
template_model_id_or_path=None,
enable_lora_hot_loading=False,
device="cpu",
task="sft",
):
@@ -29,6 +30,7 @@ class Flux2ImageTrainingModule(DiffusionTrainingModule):
self.pipe = Flux2ImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config)
self.pipe = self.load_training_template_model(self.pipe, template_model_id_or_path, args.use_gradient_checkpointing, args.use_gradient_checkpointing_offload)
self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model)
if enable_lora_hot_loading: self.pipe.dit = self.pipe.enable_lora_hot_loading(self.pipe.dit)
# Training mode
self.switch_pipe_to_training_mode(
@@ -129,6 +131,7 @@ if __name__ == "__main__":
fp8_models=args.fp8_models,
offload_models=args.offload_models,
template_model_id_or_path=args.template_model_id_or_path,
enable_lora_hot_loading=args.enable_lora_hot_loading,
task=args.task,
device="cpu" if args.initialize_model_on_cpu else accelerator.device,
)