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

add a new model

Browse Source
This commit is contained in:
Artiprocher
2026-04-20 10:56:29 +08:00
parent f58ba5a784
commit 13f2618da2
19 changed files with 433 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -29,7 +29,7 @@ class Flux2ImageTrainingModule(DiffusionTrainingModule):
tokenizer_config = self.parse_path_or_model_id(tokenizer_path, default_value=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"))
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)
self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model, remove_unnecessary_params=True)
if enable_lora_hot_loading: self.pipe.dit = self.pipe.enable_lora_hot_loading(self.pipe.dit)
# Training mode