ModelScope AIGC Series Course - Controllable Generation Technology
This experiment uses Diffusion-Templates as the framework to systematically introduce various controllable generation techniques for image generation models, and demonstrate how to train a controllable generation module from scratch.
Related Resources:
Open Source Code: DiffSynth-Studio
Technical Report: arXiv
Project Homepage: GitHub
Documentation: English Version、中文版
Online Demo: ModelScope Studio
Model Collection: ModelScope、ModelScope International、HuggingFace
!pip install diffsynth==2.0.15 transformers==5.8.1
from diffsynth.diffusion.template import TemplatePipeline
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
import torch
from modelscope import dataset_snapshot_download, snapshot_download
from PIL import Image
import numpy as np
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",
}
def show_images(images, resolution):
images = [i.resize((resolution, resolution)).convert("RGB") for i in images]
images = [np.array(i) for i in images]
images = np.concat(images, axis=1)
images = Image.fromarray(images)
return images
Image Structure Control
First, load the base model black-forest-labs/FLUX.2-klein-base-4B. This is a 4B-parameter image generation model, and all controllable generation modules in this experiment will be mounted on top of this base model.
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,
)
ControlNet is one of the earliest controllable generation techniques for Diffusion models. It uses structural conditions such as depth maps, edge maps, and pose maps to achieve pixel-level control over the generated image.
By loading DiffSynth-Studio/Template-KleinBase4B-ControlNet in Template format, you can generate images with different styles using different prompts while preserving the input structure.
template = TemplatePipeline.from_pretrained(
torch_dtype=torch.bfloat16,
device="cuda",
model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-ControlNet")],
lazy_loading=True,
)
dataset_snapshot_download(
"DiffSynth-Studio/examples_in_diffsynth",
allow_file_pattern=["templates/*"],
local_dir="data/examples",
)
image = template(
pipe,
prompt="A cat is sitting on a stone, bathed in bright sunshine.",
seed=0, cfg_scale=4, num_inference_steps=50,
template_inputs = [{
"image": Image.open("data/examples/templates/image_depth.jpg"),
"prompt": "A cat is sitting on a stone, bathed in bright sunshine.",
}],
negative_template_inputs = [{
"image": Image.open("data/examples/templates/image_depth.jpg"),
"prompt": "",
}],
)
image.save("image_ControlNet_sunshine.jpg")
image = template(
pipe,
prompt="A cat is sitting on a stone, surrounded by colorful magical particles.",
seed=0, cfg_scale=4, num_inference_steps=50,
template_inputs = [{
"image": Image.open("data/examples/templates/image_depth.jpg"),
"prompt": "A cat is sitting on a stone, surrounded by colorful magical particles.",
}],
negative_template_inputs = [{
"image": Image.open("data/examples/templates/image_depth.jpg"),
"prompt": "",
}],
)
image.save("image_ControlNet_magic.jpg")
show_images([
Image.open("data/examples/templates/image_depth.jpg"),
Image.open("image_ControlNet_sunshine.jpg"),
Image.open("image_ControlNet_magic.jpg"),
], resolution=256)
Attribute Value Control
AttriCtrl is a type of controllable generation model capable of injecting continuous value attributes as control conditions into the generation process.
Run the following code to load DiffSynth-Studio/Template-KleinBase4B-SoftRGB and precisely control the overall color tone of the image by inputting R/G/B values.
template = TemplatePipeline.from_pretrained(
torch_dtype=torch.bfloat16,
device="cuda",
model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-SoftRGB")],
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 = [{"R": 128/255, "G": 128/255, "B": 128/255}],
)
image.save("image_rgb_normal.jpg")
image = template(
pipe,
prompt="A cat is sitting on a stone.",
seed=0, cfg_scale=4, num_inference_steps=50,
template_inputs = [{"R": 208/255, "G": 185/255, "B": 138/255}],
)
image.save("image_rgb_warm.jpg")
image = template(
pipe,
prompt="A cat is sitting on a stone.",
seed=0, cfg_scale=4, num_inference_steps=50,
template_inputs = [{"R": 94/255, "G": 163/255, "B": 174/255}],
)
image.save("image_rgb_cold.jpg")
show_images([
Image.open("image_rgb_normal.jpg"),
Image.open("image_rgb_warm.jpg"),
Image.open("image_rgb_cold.jpg"),
], resolution=256)
Image Editing
Image editing models are a type of highly versatile controllable generation model: given an original image and an editing instruction, the model can make partial or overall modifications to the original image.
Run the following code to load DiffSynth-Studio/Template-KleinBase4B-Edit. This model uses KV-Cache to reuse the attention key-value pairs of the input image, enabling fast editing with quick inference speed.
template = TemplatePipeline.from_pretrained(
torch_dtype=torch.bfloat16,
device="cuda",
model_configs=[ModelConfig(model_id="DiffSynth-Studio/Template-KleinBase4B-Edit")],
lazy_loading=True,
)
dataset_snapshot_download(
"DiffSynth-Studio/examples_in_diffsynth",
allow_file_pattern=["templates/*"],
local_dir="data/examples",
)
image = template(
pipe,
prompt="Put a hat on this cat.",
seed=0, cfg_scale=4, num_inference_steps=50,
template_inputs = [{
"image": Image.open("data/examples/templates/image_reference.jpg"),
"prompt": "Put a hat on this cat.",
}],
negative_template_inputs = [{
"image": Image.open("data/examples/templates/image_reference.jpg"),
"prompt": "",
}],
)
image.save("image_Edit_hat.jpg")
image = template(
pipe,
prompt="Make the cat turn its head to look to the right.",
seed=0, cfg_scale=4, num_inference_steps=50,
template_inputs = [{
"image": Image.open("data/examples/templates/image_reference.jpg"),
"prompt": "Make the cat turn its head to look to the right.",
}],
negative_template_inputs = [{
"image": Image.open("data/examples/templates/image_reference.jpg"),
"prompt": "",
}],
)
image.save("image_Edit_head.jpg")
show_images([
Image.open("data/examples/templates/image_reference.jpg"),
Image.open("image_Edit_hat.jpg"),
Image.open("image_Edit_head.jpg"),
], resolution=256)
Style Control
The most straightforward way to achieve image style control is to train a style LoRA — however, each style requires separate training, which is costly. To address this, we trained a special Image-to-LoRA model that can generate LoRA weights on-demand from input reference images, eliminating the traditional style training process.
Run the following code to load DiffSynth-Studio/KleinBase4B-i2L-v2 and dynamically generate LoRA from reference images to control the image style.
from modelscope import snapshot_download
template = TemplatePipeline.from_pretrained(
torch_dtype=torch.bfloat16,
device="cuda",
model_configs=[ModelConfig(model_id="DiffSynth-Studio/KleinBase4B-i2L-v2")],
lazy_loading=True,
)
snapshot_download("DiffSynth-Studio/KleinBase4B-i2L-v2", allow_file_pattern="assets/*", local_dir="data")
images = [Image.open(f"data/assets/image_1_{i}.jpg") for i in range(4)]
image = template(
pipe,
prompt="A cat is sitting on a stone",
seed=42, cfg_scale=4, num_inference_steps=50,
template_inputs = [{"image": images}],
negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}],
)
image.save("image_KleinBase4B-i2L-v2_1.jpg")
images = [Image.open(f"data/assets/image_3_{i}.jpg") for i in range(4)]
image = template(
pipe,
prompt="A cat is sitting on a stone",
seed=42, cfg_scale=4, num_inference_steps=50,
template_inputs = [{"image": images}],
negative_template_inputs = [{"image": [Image.fromarray(np.zeros_like(np.array(i)) + 128) for i in images]}],
)
image.save("image_KleinBase4B-i2L-v2_2.jpg")
show_images([
Image.open("data/assets/image_1_2.jpg"),
Image.open("image_KleinBase4B-i2L-v2_1.jpg"),
Image.open("data/assets/image_3_0.jpg"),
Image.open("image_KleinBase4B-i2L-v2_2.jpg"),
], resolution=256)
Training Controllable Generation Models
The Diffusion-Templates framework allows developers to train controllable generation models of any structure — as long as you provide the model definition, data processing logic, and dataset, you can integrate into a unified training workflow. Below, we train a brightness control model from scratch, allowing images to be generated with specified brightness values.
Step 1: Write the model structure code (including the numerical encoder, KV-Cache generation backbone, and data annotator):
code = """
import torch, math, os
from PIL import Image
import numpy as np
class SingleValueEncoder(torch.nn.Module):
def __init__(self, dim_in=256, dim_out=4096, length=32):
super().__init__()
self.length = length
self.prefer_value_embedder = torch.nn.Sequential(torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out))
self.positional_embedding = torch.nn.Parameter(torch.randn(self.length, dim_out))
def get_timestep_embedding(self, timesteps, embedding_dim, max_period=10000):
half_dim = embedding_dim // 2
exponent = -math.log(max_period) * torch.arange(0, half_dim, dtype=torch.float32, device=timesteps.device) / half_dim
emb = timesteps[:, None].float() * torch.exp(exponent)[None, :]
emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1)
return emb
def forward(self, value, dtype):
emb = self.get_timestep_embedding(value * 1000, 256).to(dtype)
emb = self.prefer_value_embedder(emb).squeeze(0)
base_embeddings = emb.expand(self.length, -1)
positional_embedding = self.positional_embedding.to(dtype=base_embeddings.dtype, device=base_embeddings.device)
learned_embeddings = base_embeddings + positional_embedding
return learned_embeddings
# Backbone model structure (converts input values into KV-Cache vectors)
class ValueFormatModel(torch.nn.Module):
def __init__(self, num_double_blocks=5, num_single_blocks=20, dim=3072, num_heads=24, length=512):
super().__init__()
self.block_names = [f"double_{i}" for i in range(num_double_blocks)] + [f"single_{i}" for i in range(num_single_blocks)]
self.proj_k = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names})
self.proj_v = torch.nn.ModuleDict({block_name: SingleValueEncoder(dim_out=dim, length=length) for block_name in self.block_names})
self.num_heads = num_heads
self.length = length
@torch.no_grad()
def process_inputs(self, pipe, scale, **kwargs):
return {"value": torch.Tensor([scale]).to(dtype=pipe.torch_dtype, device=pipe.device)}
def forward(self, value, **kwargs):
kv_cache = {}
for block_name in self.block_names:
k = self.proj_k[block_name](value, value.dtype)
k = k.view(1, self.length, self.num_heads, -1)
v = self.proj_v[block_name](value, value.dtype)
v = v.view(1, self.length, self.num_heads, -1)
kv_cache[block_name] = (k, v)
return {"kv_cache": kv_cache}
# Converts image data to model input (calculates brightness from RGB values in the image)
class DataAnnotator(torch.nn.Module):
def __init__(self):
pass
def __call__(self, image, **kwargs):
image = Image.open(image)
image = np.array(image)
return {"scale": image.astype(np.float32).mean() / 255}
TEMPLATE_MODEL = ValueFormatModel
TEMPLATE_MODEL_PATH = "model.safetensors" if "model.safetensors" in os.listdir(os.path.dirname(__file__)) else None
TEMPLATE_DATA_PROCESSOR = DataAnnotator
"""
import os
os.makedirs("models/template_brightness", exist_ok=True)
with open("models/template_brightness/model.py", "w", encoding="utf-8") as f:
f.write(code.strip())
Step 2: Download and preprocess the dataset, while generating the metadata required for training:
import json, os
from modelscope import dataset_snapshot_download
# Download dataset
dataset_snapshot_download(
"DiffSynth-Studio/ImagePulseV2-TextImage",
local_dir="data/ImagePulseV2-TextImage",
allow_file_pattern="data/1770381050168240056.tar.gz"
)
# Extract dataset
os.makedirs("data/dataset", exist_ok=True)
os.system("tar zxvf data/ImagePulseV2-TextImage/data/1770381050168240056.tar.gz -C data/dataset")
# Generate dataset metadata
dataset_path = "data/dataset/1770381050168240056"
metadata = []
for file_name in os.listdir(dataset_path):
if file_name.endswith(".json"):
with open(os.path.join(dataset_path, file_name), "r") as f:
data = json.load(f)
data["template_inputs"] = {"image": os.path.join(dataset_path, data["image"])}
metadata.append(data)
with open("data/dataset/metadata.json", "w") as f:
json.dump(metadata, f, indent=4, ensure_ascii=False)
Step 3: Start training:
import os
# Training script
code = """
import torch, os, argparse, accelerate
from diffsynth.core import UnifiedDataset
from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
from diffsynth.diffusion import *
os.environ["TOKENIZERS_PARALLELISM"] = "false"
class Flux2ImageTrainingModule(DiffusionTrainingModule):
def __init__(
self,
model_paths=None, model_id_with_origin_paths=None,
tokenizer_path=None,
trainable_models=None,
lora_base_model=None, lora_target_modules="", lora_rank=32, lora_checkpoint=None,
preset_lora_path=None, preset_lora_model=None,
use_gradient_checkpointing=True,
use_gradient_checkpointing_offload=False,
extra_inputs=None,
fp8_models=None,
offload_models=None,
template_model_id_or_path=None,
resume_from_checkpoint=None, remove_prefix_in_ckpt=None,
enable_lora_hot_loading=False,
device="cpu",
task="sft",
):
super().__init__()
# Load models
model_configs = self.parse_model_configs(model_paths, model_id_with_origin_paths, fp8_models=fp8_models, offload_models=offload_models, device=device)
tokenizer_config = self.parse_path_or_model_id(tokenizer_path, default_value=ModelConfig(model_id="black-forest-labs/FLUX.2-dev", origin_file_pattern="tokenizer/"))
self.pipe = Flux2ImagePipeline.from_pretrained(torch_dtype=torch.bfloat16, device=device, model_configs=model_configs, tokenizer_config=tokenizer_config)
self.pipe = self.load_training_template_model(self.pipe, template_model_id_or_path, use_gradient_checkpointing, use_gradient_checkpointing_offload)
self.pipe = self.split_pipeline_units(task, self.pipe, trainable_models, lora_base_model, remove_unnecessary_params=True)
self.resume_from_checkpoint(resume_from_checkpoint, remove_prefix_in_ckpt)
if enable_lora_hot_loading: self.pipe.dit = self.pipe.enable_lora_hot_loading(self.pipe.dit)
# Training mode
self.switch_pipe_to_training_mode(
self.pipe, trainable_models,
lora_base_model, lora_target_modules, lora_rank, lora_checkpoint,
preset_lora_path, preset_lora_model,
task=task,
)
# Other configs
self.use_gradient_checkpointing = use_gradient_checkpointing
self.use_gradient_checkpointing_offload = use_gradient_checkpointing_offload
self.extra_inputs = extra_inputs.split(",") if extra_inputs is not None else []
self.fp8_models = fp8_models
self.task = task
self.task_to_loss = {
"sft:data_process": lambda pipe, *args: args,
"direct_distill:data_process": lambda pipe, *args: args,
"sft": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
"sft:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: FlowMatchSFTLoss(pipe, **inputs_shared, **inputs_posi),
"direct_distill": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi),
"direct_distill:train": lambda pipe, inputs_shared, inputs_posi, inputs_nega: DirectDistillLoss(pipe, **inputs_shared, **inputs_posi),
}
def get_pipeline_inputs(self, data):
inputs_posi = {"prompt": data["prompt"]}
inputs_nega = {"negative_prompt": ""}
inputs_shared = {
# Assume you are using this pipeline for inference,
# please fill in the input parameters.
"input_image": data["image"],
"height": data["image"].size[1],
"width": data["image"].size[0],
# Please do not modify the following parameters
# unless you clearly know what this will cause.
"embedded_guidance": 1.0,
"cfg_scale": 1,
"rand_device": self.pipe.device,
"use_gradient_checkpointing": self.use_gradient_checkpointing,
"use_gradient_checkpointing_offload": self.use_gradient_checkpointing_offload,
}
inputs_shared = self.parse_extra_inputs(data, self.extra_inputs, inputs_shared)
return inputs_shared, inputs_posi, inputs_nega
def forward(self, data, inputs=None):
if inputs is None: inputs = self.get_pipeline_inputs(data)
inputs = self.transfer_data_to_device(inputs, self.pipe.device, self.pipe.torch_dtype)
for unit in self.pipe.units:
inputs = self.pipe.unit_runner(unit, self.pipe, *inputs)
loss = self.task_to_loss[self.task](self.pipe, *inputs)
return loss
def flux2_parser():
parser = argparse.ArgumentParser(description="Simple example of a training script.")
parser = add_general_config(parser)
parser = add_image_size_config(parser)
parser.add_argument("--tokenizer_path", type=str, default=None, help="Path to tokenizer.")
parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.")
return parser
if __name__ == "__main__":
parser = flux2_parser()
args = parser.parse_args()
accelerator = accelerate.Accelerator(
gradient_accumulation_steps=args.gradient_accumulation_steps,
kwargs_handlers=[accelerate.DistributedDataParallelKwargs(find_unused_parameters=args.find_unused_parameters)],
)
dataset = UnifiedDataset(
base_path=args.dataset_base_path,
metadata_path=args.dataset_metadata_path,
repeat=args.dataset_repeat,
data_file_keys=args.data_file_keys.split(","),
main_data_operator=UnifiedDataset.default_image_operator(
base_path=args.dataset_base_path,
max_pixels=args.max_pixels,
height=args.height,
width=args.width,
height_division_factor=16,
width_division_factor=16,
)
)
model = Flux2ImageTrainingModule(
model_paths=args.model_paths,
model_id_with_origin_paths=args.model_id_with_origin_paths,
tokenizer_path=args.tokenizer_path,
trainable_models=args.trainable_models,
lora_base_model=args.lora_base_model,
lora_target_modules=args.lora_target_modules,
lora_rank=args.lora_rank,
lora_checkpoint=args.lora_checkpoint,
preset_lora_path=args.preset_lora_path,
preset_lora_model=args.preset_lora_model,
use_gradient_checkpointing=args.use_gradient_checkpointing,
use_gradient_checkpointing_offload=args.use_gradient_checkpointing_offload,
extra_inputs=args.extra_inputs,
fp8_models=args.fp8_models,
offload_models=args.offload_models,
template_model_id_or_path=args.template_model_id_or_path,
resume_from_checkpoint=args.resume_from_checkpoint,
remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
enable_lora_hot_loading=args.enable_lora_hot_loading,
task=args.task,
device="cpu" if (args.initialize_model_on_cpu or args.enable_model_cpu_offload) else accelerator.device,
)
model_logger = ModelLogger(
args.output_path,
remove_prefix_in_ckpt=args.remove_prefix_in_ckpt,
enable_tensorboard_log=args.enable_tensorboard_log,
enable_swanlab_log=args.enable_swanlab_log,
swanlab_project=args.swanlab_project,
enable_wandb_log=args.enable_wandb_log,
wandb_project=args.wandb_project,
)
launcher_map = {
"sft:data_process": launch_data_process_task,
"direct_distill:data_process": launch_data_process_task,
"sft": launch_training_task,
"sft:train": launch_training_task,
"direct_distill": launch_training_task,
"direct_distill:train": launch_training_task,
}
launcher_map[args.task](accelerator, dataset, model, model_logger, args=args)
""".strip()
with open("train.py", "w", encoding="utf-8") as f:
f.write(code)
# Start training task
cmd = """
accelerate launch train.py \
--dataset_base_path data/dataset/1770381050168240056 \
--dataset_metadata_path data/dataset/metadata.json \
--extra_inputs "template_inputs" \
--max_pixels 1048576 \
--dataset_repeat 1 \
--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-Brightness:" \
--tokenizer_path "black-forest-labs/FLUX.2-klein-4B:tokenizer/" \
--learning_rate 1e-4 \
--num_epochs 1 \
--remove_prefix_in_ckpt "pipe.template_model." \
--output_path "models/template_brightness_training" \
--trainable_models "template_model" \
--use_gradient_checkpointing \
--find_unused_parameters \
--fp8_models "black-forest-labs/FLUX.2-klein-4B:text_encoder/*.safetensors,black-forest-labs/FLUX.2-klein-base-4B:transformer/*.safetensors"
"""
os.system(cmd)
After training is complete, package the obtained weights together with the model definition written earlier into the models/template_brightness directory to form a complete Template model:
import shutil
shutil.copy(
"models/template_brightness_training/epoch-0.safetensors",
"models/template_brightness/model.safetensors",
)
Load the trained model and generate images with different brightness by passing different scale values:
template = TemplatePipeline.from_pretrained(
torch_dtype=torch.bfloat16,
device="cuda",
model_configs=[ModelConfig("models/template_brightness")],
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.7}],
negative_template_inputs = [{"scale": 0.5}]
)
image.save("image_Brightness_light.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.3}],
negative_template_inputs = [{"scale": 0.5}]
)
image.save("image_Brightness_dark.jpg")
show_images([
Image.open("image_Brightness_light.jpg"),
Image.open("image_Brightness_dark.jpg"),
], resolution=256)