This model is one of the Diffusion Templates series models open-sourced in
DiffSynth-Studio . By adjusting the
scale parameter, this model can precisely control the sharpness and detail expressiveness of generated images.
git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio
pip install -e .
1 from diffsynth . diffusion . template import TemplatePipeline
2 from diffsynth . pipelines . flux2_image import Flux2ImagePipeline , ModelConfig
3 import torch
4
5
6 pipe = Flux2ImagePipeline . from_pretrained (
7 torch_dtype = torch . bfloat16 ,
8 device = "cuda" ,
9 model_configs = [
10 ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-base-4B" , origin_file_pattern = "transformer/*.safetensors" ) ,
11 ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-4B" , origin_file_pattern = "text_encoder/*.safetensors" ) ,
12 ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-4B" , origin_file_pattern = "vae/diffusion_pytorch_model.safetensors" ) ,
13 ] ,
14 tokenizer_config = ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-4B" , origin_file_pattern = "tokenizer/" ) ,
15 )
16 template = TemplatePipeline . from_pretrained (
17 torch_dtype = torch . bfloat16 ,
18 device = "cuda" ,
19 model_configs = [ ModelConfig ( model_id = "DiffSynth-Studio/Template-KleinBase4B-Sharpness" ) ] ,
20 )
21 image = template (
22 pipe ,
23 prompt = "A cat is sitting on a stone." ,
24 seed = 0 , cfg_scale = 4 , num_inference_steps = 50 ,
25 template_inputs = [ { "scale" : 0.1 } ] ,
26 negative_template_inputs = [ { "scale" : 0.5 } ] ,
27 )
28 image . save ( "image_Sharpness_0.1.jpg" )
29 image = template (
30 pipe ,
31 prompt = "A cat is sitting on a stone." ,
32 seed = 0 , cfg_scale = 4 , num_inference_steps = 50 ,
33 template_inputs = [ { "scale" : 0.8 } ] ,
34 negative_template_inputs = [ { "scale" : 0.5 } ] ,
35 )
36 image . save ( "image_Sharpness_0.8.jpg" )
1 from diffsynth . diffusion . template import TemplatePipeline
2 from diffsynth . pipelines . flux2_image import Flux2ImagePipeline , ModelConfig
3 import torch
4
5 vram_config = {
6 "offload_dtype" : "disk" ,
7 "offload_device" : "disk" ,
8 "onload_dtype" : torch . float8_e4m3fn ,
9 "onload_device" : "cpu" ,
10 "preparing_dtype" : torch . float8_e4m3fn ,
11 "preparing_device" : "cuda" ,
12 "computation_dtype" : torch . bfloat16 ,
13 "computation_device" : "cuda" ,
14 }
15 pipe = Flux2ImagePipeline . from_pretrained (
16 torch_dtype = torch . bfloat16 ,
17 device = "cuda" ,
18 model_configs = [
19 ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-base-4B" , origin_file_pattern = "transformer/*.safetensors" , ** vram_config ) ,
20 ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-4B" , origin_file_pattern = "text_encoder/*.safetensors" , ** vram_config ) ,
21 ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-4B" , origin_file_pattern = "vae/diffusion_pytorch_model.safetensors" ) ,
22 ] ,
23 tokenizer_config = ModelConfig ( model_id = "black-forest-labs/FLUX.2-klein-4B" , origin_file_pattern = "tokenizer/" ) ,
24 vram_limit = torch . cuda . mem_get_info ( "cuda" ) [ 1 ] / ( 1024 ** 3 ) - 0.5 ,
25 )
26 template = TemplatePipeline . from_pretrained (
27 torch_dtype = torch . bfloat16 ,
28 device = "cuda" ,
29 model_configs = [ ModelConfig ( model_id = "DiffSynth-Studio/Template-KleinBase4B-Sharpness" ) ] ,
30 lazy_loading = True ,
31 )
32 image = template (
33 pipe ,
34 prompt = "A cat is sitting on a stone." ,
35 seed = 0 , cfg_scale = 4 , num_inference_steps = 50 ,
36 template_inputs = [ { "scale" : 0.1 } ] ,
37 negative_template_inputs = [ { "scale" : 0.5 } ] ,
38 )
39 image . save ( "image_Sharpness_0.1.jpg" )
40 image = template (
41 pipe ,
42 prompt = "A cat is sitting on a stone." ,
43 seed = 0 , cfg_scale = 4 , num_inference_steps = 50 ,
44 template_inputs = [ { "scale" : 0.8 } ] ,
45 negative_template_inputs = [ { "scale" : 0.5 } ] ,
46 )
47 image . save ( "image_Sharpness_0.8.jpg" )
After installing DiffSynth-Studio, use the following script to start training. For more information, please refer to the
DiffSynth-Studio Documentation .
1 modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Sharpness/*" --local_dir ./data/diffsynth_example_dataset
2
3 accelerate launch examples/flux2/model_training/train.py \
4 --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness \
5 --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness/metadata.jsonl \
6 --extra_inputs "template_inputs" \
7 --max_pixels 1048576 \
8 --dataset_repeat 50 \
9 --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" \
10 --template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Sharpness:" \
11 --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \
12 --learning_rate 1e-4 \
13 --num_epochs 2 \
14 --remove_prefix_in_ckpt "pipe.template_model." \
15 --output_path "./models/train/Template-KleinBase4B-Sharpness_full" \
16 --trainable_models "template_model" \
17 --use_gradient_checkpointing \
18 --find_unused_parameters