--- frameworks: - Pytorch license: Apache License 2.0 tags: [] tasks: - text-to-image-synthesis --- # Templates-锐利激发(FLUX.2-klein-base-4B) 本模型是 [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) 开源的 Diffusion Templates 系列模型之一。该模型能够通过调整 `scale` 参数,精准控制生成图像的锐度与细节表现力。 ## 效果展示 > **Prompt:** A cat is sitting on a stone. | scale=0.1 | scale=0.8 | |:---:|:---:| | ![](./assets/cat_Sharpness_0.1.jpg) | ![](./assets/cat_Sharpness_0.8.jpg) | --- > **Prompt:** A close-up of a person's eyes, looking at the camera, reflections in the pupils, highly aesthetic. | scale=0.1 | scale=0.8 | |:---:|:---:| | ![](./assets/eye_Sharpness_0.1.jpg) | ![](./assets/eye_Sharpness_0.8.jpg) | --- > **Prompt:** A beautifully decorated frosted cupcake. | scale=0.1 | scale=0.8 | |:---:|:---:| | ![](./assets/cake_Sharpness_0.1.jpg) | ![](./assets/cake_Sharpness_0.8.jpg) | ## 推理代码 * 安装 [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) ``` git clone https://github.com/modelscope/DiffSynth-Studio.git cd DiffSynth-Studio pip install -e . ``` * 直接推理,需 40G 显存 ```python from diffsynth.diffusion.template import TemplatePipeline from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig import torch pipe = Flux2ImagePipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"), ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"), ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), ], tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), ) template = TemplatePipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")], ) image = template( pipe, prompt="A cat is sitting on a stone.", seed=0, cfg_scale=4, num_inference_steps=50, template_inputs = [{"scale": 0.1}], negative_template_inputs = [{"scale": 0.5}], ) image.save("image_Sharpness_0.1.jpg") image = template( pipe, prompt="A cat is sitting on a stone.", seed=0, cfg_scale=4, num_inference_steps=50, template_inputs = [{"scale": 0.8}], negative_template_inputs = [{"scale": 0.5}], ) image.save("image_Sharpness_0.8.jpg") ``` * 开启惰性加载和显存管理,需 24G 显存 ```python from diffsynth.diffusion.template import TemplatePipeline from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig import torch vram_config = { "offload_dtype": "disk", "offload_device": "disk", "onload_dtype": torch.float8_e4m3fn, "onload_device": "cpu", "preparing_dtype": torch.float8_e4m3fn, "preparing_device": "cuda", "computation_dtype": torch.bfloat16, "computation_device": "cuda", } pipe = Flux2ImagePipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config), ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config), ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"), ], tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"), vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5, ) template = TemplatePipeline.from_pretrained( torch_dtype=torch.bfloat16, device="cuda", model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")], lazy_loading=True, ) image = template( pipe, prompt="A cat is sitting on a stone.", seed=0, cfg_scale=4, num_inference_steps=50, template_inputs = [{"scale": 0.1}], negative_template_inputs = [{"scale": 0.5}], ) image.save("image_Sharpness_0.1.jpg") image = template( pipe, prompt="A cat is sitting on a stone.", seed=0, cfg_scale=4, num_inference_steps=50, template_inputs = [{"scale": 0.8}], negative_template_inputs = [{"scale": 0.5}], ) image.save("image_Sharpness_0.8.jpg") ``` ## 训练代码 安装 DiffSynth-Studio 后,使用以下脚本可开启训练,更多信息请参考 [DiffSynth-Studio 文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/)。 ```shell modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Sharpness/*" --local_dir ./data/diffsynth_example_dataset accelerate launch examples/flux2/model_training/train.py \ --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness \ --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness/metadata.jsonl \ --extra_inputs "template_inputs" \ --max_pixels 1048576 \ --dataset_repeat 50 \ --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 "DiffSynth-Studio/Template-KleinBase4B-Sharpness:" \ --tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \ --learning_rate 1e-4 \ --num_epochs 2 \ --remove_prefix_in_ckpt "pipe.template_model." \ --output_path "./models/train/Template-KleinBase4B-Sharpness_full" \ --trainable_models "template_model" \ --use_gradient_checkpointing \ --find_unused_parameters ```