中文Sign up

DiffSynth-Studio/Template-KleinBase4B-Aesthetic

🤗 Hugging Face sourceapache-2.0664M params1.3 GBsafetensors✓ 7 checksumsupdated today
Submit in one command

Run it next to your model folder. It makes the torrent, checks your files against Hugging Face, and submits it. You just start seeding and paste your key from your account. It only reads your files and never changes them. Read the script first if you like.

curl -fsSL https://pirateface.co/package.sh | bash -s -- --repo DiffSynth-Studio/Template-KleinBase4B-Aesthetic ./model-folder
Needs a seeder →

Templates - Aesthetic Alignment (FLUX.2-klein-base-4B)

This model is one of the open-source Diffusion Templates series models from 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

Prompt: A cat is sitting on a stone.

base model scale=1.0 scale=2.5

Prompt: A cute anime girl with pink hair and cat ears, pastel colors.

base model scale=1.0 scale=2.5

Prompt: A cyberpunk apartment with a view of neon lights.

base model scale=1.0 scale=2.5

Inference Code

git clone https://github.com/modelscope/DiffSynth-Studio.git
cd DiffSynth-Studio
pip install -e .
  • Direct inference (requires 40G GPU memory)
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
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 Script

After installing DiffSynth-Studio, use the following script to start training. For more information, please refer to the DiffSynth-Studio Documentation.

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