Upload folder using huggingface_hub
Browse files- .gitattributes +3 -0
- README.md +159 -0
- README_from_modelscope.md +163 -0
- assets/cake_Sharpness_0.1.jpg +0 -0
- assets/cake_Sharpness_0.8.jpg +0 -0
- assets/cat_Sharpness_0.1.jpg +3 -0
- assets/cat_Sharpness_0.8.jpg +3 -0
- assets/eye_Sharpness_0.1.jpg +0 -0
- assets/eye_Sharpness_0.8.jpg +3 -0
- configuration.json +1 -0
- model.py +76 -0
- model.safetensors +3 -0
- scores.json +0 -0
.gitattributes
CHANGED
|
@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
assets/cat_Sharpness_0.1.jpg filter=lfs diff=lfs merge=lfs -text
|
| 37 |
+
assets/cat_Sharpness_0.8.jpg filter=lfs diff=lfs merge=lfs -text
|
| 38 |
+
assets/eye_Sharpness_0.8.jpg filter=lfs diff=lfs merge=lfs -text
|
README.md
ADDED
|
@@ -0,0 +1,159 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: apache-2.0
|
| 3 |
+
---
|
| 4 |
+
# Templates - Sharpness Activation (FLUX.2-klein-base-4B)
|
| 5 |
+
|
| 6 |
+
This model is one of the Diffusion Templates series models open-sourced in [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio). By adjusting the `scale` parameter, this model can precisely control the sharpness and detail expressiveness of generated images.
|
| 7 |
+
|
| 8 |
+
## Result Examples
|
| 9 |
+
|
| 10 |
+
> **Prompt:** A cat is sitting on a stone.
|
| 11 |
+
|
| 12 |
+
| scale=0.1 | scale=0.8 |
|
| 13 |
+
|:---:|:---:|
|
| 14 |
+
|  |  |
|
| 15 |
+
|
| 16 |
+
---
|
| 17 |
+
|
| 18 |
+
> **Prompt:** A close-up of a person's eyes, looking at the camera, reflections in the pupils, highly aesthetic.
|
| 19 |
+
|
| 20 |
+
| scale=0.1 | scale=0.8 |
|
| 21 |
+
|:---:|:---:|
|
| 22 |
+
|  |  |
|
| 23 |
+
|
| 24 |
+
---
|
| 25 |
+
|
| 26 |
+
> **Prompt:** A beautifully decorated frosted cupcake.
|
| 27 |
+
|
| 28 |
+
| scale=0.1 | scale=0.8 |
|
| 29 |
+
|:---:|:---:|
|
| 30 |
+
|  |  |
|
| 31 |
+
|
| 32 |
+
## Inference Code
|
| 33 |
+
|
| 34 |
+
* Install [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio)
|
| 35 |
+
|
| 36 |
+
```
|
| 37 |
+
git clone https://github.com/modelscope/DiffSynth-Studio.git
|
| 38 |
+
cd DiffSynth-Studio
|
| 39 |
+
pip install -e .
|
| 40 |
+
```
|
| 41 |
+
|
| 42 |
+
* Direct inference (requires 40G GPU memory)
|
| 43 |
+
|
| 44 |
+
```python
|
| 45 |
+
from diffsynth.diffusion.template import TemplatePipeline
|
| 46 |
+
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
|
| 47 |
+
import torch
|
| 48 |
+
```
|
| 49 |
+
|
| 50 |
+
pipe = Flux2ImagePipeline.from_pretrained(
|
| 51 |
+
torch_dtype=torch.bfloat16,
|
| 52 |
+
device="cuda",
|
| 53 |
+
model_configs=[
|
| 54 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"),
|
| 55 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"),
|
| 56 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
|
| 57 |
+
],
|
| 58 |
+
tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"),
|
| 59 |
+
)
|
| 60 |
+
template = TemplatePipeline.from_pretrained(
|
| 61 |
+
torch_dtype=torch.bfloat16,
|
| 62 |
+
device="cuda",
|
| 63 |
+
model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")],
|
| 64 |
+
)
|
| 65 |
+
image = template(
|
| 66 |
+
pipe,
|
| 67 |
+
prompt="A cat is sitting on a stone.",
|
| 68 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 69 |
+
template_inputs = [{"scale": 0.1}],
|
| 70 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 71 |
+
)
|
| 72 |
+
image.save("image_Sharpness_0.1.jpg")
|
| 73 |
+
image = template(
|
| 74 |
+
pipe,
|
| 75 |
+
prompt="A cat is sitting on a stone.",
|
| 76 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 77 |
+
template_inputs = [{"scale": 0.8}],
|
| 78 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 79 |
+
)
|
| 80 |
+
image.save("image_Sharpness_0.8.jpg")
|
| 81 |
+
```
|
| 82 |
+
|
| 83 |
+
* Enable lazy loading and memory management, requires 24G GPU memory
|
| 84 |
+
|
| 85 |
+
```python
|
| 86 |
+
from diffsynth.diffusion.template import TemplatePipeline
|
| 87 |
+
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
|
| 88 |
+
import torch
|
| 89 |
+
|
| 90 |
+
```python
|
| 91 |
+
vram_config = {
|
| 92 |
+
"offload_dtype": "disk",
|
| 93 |
+
"offload_device": "disk",
|
| 94 |
+
"onload_dtype": torch.float8_e4m3fn,
|
| 95 |
+
"onload_device": "cpu",
|
| 96 |
+
"preparing_dtype": torch.float8_e4m3fn,
|
| 97 |
+
"preparing_device": "cuda",
|
| 98 |
+
"computation_dtype": torch.bfloat16,
|
| 99 |
+
"computation_device": "cuda",
|
| 100 |
+
}
|
| 101 |
+
pipe = Flux2ImagePipeline.from_pretrained(
|
| 102 |
+
torch_dtype=torch.bfloat16,
|
| 103 |
+
device="cuda",
|
| 104 |
+
model_configs=[
|
| 105 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config),
|
| 106 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config),
|
| 107 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
|
| 108 |
+
],
|
| 109 |
+
tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"),
|
| 110 |
+
vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
|
| 111 |
+
)
|
| 112 |
+
template = TemplatePipeline.from_pretrained(
|
| 113 |
+
torch_dtype=torch.bfloat16,
|
| 114 |
+
device="cuda",
|
| 115 |
+
model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")],
|
| 116 |
+
lazy_loading=True,
|
| 117 |
+
)
|
| 118 |
+
image = template(
|
| 119 |
+
pipe,
|
| 120 |
+
prompt="A cat is sitting on a stone.",
|
| 121 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 122 |
+
template_inputs = [{"scale": 0.1}],
|
| 123 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 124 |
+
)
|
| 125 |
+
image.save("image_Sharpness_0.1.jpg")
|
| 126 |
+
image = template(
|
| 127 |
+
pipe,
|
| 128 |
+
prompt="A cat is sitting on a stone.",
|
| 129 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 130 |
+
template_inputs = [{"scale": 0.8}],
|
| 131 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 132 |
+
)
|
| 133 |
+
image.save("image_Sharpness_0.8.jpg")
|
| 134 |
+
```
|
| 135 |
+
|
| 136 |
+
## Training Code
|
| 137 |
+
|
| 138 |
+
After installing DiffSynth-Studio, use the following script to start training. For more information, please refer to the [DiffSynth-Studio Documentation](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/).
|
| 139 |
+
|
| 140 |
+
```shell
|
| 141 |
+
modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Sharpness/*" --local_dir ./data/diffsynth_example_dataset
|
| 142 |
+
|
| 143 |
+
accelerate launch examples/flux2/model_training/train.py \
|
| 144 |
+
--dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness \
|
| 145 |
+
--dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness/metadata.jsonl \
|
| 146 |
+
--extra_inputs "template_inputs" \
|
| 147 |
+
--max_pixels 1048576 \
|
| 148 |
+
--dataset_repeat 50 \
|
| 149 |
+
--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" \
|
| 150 |
+
--template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Sharpness:" \
|
| 151 |
+
--tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \
|
| 152 |
+
--learning_rate 1e-4 \
|
| 153 |
+
--num_epochs 2 \
|
| 154 |
+
--remove_prefix_in_ckpt "pipe.template_model." \
|
| 155 |
+
--output_path "./models/train/Template-KleinBase4B-Sharpness_full" \
|
| 156 |
+
--trainable_models "template_model" \
|
| 157 |
+
--use_gradient_checkpointing \
|
| 158 |
+
--find_unused_parameters
|
| 159 |
+
```
|
README_from_modelscope.md
ADDED
|
@@ -0,0 +1,163 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
frameworks:
|
| 3 |
+
- Pytorch
|
| 4 |
+
license: Apache License 2.0
|
| 5 |
+
tags: []
|
| 6 |
+
tasks:
|
| 7 |
+
- text-to-image-synthesis
|
| 8 |
+
---
|
| 9 |
+
|
| 10 |
+
# Templates-锐利激发(FLUX.2-klein-base-4B)
|
| 11 |
+
|
| 12 |
+
本模型是 [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) 开源的 Diffusion Templates 系列模型之一。该模型能够通过调整 `scale` 参数,精准控制生成图像的锐度与细节表现力。
|
| 13 |
+
|
| 14 |
+
## 效果展示
|
| 15 |
+
|
| 16 |
+
> **Prompt:** A cat is sitting on a stone.
|
| 17 |
+
|
| 18 |
+
| scale=0.1 | scale=0.8 |
|
| 19 |
+
|:---:|:---:|
|
| 20 |
+
|  |  |
|
| 21 |
+
|
| 22 |
+
---
|
| 23 |
+
|
| 24 |
+
> **Prompt:** A close-up of a person's eyes, looking at the camera, reflections in the pupils, highly aesthetic.
|
| 25 |
+
|
| 26 |
+
| scale=0.1 | scale=0.8 |
|
| 27 |
+
|:---:|:---:|
|
| 28 |
+
|  |  |
|
| 29 |
+
|
| 30 |
+
---
|
| 31 |
+
|
| 32 |
+
> **Prompt:** A beautifully decorated frosted cupcake.
|
| 33 |
+
|
| 34 |
+
| scale=0.1 | scale=0.8 |
|
| 35 |
+
|:---:|:---:|
|
| 36 |
+
|  |  |
|
| 37 |
+
|
| 38 |
+
## 推理代码
|
| 39 |
+
|
| 40 |
+
* 安装 [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio)
|
| 41 |
+
|
| 42 |
+
```
|
| 43 |
+
git clone https://github.com/modelscope/DiffSynth-Studio.git
|
| 44 |
+
cd DiffSynth-Studio
|
| 45 |
+
pip install -e .
|
| 46 |
+
```
|
| 47 |
+
|
| 48 |
+
* 直接推理,需 40G 显存
|
| 49 |
+
|
| 50 |
+
```python
|
| 51 |
+
from diffsynth.diffusion.template import TemplatePipeline
|
| 52 |
+
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
|
| 53 |
+
import torch
|
| 54 |
+
|
| 55 |
+
pipe = Flux2ImagePipeline.from_pretrained(
|
| 56 |
+
torch_dtype=torch.bfloat16,
|
| 57 |
+
device="cuda",
|
| 58 |
+
model_configs=[
|
| 59 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"),
|
| 60 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors"),
|
| 61 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
|
| 62 |
+
],
|
| 63 |
+
tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"),
|
| 64 |
+
)
|
| 65 |
+
template = TemplatePipeline.from_pretrained(
|
| 66 |
+
torch_dtype=torch.bfloat16,
|
| 67 |
+
device="cuda",
|
| 68 |
+
model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")],
|
| 69 |
+
)
|
| 70 |
+
image = template(
|
| 71 |
+
pipe,
|
| 72 |
+
prompt="A cat is sitting on a stone.",
|
| 73 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 74 |
+
template_inputs = [{"scale": 0.1}],
|
| 75 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 76 |
+
)
|
| 77 |
+
image.save("image_Sharpness_0.1.jpg")
|
| 78 |
+
image = template(
|
| 79 |
+
pipe,
|
| 80 |
+
prompt="A cat is sitting on a stone.",
|
| 81 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 82 |
+
template_inputs = [{"scale": 0.8}],
|
| 83 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 84 |
+
)
|
| 85 |
+
image.save("image_Sharpness_0.8.jpg")
|
| 86 |
+
```
|
| 87 |
+
|
| 88 |
+
* 开启惰性加载和显存管理,需 24G 显存
|
| 89 |
+
|
| 90 |
+
```python
|
| 91 |
+
from diffsynth.diffusion.template import TemplatePipeline
|
| 92 |
+
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
|
| 93 |
+
import torch
|
| 94 |
+
|
| 95 |
+
vram_config = {
|
| 96 |
+
"offload_dtype": "disk",
|
| 97 |
+
"offload_device": "disk",
|
| 98 |
+
"onload_dtype": torch.float8_e4m3fn,
|
| 99 |
+
"onload_device": "cpu",
|
| 100 |
+
"preparing_dtype": torch.float8_e4m3fn,
|
| 101 |
+
"preparing_device": "cuda",
|
| 102 |
+
"computation_dtype": torch.bfloat16,
|
| 103 |
+
"computation_device": "cuda",
|
| 104 |
+
}
|
| 105 |
+
pipe = Flux2ImagePipeline.from_pretrained(
|
| 106 |
+
torch_dtype=torch.bfloat16,
|
| 107 |
+
device="cuda",
|
| 108 |
+
model_configs=[
|
| 109 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors", **vram_config),
|
| 110 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="text_encoder/*.safetensors", **vram_config),
|
| 111 |
+
ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
|
| 112 |
+
],
|
| 113 |
+
tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-4B", origin_file_pattern="tokenizer/"),
|
| 114 |
+
vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 0.5,
|
| 115 |
+
)
|
| 116 |
+
template = TemplatePipeline.from_pretrained(
|
| 117 |
+
torch_dtype=torch.bfloat16,
|
| 118 |
+
device="cuda",
|
| 119 |
+
model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Sharpness")],
|
| 120 |
+
lazy_loading=True,
|
| 121 |
+
)
|
| 122 |
+
image = template(
|
| 123 |
+
pipe,
|
| 124 |
+
prompt="A cat is sitting on a stone.",
|
| 125 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 126 |
+
template_inputs = [{"scale": 0.1}],
|
| 127 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 128 |
+
)
|
| 129 |
+
image.save("image_Sharpness_0.1.jpg")
|
| 130 |
+
image = template(
|
| 131 |
+
pipe,
|
| 132 |
+
prompt="A cat is sitting on a stone.",
|
| 133 |
+
seed=0, cfg_scale=4, num_inference_steps=50,
|
| 134 |
+
template_inputs = [{"scale": 0.8}],
|
| 135 |
+
negative_template_inputs = [{"scale": 0.5}],
|
| 136 |
+
)
|
| 137 |
+
image.save("image_Sharpness_0.8.jpg")
|
| 138 |
+
```
|
| 139 |
+
|
| 140 |
+
## 训练代码
|
| 141 |
+
|
| 142 |
+
安装 DiffSynth-Studio 后,使用以下脚本可开启训练,更多信息请参考 [DiffSynth-Studio 文档](https://diffsynth-studio-doc.readthedocs.io/zh-cn/latest/)。
|
| 143 |
+
|
| 144 |
+
```shell
|
| 145 |
+
modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Sharpness/*" --local_dir ./data/diffsynth_example_dataset
|
| 146 |
+
|
| 147 |
+
accelerate launch examples/flux2/model_training/train.py \
|
| 148 |
+
--dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness \
|
| 149 |
+
--dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Sharpness/metadata.jsonl \
|
| 150 |
+
--extra_inputs "template_inputs" \
|
| 151 |
+
--max_pixels 1048576 \
|
| 152 |
+
--dataset_repeat 50 \
|
| 153 |
+
--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" \
|
| 154 |
+
--template_model_id_or_path "DiffSynth-Studio/Template-KleinBase4B-Sharpness:" \
|
| 155 |
+
--tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \
|
| 156 |
+
--learning_rate 1e-4 \
|
| 157 |
+
--num_epochs 2 \
|
| 158 |
+
--remove_prefix_in_ckpt "pipe.template_model." \
|
| 159 |
+
--output_path "./models/train/Template-KleinBase4B-Sharpness_full" \
|
| 160 |
+
--trainable_models "template_model" \
|
| 161 |
+
--use_gradient_checkpointing \
|
| 162 |
+
--find_unused_parameters
|
| 163 |
+
```
|
assets/cake_Sharpness_0.1.jpg
ADDED
|
assets/cake_Sharpness_0.8.jpg
ADDED
|
assets/cat_Sharpness_0.1.jpg
ADDED
|
Git LFS Details
|
assets/cat_Sharpness_0.8.jpg
ADDED
|
Git LFS Details
|
assets/eye_Sharpness_0.1.jpg
ADDED
|
assets/eye_Sharpness_0.8.jpg
ADDED
|
Git LFS Details
|
configuration.json
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
{"framework":"Pytorch","task":"text-to-image-synthesis"}
|
model.py
ADDED
|
@@ -0,0 +1,76 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch, math, os, json
|
| 2 |
+
from PIL import Image
|
| 3 |
+
import numpy as np
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
class SingleValueEncoder(torch.nn.Module):
|
| 7 |
+
def __init__(self, dim_in=256, dim_out=4096, length=32):
|
| 8 |
+
super().__init__()
|
| 9 |
+
self.length = length
|
| 10 |
+
self.prefer_value_embedder = torch.nn.Sequential(torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out))
|
| 11 |
+
self.positional_embedding = torch.nn.Parameter(torch.randn(self.length, dim_out))
|
| 12 |
+
|
| 13 |
+
def get_timestep_embedding(self, timesteps, embedding_dim, max_period=10000):
|
| 14 |
+
half_dim = embedding_dim // 2
|
| 15 |
+
exponent = -math.log(max_period) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / half_dim
|
| 16 |
+
emb = timesteps[:, None].float() * torch.exp(exponent)[None, :]
|
| 17 |
+
emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1)
|
| 18 |
+
return emb
|
| 19 |
+
|
| 20 |
+
def forward(self, value, dtype):
|
| 21 |
+
emb = self.get_timestep_embedding(value * 1000, 256).to(dtype)
|
| 22 |
+
emb = self.prefer_value_embedder(emb).squeeze(0)
|
| 23 |
+
base_embeddings = emb.expand(self.length, -1)
|
| 24 |
+
positional_embedding = self.positional_embedding.to(dtype=base_embeddings.dtype, device=base_embeddings.device)
|
| 25 |
+
learned_embeddings = base_embeddings + positional_embedding
|
| 26 |
+
return learned_embeddings
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class ValueFormatModel(torch.nn.Module):
|
| 30 |
+
def __init__(self, num_double_blocks=5, num_single_blocks=20, dim=3072, num_heads=24, length=512):
|
| 31 |
+
super().__init__()
|
| 32 |
+
self.block_names = [f"double_{i}" for i in range(num_double_blocks)] + [f"single_{i}" for i in range(num_single_blocks)]
|
| 33 |
+
self.proj_k = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names})
|
| 34 |
+
self.proj_v = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names})
|
| 35 |
+
self.num_heads = num_heads
|
| 36 |
+
self.length = length
|
| 37 |
+
|
| 38 |
+
@torch.no_grad()
|
| 39 |
+
def process_inputs(self, pipe, scale, **kwargs):
|
| 40 |
+
return {"value": torch.Tensor([scale]).to(dtype=pipe.torch_dtype, device=pipe.device)}
|
| 41 |
+
|
| 42 |
+
def forward(self, value, **kwargs):
|
| 43 |
+
kv_cache = {}
|
| 44 |
+
for block_name in self.block_names:
|
| 45 |
+
k = self.proj_k[block_name](value, value.dtype)
|
| 46 |
+
k = k.view(1, self.length, self.num_heads, -1)
|
| 47 |
+
v = self.proj_v[block_name](value, value.dtype)
|
| 48 |
+
v = v.view(1, self.length, self.num_heads, -1)
|
| 49 |
+
kv_cache[block_name] = (k, v)
|
| 50 |
+
return {"kv_cache": kv_cache}
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
class DataAnnotator:
|
| 54 |
+
def __init__(self):
|
| 55 |
+
with open(os.path.join(os.path.dirname(__file__), "scores.json"), "r") as f:
|
| 56 |
+
self.scores = json.load(f)
|
| 57 |
+
|
| 58 |
+
def get_score(self, x):
|
| 59 |
+
l, r = 0, len(self.scores)
|
| 60 |
+
while l < r:
|
| 61 |
+
m = (l + r) // 2
|
| 62 |
+
if self.scores[m] < x: l = m + 1
|
| 63 |
+
else: r = m
|
| 64 |
+
return l / len(self.scores)
|
| 65 |
+
|
| 66 |
+
def __call__(self, image, **kwargs):
|
| 67 |
+
import cv2
|
| 68 |
+
image = cv2.imread(image, cv2.IMREAD_GRAYSCALE)
|
| 69 |
+
edges = cv2.Canny(image, 100, 200)
|
| 70 |
+
scale = edges.astype(np.float32).mean().tolist()
|
| 71 |
+
return {"scale": self.get_score(scale)}
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
TEMPLATE_MODEL = ValueFormatModel
|
| 75 |
+
TEMPLATE_MODEL_PATH = "model.safetensors"
|
| 76 |
+
TEMPLATE_DATA_PROCESSOR = DataAnnotator
|
model.safetensors
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:c95fb95b0522ced881f25f7b380eaeee0f2668a5ddcb414af4fccdd14439b035
|
| 3 |
+
size 1180292000
|
scores.json
ADDED
|
The diff for this file is too large to render.
See raw diff
|
|
|