Diffusion Templates framework

examples/flux2/model_training/full/Template-KleinBase4B.sh

@@ -0,0 +1,17 @@
+accelerate launch examples/flux2/model_training/train.py \
+  --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" \
+  --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.template_model." \
+  --output_path "./models/train/Template-KleinBase4B_full" \
+  --trainable_models "template_model" \
+  --save_steps 1000 \
+  --use_gradient_checkpointing \
+  --find_unused_parameters