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

Diffusion Templates framework

Browse Source
This commit is contained in:
Artiprocher
2026-04-08 15:25:33 +08:00
parent f88b99cb4f
commit 9f8c352a15
10 changed files with 526 additions and 241 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
examples/flux2/model_training/full/FLUX.2-klein-base-4B-skills.sh → examples/flux2/model_training/full/Template-KleinBase4B.sh
View File

@@ -1,16 +1,17 @@
accelerate launch examples/flux2/model_training/train.py \
--dataset_base_path /mnt/nas1/duanzhongjie.dzj/dataset/ImagePulseV2 \
--dataset_metadata_path /mnt/nas1/duanzhongjie.dzj/dataset/ImagePulseV2/metadata_example_ti2ti.jsonl \
--extra_inputs "skill_inputs" \
--dataset_base_path xxx \
--dataset_metadata_path xxx/metadata.jsonl \
--extra_inputs "template_inputs" \
--max_pixels 1048576 \
--dataset_repeat 1 \
--model_id_with_origin_paths "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors,black-forest-labs/FLUX.2-klein-4B:vae/diffusion_pytorch_model.safetensors" \
--skill_model_id_or_path "models/base" \
--template_model_id_or_path "xxx" \
--tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \
--learning_rate 1e-4 \
--num_epochs 999 \
--remove_prefix_in_ckpt "pipe.skill_model." \
--output_path "./models/train/FLUX.2-klein-base-4B-skills_full" \
--trainable_models "skill_model" \
--remove_prefix_in_ckpt "pipe.template_model." \
--output_path "./models/train/Template-KleinBase4B_full" \
--trainable_models "template_model" \
--save_steps 1000 \
--use_gradient_checkpointing \
--save_steps 200
--find_unused_parameters

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

@@ -18,7 +18,7 @@ class Flux2ImageTrainingModule(DiffusionTrainingModule):
extra_inputs=None,
fp8_models=None,
offload_models=None,
skill_model_id_or_path=None,
template_model_id_or_path=None,
device="cpu",
task="sft",
):
@@ -27,7 +27,7 @@ class Flux2ImageTrainingModule(DiffusionTrainingModule):
model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, device=device)
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_skill_model(self.pipe, skill_model_id_or_path)
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)
# Training mode
@@ -128,7 +128,7 @@ if __name__ == "__main__":
extra_inputs=args.extra_inputs,
fp8_models=args.fp8_models,
offload_models=args.offload_models,
skill_model_id_or_path=args.skill_model_id_or_path,
template_model_id_or_path=args.template_model_id_or_path,
task=args.task,
device="cpu" if args.initialize_model_on_cpu else accelerator.device,
)