EasyControl 的实现

EasyControl：为扩散 Transformer 添加高效灵活的控制机制

张宇轩、袁一睿、宋怡仁、王浩凡、刘嘉铭]
香港中文大学、上海科技大学、Tiamat AI、阿里巴巴集团、新加坡国立大学、Liblib AI

特性

• 动机: 扩散模型的架构正从基于 Unet 转向 DiT（扩散 Transformer）。然而，DiT 生态系统缺乏成熟的插件支持，并面临效率瓶颈、多条件协调冲突以及模型适应性不足等挑战。
• 贡献: 我们提出了 EasyControl，一个高效且灵活的统一条件 DiT 框架。通过引入轻量级条件注入 LoRA 模块、位置感知训练范式，以及因果注意力机制与 KV 缓存技术的结合，我们显著提升了 模型兼容性（实现即插即用功能和风格无损控制）、生成灵活性（支持多种分辨率、宽高比及多条件组合）和 推理效率。

新闻

• 2025-04-11: 🔥🔥🔥 训练代码已发布。推荐硬件：至少 1 张 NVIDIA H100/H800/A100 显卡，显存约 80GB。

• 2025-04-09: ⭐️ 简易 API 的代码已发布。如果您希望在个人设备上运行这些模型，请前往 simple_api 分支获取相关资源。

• 2025-04-07: 🔥 感谢 CFG-Zero* 团队的出色工作，EasyControl 现已与 CFG-Zero* 集成！！只需几行代码，您就可以大幅提升图像质量和可控性！！您可以从 此链接 下载修改后的代码并尝试。

源图像	CFG	CFG-Zero*

• 2025-04-03: 感谢 jax-explorer，现已支持 Ghibli Img2Img 控制 ComfyUI 节点！

• 2025-04-01: 🔥 新的 风格化 Img2Img 控制模型现已发布！！使用此 LoRA 模型将人像转化为吉卜力工作室风格的艺术作品。该模型仅基于 100 张真实的亚洲人脸 和 由 GPT-4o 生成的吉卜力风格对应图 进行训练，能够在保留面部特征的同时应用标志性的动漫美学。

• 2025-03-19: 🔥 我们已在 Hugging Face 上发布了 演示空间! 您现在可以通过 Hugging Face 空间试用 EasyControl，尽情享受吧！

• 2025-03-18: 🔥 我们已在 Hugging Face 上发布了 预训练检查点！您现在可以使用官方权重试用 EasyControl。
• 2025-03-12: ⭐️ 推理代码已发布。待确认一切正常后，新模型将合并到本仓库中。敬请期待后续更新！😊

安装

我们建议使用 Python 3.10 和支持 CUDA 的 PyTorch。搭建环境的步骤如下：

# 创建一个新的 conda 环境
conda create -n easycontrol python=3.10
conda activate easycontrol

# 安装其他依赖
pip install -r requirements.txt

下载

您可以直接从 Hugging Face 下载模型。 或者使用 Python 脚本下载：

from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/canny.safetensors", local_dir="./")
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/depth.safetensors", local_dir="./")
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/hedsketch.safetensors", local_dir="./")
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/inpainting.safetensors", local_dir="./")
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/pose.safetensors", local_dir="./")
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/seg.safetensors", local_dir="./")
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/subject.safetensors", local_dir="./")
hf_hub_download(repo_id="Xiaojiu-Z/EasyControl", filename="models/Ghibli.safetensors", local_dir="./")

如果您无法访问 Hugging Face，可以使用 hf-mirror 下载模型：

export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download --resume-download Xiaojiu-Z/EasyControl --local-dir checkpoints --local-dir-use-symlinks False

使用方法

以下是使用 EasyControl 的基本示例：

模型初始化

import torch
from PIL import Image
from src.pipeline import FluxPipeline
from src.transformer_flux import FluxTransformer2DModel
from src.lora_helper import set_single_lora, set_multi_lora

def clear_cache(transformer):
    for name, attn_processor in transformer.attn_processors.items():
        attn_processor.bank_kv.clear()

# 初始化模型
device = "cuda"
base_path = "FLUX.1-dev"  # 您的基础模型路径
pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16, device=device)
transformer = FluxTransformer2DModel.from_pretrained(
    base_path, 
    subfolder="transformer",
    torch_dtype=torch.bfloat16, 
    device=device
)
pipe.transformer = transformer
pipe.to(device)

# 加载控制模型
lora_path = "./checkpoints/models"
control_models = {
    "canny": f"{lora_path}/canny.safetensors",
    "depth": f"{lora_path}/depth.safetensors",
    "hedsketch": f"{lora_path}/hedsketch.safetensors",
    "pose": f"{lora_path}/pose.safetensors",
    "seg": f"{lora_path}/seg.safetensors",
    "inpainting": f"{lora_path}/inpainting.safetensors",
    "subject": f"{lora_path}/subject.safetensors",
}

单条件控制

# 单空间条件控制示例
path = control_models["canny"]
set_single_lora(pipe.transformer, path, lora_weights=[1], cond_size=512)

# 生成图像
prompt = "海滩上的一辆漂亮的汽车"
spatial_image = Image.open("./https://oss.gittoolsai.com/images/Xiaojiu-z_EasyControl_readme_dd888779b6f7.png").convert("RGB")

image = pipe(
    prompt,
    height=720,
    width=992,
    guidance_scale=3.5,
    num_inference_steps=25,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(5),
    spatial_images=[spatial_image],
    cond_size=512,
).images[0]

# 生成后清除缓存
clear_cache(pipe.transformer)

# 单主体条件控制示例
path = control_models["subject"]
set_single_lora(pipe.transformer, path, lora_weights=[1], cond_size=512)

# 生成图像
prompt = "图书馆里的SKS步枪"
subject_image = Image.open("./https://oss.gittoolsai.com/images/Xiaojiu-z_EasyControl_readme_93f777c8939a.png").convert("RGB")

image = pipe(
    prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    num_inference_steps=25,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(5),
    subject_images=[subject_image],
    cond_size=512,
).images[0]

# 生成后清除缓存
clear_cache(pipe.transformer)

多条件控制

# 多条件控制示例
paths = [control_models["subject"], control_models["inpainting"]]
set_multi_lora(pipe.transformer, paths, lora_weights=[[1], [1]], cond_size=512)

prompt = "车上的SKS步枪"
subject_images = [Image.open("./https://oss.gittoolsai.com/images/Xiaojiu-z_EasyControl_readme_1ba5048b3d95.png").convert("RGB")]
spatial_images = [Image.open("./https://oss.gittoolsai.com/images/Xiaojiu-z_EasyControl_readme_ca5e99b867e2.png").convert("RGB")]

image = pipe(
    prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    num_inference_steps=25,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(42),
    subject_images=subject_images,
    spatial_images=spatial_images,
    cond_size=512,
).images[0]

# 生成后清除缓存
clear_cache(pipe.transformer)

吉卜力风格人像生成

import spaces
import os
import json
import time
import torch
from PIL import Image
from tqdm import tqdm
import gradio as gr

from safetensors.torch import save_file
from src.pipeline import FluxPipeline
from src.transformer_flux import FluxTransformer2DModel
from src.lora_helper import set_single_lora, set_multi_lora, unset_lora

# 初始化图像处理器
base_path = "black-forest-labs/FLUX.1-dev"    
lora_base_path = "./checkpoints/models"


pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16)
transformer = FluxTransformer2DModel.from_pretrained(base_path, subfolder="transformer", torch_dtype=torch.bfloat16)
pipe.transformer = transformer
pipe.to("cuda")

def clear_cache(transformer):
    for name, attn_processor in transformer.attn_processors.items():
        attn_processor.bank_kv.clear()

# 定义 Gradio 界面
@spaces.GPU()
def single_condition_generate_image(prompt, spatial_img, height, width, seed, control_type):
    # 设置控制类型
    if control_type == "Ghibli":
        lora_path = os.path.join(lora_base_path, "Ghibli.safetensors")
    set_single_lora(pipe.transformer, lora_path, lora_weights=[1], cond_size=512)
    
    # 处理图像
    spatial_imgs = [spatial_img] if spatial_img else []
    image = pipe(
        prompt,
        height=int(height),
        width=int(width),
        guidance_scale=3.5,
        num_inference_steps=25,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(seed), 
        subject_images=[],
        spatial_images=spatial_imgs,
        cond_size=512,
    ).images[0]
    clear_cache(pipe.transformer)
    return image

# 定义 Gradio 界面组件
control_types = ["Ghibli"]

# 创建 Gradio Blocks 界面
with gr.Blocks() as demo:
    gr.Markdown("# 使用 EasyControl 控制吉卜力工作室风格图像生成")
    gr.Markdown("该模型仅基于**100张真实的亚洲人脸**与**由 GPT-4o 生成的吉卜力风格对应图像**进行训练，在保留面部特征的同时，能够应用标志性的动漫美学风格。")
    gr.Markdown("使用带有吉卜力控制 LoRA 的 EasyControl 生成图像。（由于硬件限制，目前只能生成低分辨率图像。如需高分辨率（1024+），请自行搭建环境。）")
    
    gr.Markdown("**[注意！！]**：推荐用于吉卜力控制 LoRA 的提示词应包含触发词：`吉卜力工作室风格、迷人手绘动漫风格插画`")
    gr.Markdown("😊😊如果您喜欢这个演示，请给我们点个赞（GitHub：[EasyControl](https://github.com/Xiaojiu-z/EasyControl))")

    with gr.Tab("吉卜力条件生成"):
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="提示词", value="吉卜力工作室风格、迷人手绘动漫风格插画")
                spatial_img = gr.Image(label="吉卜力图像", type="pil")  # 上传图像文件
                height = gr.Slider(minimum=256, maximum=1024, step=64, label="高度", value=768)
                width = gr.Slider(minimum=256, maximum=1024, step=64, label="宽度", value=768)
                seed = gr.Number(label="种子", value=42)
                control_type = gr.Dropdown(choices=control_types, label="控制类型")
                single_generate_btn = gr.Button("生成图像")
            with gr.Column():
                single_output_image = gr.Image(label="生成图像")


    # 将按钮与函数关联
    single_generate_btn.click(
        single_condition_generate_image,
        inputs=[prompt, spatial_img, height, width, seed, control_type],
        outputs=single_output_image
    )

# 启动 Gradio 应用
demo.queue().launch()

使用提示

• 每次生成后请使用 clear_cache(pipe.transformer) 清除缓存。
• 为获得最佳效果：
  • 建议从 guidance_scale=3.5 开始，并根据结果调整。
  • 使用 num_inference_steps=25 可在质量和速度之间取得良好平衡。
• 使用 set_multi_lora API 时，请确保主体 LoRA 路径（subject）位于空间 LoRA 路径（canny、depth、hedsketch 等）之前。

待办事项清单

1. • 推理代码
2. • 空间预训练权重
3. • 主体预训练权重
4. • 训练代码

星标历史

免责声明

EasyControl 的代码采用 Apache 许可证 发布，适用于学术和商业用途。我们发布的检查点仅供研究目的使用。用户可以自由地使用此工具生成图像，但必须遵守当地法律并负责任地使用。开发者对用户的任何潜在滥用行为不承担任何责任。

招聘/合作

• 2025年4月3日：如果您对 EasyControl 及其相关技术感兴趣，或者希望以低成本方式构建类似 4o 的能力，我们可以在上海、北京、香港、新加坡等地线下合作，或通过线上方式进行交流。 联系方式：jmliu1217@gmail.com（微信：jiaming068870）

引用

@article{zhang2025easycontrol,
  title={EasyControl: 为扩散 Transformer 添加高效灵活的控制},
  author={Zhang, Yuxuan 和 Yuan, Yirui 和 Song, Yiren 和 Wang, Haofan 和 Liu, Jiaming},
  journal={arXiv 预印本 arXiv:2503.07027},
  year={2025}
}

EasyControl 快速上手指南

EasyControl 是一个高效且灵活的扩散 Transformer (DiT) 控制框架，旨在为 FLUX.1 等模型提供插件式的条件控制（如边缘、深度、姿态、主体保持等），支持多条件组合与高分辨率生成。

环境准备

• 操作系统: Linux / Windows (推荐 Linux)
• Python 版本: 3.10
• GPU 要求: 支持 CUDA 的 NVIDIA 显卡
  • 推理: 建议显存 ≥ 24GB (取决于分辨率)
  • 训练: 建议至少 1x NVIDIA H100/H800/A100 (约 80GB 显存)
• 前置依赖: PyTorch (带 CUDA 支持), Conda (推荐)

安装步骤

1. 创建虚拟环境

推荐使用 Conda 创建独立的 Python 3.10 环境：

conda create -n easycontrol python=3.10
conda activate easycontrol

2. 安装依赖库

进入项目目录后安装所需包：

pip install -r requirements.txt

3. 下载模型权重

你可以从 Hugging Face 下载预训练模型。如果直接访问受限，强烈推荐使用国内镜像源加速下载。

方式一：使用 Python 脚本下载 (推荐)

from huggingface_hub import hf_hub_download

# 定义需要下载的模型列表
models = [
    "canny.safetensors", "depth.safetensors", "hedsketch.safetensors",
    "inpainting.safetensors", "pose.safetensors", "seg.safetensors",
    "subject.safetensors", "Ghibli.safetensors"
]

for model_name in models:
    hf_hub_download(
        repo_id="Xiaojiu-Z/EasyControl", 
        filename=f"models/{model_name}", 
        local_dir="./checkpoints"
    )

方式二：使用命令行 + 国内镜像 (hf-mirror) 如果无法访问官方源，请设置镜像地址后下载：

export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download --resume-download Xiaojiu-Z/EasyControl --local-dir checkpoints --local-dir-use-symlinks False

注意: 使用前请确保已下载基础模型 FLUX.1-dev，并在代码中指定其路径。

基本使用

以下是最简单的单条件控制（以 Canny 边缘检测为例）生成流程。

1. 初始化模型

加载 FLUX.1 基础模型并注入 EasyControl 组件。

import torch
from PIL import Image
from src.pipeline import FluxPipeline
from src.transformer_flux import FluxTransformer2DModel
from src.lora_helper import set_single_lora

def clear_cache(transformer):
    for name, attn_processor in transformer.attn_processors.items():
        attn_processor.bank_kv.clear()

# 配置设备与路径
device = "cuda"
base_path = "FLUX.1-dev"  # 替换为你本地的 FLUX.1-dev 模型路径
lora_path = "./checkpoints/models"

# 初始化 Pipeline 和 Transformer
pipe = FluxPipeline.from_pretrained(base_path, torch_dtype=torch.bfloat16, device=device)
transformer = FluxTransformer2DModel.from_pretrained(
    base_path, 
    subfolder="transformer",
    torch_dtype=torch.bfloat16, 
    device=device
)
pipe.transformer = transformer
pipe.to(device)

# 加载控制模型 (例如 Canny)
control_model_path = f"{lora_path}/canny.safetensors"
set_single_lora(pipe.transformer, control_model_path, lora_weights=[1], cond_size=512)

2. 生成图像

准备一张条件图（如边缘图），输入提示词进行生成。

# 准备输入
prompt = "A nice car on the beach"
spatial_image = Image.open("./test_imgs/canny.png").convert("RGB") # 替换为你的条件图路径

# 执行生成
image = pipe(
    prompt,
    height=720,
    width=992,
    guidance_scale=3.5,
    num_inference_steps=25,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(5),
    spatial_images=[spatial_image], # 传入条件图像列表
    cond_size=512,
).images[0]

# 生成完成后清理缓存 (重要)
clear_cache(pipe.transformer)

# 保存结果
image.save("result.png")

3. 多条件控制 (可选)

EasyControl 支持同时叠加多个条件（例如：主体保持 + 局部重绘）。

from src.lora_helper import set_multi_lora

# 加载多个 LoRA (例如：主体 + 重绘)
paths = [f"{lora_path}/subject.safetensors", f"{lora_path}/inpainting.safetensors"]
set_multi_lora(pipe.transformer, paths, lora_weights=[[1], [1]], cond_size=512)

prompt = "A SKS on the car"
subject_images = [Image.open("./test_imgs/subject_1.png").convert("RGB")]
spatial_images = [Image.open("./test_imgs/inpainting.png").convert("RGB")]

image = pipe(
    prompt,
    height=1024,
    width=1024,
    guidance_scale=3.5,
    num_inference_steps=25,
    max_sequence_length=512,
    generator=torch.Generator("cpu").manual_seed(42),
    subject_images=subject_images,
    spatial_images=spatial_images,
    cond_size=512,
).images[0]

clear_cache(pipe.transformer)
image.save("multi_condition_result.png")