魔搭社区 AIGC 系列课程 - 可控生成技术
本实验以 Diffusion-Templates 为框架，系统介绍图像生成模型的多种可控生成技术，并演示如何自行训练一个可控生成模块。
相关资料：
开源代码：DiffSynth-Studio
技术报告：arXiv
项目主页：GitHub
文档参考：English Version、中文版
在线体验：魔搭社区创空间
!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
图像结构控制
首先，加载基础模型 black-forest-labs/FLUX.2-klein-base-4B。这是一个参数量为 4B 的图像生成模型，本实验后续所有可控生成模块都会挂载到这个基础模型之上。
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 是最早的一批 Diffusion 可控生成技术，可用深度图、边缘图、姿态图等结构性条件对生成画面进行逐像素级的控制。
以 Template 格式加载 DiffSynth-Studio/Template-KleinBase4B-ControlNet，即可在保留输入结构的前提下，用不同的提示词生成不同风格的画面。
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)
数值属性控制
AttriCtrl 是一类数值型可控生成模型，能够将连续的数值属性作为控制条件注入生成过程。
运行以下代码，加载 DiffSynth-Studio/Template-KleinBase4B-SoftRGB，通过输入 R/G/B 数值精确控制画面的整体色调。
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)
图像编辑
图像编辑模型是一类通用性较强的可控生成模型：给定一张原图和一段编辑指令，即可对原图进行局部或整体修改。
运行以下代码，加载 DiffSynth-Studio/Template-KleinBase4B-Edit。该模型通过 KV-Cache 复用输入图像的注意力键值，从而快速完成编辑，推理速度较快。
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)
风格控制
实现图像风格控制的最直接方式，是训练一个风格 LoRA——但每种风格都需要单独训练，成本较高。为此我们训练了一个特殊的 Image-to-LoRA 模型，它可以根据输入的参考图像即时生成一份 LoRA 权重，免去了传统的风格训练过程。
运行以下代码，加载 DiffSynth-Studio/KleinBase4B-i2L-v2，用参考图像动态生成 LoRA，从而控制画面风格。
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)
训练可控生成模型
Diffusion-Templates 框架允许开发者训练任意结构的可控生成模型——只要给定模型定义、数据处理逻辑和数据集，即可接入统一的训练流程。下面我们从零训练一个亮度控制模型，让画面按指定的亮度数值生成。
第一步，编写模型结构代码（包含数值编码器、KV-Cache 生成主干和数据标注器）：
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
# 主干模型结构（将输入的数值转换为 KV-Cache 向量）
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}
# 将图像数据转换为模型输入（根据图像中的 RGB 数值计算亮度）
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())
第二步，下载并预处理数据集，同时生成训练所需的 metadata：
import json, os
from modelscope import dataset_snapshot_download
# 下载数据集
dataset_snapshot_download(
"DiffSynth-Studio/ImagePulseV2-TextImage",
local_dir="data/ImagePulseV2-TextImage",
allow_file_pattern="data/1770381050168240056.tar.gz"
)
# 解压数据集
os.makedirs("data/dataset", exist_ok=True)
os.system("tar zxvf data/ImagePulseV2-TextImage/data/1770381050168240056.tar.gz -C data/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)
第三步，启动训练：
import os
# 训练脚本
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)
# 启动训练任务
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)
训练完成后，将得到的权重与前面写好的模型定义一起打包到 models/template_brightness 目录，形成一个完整的 Template 模型：
import shutil
shutil.copy(
"models/template_brightness_training/epoch-0.safetensors",
"models/template_brightness/model.safetensors",
)
加载训练好的模型，通过传入不同的 scale 数值生成明暗不同的图像：
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)