Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/apartment_Aesthetic_1.0.jpg filter=lfs diff=lfs merge=lfs -text
+assets/apartment_Aesthetic_2.5.jpg filter=lfs diff=lfs merge=lfs -text
+assets/apartment_base.jpg filter=lfs diff=lfs merge=lfs -text
+assets/cat_Aesthetic_1.0.jpg filter=lfs diff=lfs merge=lfs -text
+assets/cat_Aesthetic_2.5.jpg filter=lfs diff=lfs merge=lfs -text
+assets/cat_base.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,192 @@
+---
+license: apache-2.0
+---
+# Templates - Aesthetic Alignment (FLUX.2-klein-base-4B)
+
+This model is one of the open-source Diffusion Templates series models from [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio). It is an Aesthetic alignment model that allows adjusting the degree of aesthetic alignment in generated images by modifying the `scale` parameter.
+
+## Results Showcase
+
+> **Prompt:** A cat is sitting on a stone.
+
+| base model | scale=1.0 | scale=2.5 |
+|:---:|:---:|:---:|
+| ![](./assets/cat_base.jpg) | ![](./assets/cat_Aesthetic_1.0.jpg) | ![](./assets/cat_Aesthetic_2.5.jpg) |
+
+---
+
+> **Prompt:** A cute anime girl with pink hair and cat ears, pastel colors.
+
+| base model | scale=1.0 | scale=2.5 |
+|:---:|:---:|:---:|
+| ![](./assets/girl_base.jpg) | ![](./assets/girl_Aesthetic_1.0.jpg) | ![](./assets/girl_Aesthetic_2.5.jpg) |
+
+---
+
+> **Prompt:** A cyberpunk apartment with a view of neon lights.
+
+| base model | scale=1.0 | scale=2.5 |
+|:---:|:---:|:---:|
+| ![](./assets/apartment_base.jpg) | ![](./assets/apartment_Aesthetic_1.0.jpg) | ![](./assets/apartment_Aesthetic_2.5.jpg) |
+
+## Inference Code
+
+* Install [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio)
+
+```
+git clone https://github.com/modelscope/DiffSynth-Studio.git
+cd DiffSynth-Studio
+pip install -e .
+```
+
+* Direct inference (requires 40GB GPU memory)
+
+```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/"),
+)
+pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important!
+template = TemplatePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Aesthetic")],
+)
+image = template(
+    pipe,
+    prompt="A cat is sitting on a stone.",
+    seed=0, cfg_scale=4, num_inference_steps=50,
+    template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_1.0.jpg")
+image = template(
+    pipe,
+    prompt="A cat is sitting on a stone.",
+    seed=0, cfg_scale=4, num_inference_steps=50,
+    template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_2.5.jpg")
+```
+
+* Enable lazy loading and memory management, requires 24G GPU memory
+
+```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-Aesthetic")],
+    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 = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_1.0.jpg")
+image = template(
+    pipe,
+    prompt="A cat is sitting on a stone.",
+    seed=0, cfg_scale=4, num_inference_steps=50,
+    template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_2.5.jpg")
+
+```
+
+## Training Code
+
+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/).
+
+```shell
+modelscope download --dataset DiffSynth-Studio/diffsynth_example_dataset --include "flux2/Template-KleinBase4B-Aesthetic/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/flux2/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Aesthetic \
+  --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Aesthetic/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-Aesthetic:" \
+  --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-Aesthetic_full" \
+  --trainable_models "template_model" \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --enable_lora_hot_loading
+```

README_from_modelscope.md

@@ -0,0 +1,199 @@
+---
+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 系列模型之一。该模型为 Aesthetic（美学）对齐模型，能够通过修改 `scale` 参数来调整图像的美学对齐程度。
+
+## 效果展示
+
+> **Prompt:** A cat is sitting on a stone.
+
+| base model | scale=1.0 | scale=2.5 |
+|:---:|:---:|:---:|
+| ![](./assets/cat_base.jpg) | ![](./assets/cat_Aesthetic_1.0.jpg) | ![](./assets/cat_Aesthetic_2.5.jpg) |
+
+---
+
+> **Prompt:** A cute anime girl with pink hair and cat ears, pastel colors.
+
+| base model | scale=1.0 | scale=2.5 |
+|:---:|:---:|:---:|
+| ![](./assets/girl_base.jpg) | ![](./assets/girl_Aesthetic_1.0.jpg) | ![](./assets/girl_Aesthetic_2.5.jpg) |
+
+---
+
+> **Prompt:** A cyberpunk apartment with a view of neon lights.
+
+| base model | scale=1.0 | scale=2.5 |
+|:---:|:---:|:---:|
+| ![](./assets/apartment_base.jpg) | ![](./assets/apartment_Aesthetic_1.0.jpg) | ![](./assets/apartment_Aesthetic_2.5.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/"),
+)
+pipe.dit = pipe.enable_lora_hot_loading(pipe.dit) # Important!
+template = TemplatePipeline.from_pretrained(
+    torch_dtype=torch.bfloat16,
+    device="cuda",
+    model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Aesthetic")],
+)
+image = template(
+    pipe,
+    prompt="A cat is sitting on a stone.",
+    seed=0, cfg_scale=4, num_inference_steps=50,
+    template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_1.0.jpg")
+image = template(
+    pipe,
+    prompt="A cat is sitting on a stone.",
+    seed=0, cfg_scale=4, num_inference_steps=50,
+    template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_2.5.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-Aesthetic")],
+    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 = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 1.0,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_1.0.jpg")
+image = template(
+    pipe,
+    prompt="A cat is sitting on a stone.",
+    seed=0, cfg_scale=4, num_inference_steps=50,
+    template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+    negative_template_inputs = [{
+        "lora_ids": list(range(1, 180, 2)),
+        "lora_scales": 2.5,
+        "merge_type": "mean",
+    }],
+)
+image.save("image_Aesthetic_2.5.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-Aesthetic/*" --local_dir ./data/diffsynth_example_dataset
+
+accelerate launch examples/flux2/model_training/train.py \
+  --dataset_base_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Aesthetic \
+  --dataset_metadata_path data/diffsynth_example_dataset/flux2/Template-KleinBase4B-Aesthetic/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-Aesthetic:" \
+  --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-Aesthetic_full" \
+  --trainable_models "template_model" \
+  --use_gradient_checkpointing \
+  --find_unused_parameters \
+  --enable_lora_hot_loading
+
+```

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"text-to-image-synthesis"}

model.py

@@ -0,0 +1,108 @@
+import torch
+
+
+class LoRALayer(torch.nn.Module):
+    def __init__(self, dim_in, dim_out, rank, initialize=False):
+        super().__init__()
+        if initialize:
+            scale = (1 / dim_in) ** 0.5
+            self.lora_A = torch.nn.Parameter(torch.rand((rank, dim_in)) * (scale * 2) - scale)
+            self.lora_B = torch.nn.Parameter(torch.zeros((dim_out, rank)))
+        else:
+            self.lora_A = torch.nn.Parameter(torch.empty((rank, dim_in)))
+            self.lora_B = torch.nn.Parameter(torch.empty((dim_out, rank)))
+
+
+class LoRA(torch.nn.Module):
+    def __init__(self, rank):
+        super().__init__()
+        self.lora_patterns = [
+            {
+                "name": "single_transformer_blocks.{block_id}.attn.to_qkv_mlp_proj",
+                "num_blocks": 20,
+                "dim_in": 3072,
+                "dim_out": 27648,
+                "rank": rank,
+            },
+            {
+                "name": "single_transformer_blocks.{block_id}.attn.to_out",
+                "num_blocks": 20,
+                "dim_in": 12288,
+                "dim_out": 3072,
+                "rank": rank,
+            },
+        ]
+        self.parse_lora_layers(self.lora_patterns)
+    
+    def parse_lora_layers(self, lora_patterns):
+        names = []
+        layers = []
+        for lora_pattern in lora_patterns:
+            for block_id in range(lora_pattern["num_blocks"]):
+                name = lora_pattern["name"].format(block_id=block_id)
+                layer = LoRALayer(lora_pattern["dim_in"], lora_pattern["dim_out"], lora_pattern["rank"])
+                names.append(name)
+                layers.append(layer)
+        self.names = names
+        self.layers = torch.nn.ModuleList(layers)
+
+    def forward(self):
+        lora = {}
+        for name, layer in zip(self.names, self.layers):
+            lora[f"{name}.lora_A.default.weight"] = layer.lora_A
+            lora[f"{name}.lora_B.default.weight"] = layer.lora_B
+        return lora
+
+
+class DualLoRA(torch.nn.Module):
+    def __init__(self, num_loras=180):
+        super().__init__()
+        self.loras = torch.nn.ModuleList([LoRA(rank=4) for _ in range(num_loras)])
+
+    @torch.no_grad()
+    def process_inputs(self, lora_ids, lora_scales, require_grads=None, merge_type="concat", **kwargs):
+        return {"lora_ids": lora_ids, "lora_scales": lora_scales, "require_grads": require_grads, "merge_type": merge_type}
+    
+    def forward(self, lora_ids, lora_scales, require_grads=None, merge_type="concat", **kwargs):
+        if isinstance(lora_scales, float):
+            lora_scales = [lora_scales] * len(lora_ids)
+        if require_grads is None:
+            require_grads = [True] * len(lora_scales)
+        loras = []
+        for lora_id, lora_scale, require_grad in zip(lora_ids, lora_scales, require_grads):
+            if not require_grad:
+                with torch.no_grad():
+                    lora_ = self.loras[lora_id]()
+            else:
+                lora_ = self.loras[lora_id]()
+            lora_ = {key: lora_[key] * (lora_scale if "lora_A" in key else 1) for key in lora_}
+            loras.append(lora_)
+        lora = {}
+        if merge_type == "concat":
+            for key in loras[0]:
+                if "lora_A" in key:
+                    lora[key] = torch.concat([lora_[key] for lora_ in loras], dim=0)
+                else:
+                    lora[key] = torch.concat([lora_[key] for lora_ in loras], dim=1)
+        elif merge_type == "sum":
+            for key in loras[0]:
+                lora[key] = torch.stack([lora_[key] for lora_ in loras]).sum(dim=0)
+        elif merge_type == "mean":
+            for key in loras[0]:
+                if "lora_A" in key:
+                    lora[key] = torch.stack([lora_[key] for lora_ in loras]).mean(dim=0)
+                else:
+                    lora[key] = torch.stack([lora_[key] for lora_ in loras]).sum(dim=0)
+        else:
+            raise ValueError(f"Unsupported merge_type: {merge_type}")
+        return {"lora": lora}
+
+
+class DataAnnotator:
+    def __call__(self, **kwargs):
+        return kwargs
+
+
+TEMPLATE_MODEL = DualLoRA
+TEMPLATE_MODEL_PATH = "model.safetensors"
+TEMPLATE_DATA_PROCESSOR = DataAnnotator

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6323dbbbcabdeb7ea9c203e0f6ed0d61a094edde0e1a8f7c134c132474ff7485
+size 1328543560