TCD SDXL LoRA

🚀 轨迹一致性蒸馏（Trajectory Consistency Distillation）

轨迹一致性蒸馏（TCD）受一致性模型（Consistency Models）启发，是一种新颖的蒸馏技术，可将预训练扩散模型的知识蒸馏到少步采样器中。本项目发布了推理代码和名为TCD - SDXL的模型，该模型从SDXL Base 1.0蒸馏而来。

🚀 快速开始

若要自行运行模型，可借助 🧨 Diffusers 库。

pip install diffusers transformers accelerate peft

然后克隆仓库：

git clone https://github.com/jabir-zheng/TCD.git
cd TCD

✨ 主要特性

✨TCD 具有以下优势：

• 灵活的函数评估次数（NFEs）：对于 TCD，NFEs 可以随意变化（与 Turbo 相比），且不会对结果质量产生不利影响（与潜在一致性模型（LCMs）相比，LCM 在高 NFEs 时质量会显著下降）。
• 优于教师模型：TCD 在高 NFEs 时保持卓越的生成质量，甚至超过了使用原始 SDXL 的 DPM - Solver++(2S)。值得注意的是，训练过程中未包含额外的判别器或 LPIPS 监督。
• 自由更改细节程度：在推理过程中，只需调整一个超参数 gamma 即可简单修改图像的细节程度，无需引入任何额外参数。
• 多功能性：集成了低秩自适应（LoRA）技术，TCD 可直接应用于各种具有相同骨干网络的模型，包括自定义社区模型、风格化 LoRA、控制网络（ControlNet）和 IP - 适配器（IP - Adapter），能在少步内加速高质量图像生成。
• 避免模式崩溃：TCD 无需对抗训练即可实现少步生成，从而避免了由生成对抗网络（GAN）目标导致的模式崩溃。与同期工作 SDXL - Lightning 相比，TCD 合成的结果更逼真、更多样，且不存在 “双面人” 伪影。

📦 安装指南

运行模型前，需安装必要的依赖库：

pip install diffusers transformers accelerate peft

然后克隆项目仓库：

git clone https://github.com/jabir-zheng/TCD.git
cd TCD

💻 使用示例

基础用法

文本到图像生成

import torch
from diffusers import StableDiffusionXLPipeline
from scheduling_tcd import TCDScheduler 

device = "cuda"
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
tcd_lora_id = "h1t/TCD-SDXL-LoRA"

pipe = StableDiffusionXLPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16, variant="fp16").to(device)
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(tcd_lora_id)
pipe.fuse_lora()

prompt = "Beautiful woman, bubblegum pink, lemon yellow, minty blue, futuristic, high-detail, epic composition, watercolor."

image = pipe(
    prompt=prompt,
    num_inference_steps=4,
    guidance_scale=0,
    # Eta（论文中称为 `gamma`）用于控制每一步的随机性。
    # 值为 0.3 通常会产生良好的结果。
    # 建议在增加推理步数时使用更高的 eta。
    eta=0.3, 
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]

高级用法

图像修复

import torch
from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid
from scheduling_tcd import TCDScheduler 

device = "cuda"
base_model_id = "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
tcd_lora_id = "h1t/TCD-SDXL-LoRA"

pipe = AutoPipelineForInpainting.from_pretrained(base_model_id, torch_dtype=torch.float16, variant="fp16").to(device)
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(tcd_lora_id)
pipe.fuse_lora()

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = load_image(img_url).resize((1024, 1024))
mask_image = load_image(mask_url).resize((1024, 1024))

prompt = "a tiger sitting on a park bench"

image = pipe(
  prompt=prompt,
  image=init_image,
  mask_image=mask_image,
  num_inference_steps=8,
  guidance_scale=0,
  eta=0.3, # Eta（论文中称为 `gamma`）用于控制每一步的随机性。值为 0.3 通常会产生良好的结果。
  strength=0.99,  # 确保使用小于 1.0 的 `strength`
  generator=torch.Generator(device=device).manual_seed(0),
).images[0]

grid_image = make_image_grid([init_image, mask_image, image], rows=1, cols=3)

适用于社区模型

import torch
from diffusers import StableDiffusionXLPipeline
from scheduling_tcd import TCDScheduler 

device = "cuda"
base_model_id = "cagliostrolab/animagine-xl-3.0"
tcd_lora_id = "h1t/TCD-SDXL-LoRA"

pipe = StableDiffusionXLPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16, variant="fp16").to(device)
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(tcd_lora_id)
pipe.fuse_lora()

prompt = "A man, clad in a meticulously tailored military uniform, stands with unwavering resolve. The uniform boasts intricate details, and his eyes gleam with determination. Strands of vibrant, windswept hair peek out from beneath the brim of his cap."

image = pipe(
    prompt=prompt,
    num_inference_steps=8,
    guidance_scale=0,
    # Eta（论文中称为 `gamma`）用于控制每一步的随机性。
    # 值为 0.3 通常会产生良好的结果。
    # 建议在增加推理步数时使用更高的 eta。
    eta=0.3, 
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]

与风格化 LoRA 结合

import torch
from diffusers import StableDiffusionXLPipeline
from scheduling_tcd import TCDScheduler 

device = "cuda"
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
tcd_lora_id = "h1t/TCD-SDXL-LoRA"
styled_lora_id = "TheLastBen/Papercut_SDXL"

pipe = StableDiffusionXLPipeline.from_pretrained(base_model_id, torch_dtype=torch.float16, variant="fp16").to(device)
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(tcd_lora_id, adapter_name="tcd")
pipe.load_lora_weights(styled_lora_id, adapter_name="style")
pipe.set_adapters(["tcd", "style"], adapter_weights=[1.0, 1.0])

prompt = "papercut of a winter mountain, snow"

image = pipe(
    prompt=prompt,
    num_inference_steps=4,
    guidance_scale=0,
    # Eta（论文中称为 `gamma`）用于控制每一步的随机性。
    # 值为 0.3 通常会产生良好的结果。
    # 建议在增加推理步数时使用更高的 eta。
    eta=0.3, 
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]

与控制网络（ControlNet）兼容

深度控制网络（Depth ControlNet）

import torch
import numpy as np
from PIL import Image
from transformers import DPTFeatureExtractor, DPTForDepthEstimation
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline
from diffusers.utils import load_image, make_image_grid
from scheduling_tcd import TCDScheduler 

device = "cuda"
depth_estimator = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas").to(device)
feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-hybrid-midas")

def get_depth_map(image):
    image = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
    with torch.no_grad(), torch.autocast(device):
        depth_map = depth_estimator(image).predicted_depth

    depth_map = torch.nn.functional.interpolate(
        depth_map.unsqueeze(1),
        size=(1024, 1024),
        mode="bicubic",
        align_corners=False,
    )
    depth_min = torch.amin(depth_map, dim=[1, 2, 3], keepdim=True)
    depth_max = torch.amax(depth_map, dim=[1, 2, 3], keepdim=True)
    depth_map = (depth_map - depth_min) / (depth_max - depth_min)
    image = torch.cat([depth_map] * 3, dim=1)

    image = image.permute(0, 2, 3, 1).cpu().numpy()[0]
    image = Image.fromarray((image * 255.0).clip(0, 255).astype(np.uint8))
    return image

base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
controlnet_id = "diffusers/controlnet-depth-sdxl-1.0"
tcd_lora_id = "h1t/TCD-SDXL-LoRA"

controlnet = ControlNetModel.from_pretrained(
    controlnet_id,
    torch_dtype=torch.float16,
    variant="fp16",
).to(device)
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    base_model_id,
    controlnet=controlnet,
    torch_dtype=torch.float16,
    variant="fp16",
).to(device)
pipe.enable_model_cpu_offload()

pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(tcd_lora_id)
pipe.fuse_lora()

prompt = "stormtrooper lecture, photorealistic"

image = load_image("https://huggingface.co/lllyasviel/sd-controlnet-depth/resolve/main/images/stormtrooper.png")
depth_image = get_depth_map(image)

controlnet_conditioning_scale = 0.5  # 建议用于良好的泛化

image = pipe(
    prompt, 
    image=depth_image, 
    num_inference_steps=4, 
    guidance_scale=0,
    eta=0.3, # 一个参数（论文中称为 `gamma`）用于控制每一步的随机性。值为 0.3 通常会产生良好的结果。
    controlnet_conditioning_scale=controlnet_conditioning_scale,
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]

grid_image = make_image_grid([depth_image, image], rows=1, cols=2)

坎尼控制网络（Canny ControlNet）

import torch
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline
from diffusers.utils import load_image, make_image_grid
from scheduling_tcd import TCDScheduler 

device = "cuda"
base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
controlnet_id = "diffusers/controlnet-canny-sdxl-1.0"
tcd_lora_id = "h1t/TCD-SDXL-LoRA"

controlnet = ControlNetModel.from_pretrained(
    controlnet_id,
    torch_dtype=torch.float16,
    variant="fp16",
).to(device)
pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
    base_model_id,
    controlnet=controlnet,
    torch_dtype=torch.float16,
    variant="fp16",
).to(device)
pipe.enable_model_cpu_offload()

pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(tcd_lora_id)
pipe.fuse_lora()

prompt = "ultrarealistic shot of a furry blue bird"

canny_image = load_image("https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png")

controlnet_conditioning_scale = 0.5  # 建议用于良好的泛化

image = pipe(
    prompt, 
    image=canny_image, 
    num_inference_steps=4, 
    guidance_scale=0,
    eta=0.3, # 一个参数（论文中称为 `gamma`）用于控制每一步的随机性。值为 0.3 通常会产生良好的结果。
    controlnet_conditioning_scale=controlnet_conditioning_scale,
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]

grid_image = make_image_grid([canny_image, image], rows=1, cols=2)

与 IP - 适配器（IP - Adapter）兼容

⚠️ 重要提示

请参考官方 仓库 获取 IP - Adapter 依赖项的安装说明。

import torch
from diffusers import StableDiffusionXLPipeline
from diffusers.utils import load_image, make_image_grid

from ip_adapter import IPAdapterXL
from scheduling_tcd import TCDScheduler 

device = "cuda"
base_model_path = "stabilityai/stable-diffusion-xl-base-1.0"
image_encoder_path = "sdxl_models/image_encoder"
ip_ckpt = "sdxl_models/ip-adapter_sdxl.bin"
tcd_lora_id = "h1t/TCD-SDXL-LoRA"

pipe = StableDiffusionXLPipeline.from_pretrained(
    base_model_path, 
    torch_dtype=torch.float16, 
    variant="fp16"
)
pipe.scheduler = TCDScheduler.from_config(pipe.scheduler.config)

pipe.load_lora_weights(tcd_lora_id)
pipe.fuse_lora()

ip_model = IPAdapterXL(pipe, image_encoder_path, ip_ckpt, device)

ref_image = load_image("https://raw.githubusercontent.com/tencent-ailab/IP-Adapter/main/assets/images/woman.png").resize((512, 512))

prompt = "best quality, high quality, wearing sunglasses"

image = ip_model.generate(
    pil_image=ref_image, 
    prompt=prompt,
    scale=0.5,
    num_samples=1, 
    num_inference_steps=4, 
    guidance_scale=0,
    eta=0.3, # 一个参数（论文中称为 `gamma`）用于控制每一步的随机性。值为 0.3 通常会产生良好的结果。
    seed=0,
)[0]

grid_image = make_image_grid([ref_image, image], rows=1, cols=2)

📚 详细文档

关于抄袭指控的郑重声明

我们遗憾地听到来自 CTM 团队的严重指控。在此次事件之前，我们已经与 CTM 的作者进行了多轮沟通，在此详细说明情况：

1. 在 arXiv 第一版 中，我们在 “A. 相关工作” 部分提供了引用和讨论：

   Kim 等人（2023）提出了一个用于一致性模型（CMs）和扩散模型（DMs）的通用框架。核心设计与我们的相似，主要区别在于我们专注于减少 CMs 中的误差，巧妙地利用概率流常微分方程（PF ODE）的半线性结构进行参数化，并避免了对抗训练的需要。

2. 在同一版本的 “D.3 定理 4.2 的证明” 中，我们指出：

   在本节中，我们的推导主要借鉴了（Kim 等人，2023；Chen 等人，2022）的证明。 我们从未试图独占功劳。正如我们在邮件中提到的，我们愿意就论文中引用明显不足的问题向 CTM 作者正式道歉，并将在修订稿中提供更多引用。

3. 在更新后的 arXiv 第二版 中，我们扩展了讨论以阐明与 CTM 框架的关系，并删除了一些之前为完整性而包含的证明。
4. CTM 和 TCD 在动机、方法和实验方面均有所不同。TCD 基于潜在一致性模型（LCM）的原理，旨在通过利用 指数积分器 设计有效的一致性函数。
5. 实验结果也无法通过任何类型的 CTM 算法获得。
   • 我们提供一个简单的验证方法：使用我们的采样器对 CTM 发布的检查点 进行采样，反之亦然。
   • CTM 也提供了训练脚本，欢迎任何人基于 CTM 算法在 SDXL 或潜在扩散模型（LDM）上重现实验。

我们认为抄袭指控不仅严重，而且损害了相关方的学术诚信，真诚希望各方能更全面地了解此事。

相关和同期工作

• Luo S, Tan Y, Huang L, 等. 潜在一致性模型：通过少步推理合成高分辨率图像. arXiv 预印本 arXiv:2310.04378, 2023.
• Luo S, Tan Y, Patil S, 等. LCM - LoRA：一种通用的稳定扩散加速模块. arXiv 预印本 arXiv:2311.05556, 2023.
• Lu C, Zhou Y, Bao F, 等. DPM - Solver：一种用于扩散概率模型采样的快速常微分方程求解器，约 10 步完成. 神经信息处理系统进展, 2022, 35: 5775 - 5787.
• Lu C, Zhou Y, Bao F, 等. DPM - solver++：用于扩散概率模型引导采样的快速求解器. arXiv 预印本 arXiv:2211.01095, 2022.
• Zhang Q, Chen Y. 使用指数积分器对扩散模型进行快速采样. ICLR 2023, 卢旺达基加利, 2023 年 5 月 1 - 5 日.
• Kim D, Lai C H, Liao W H, 等. 一致性轨迹模型：学习扩散的概率流常微分方程轨迹. ICLR 2024.

引用

如果您使用了本项目的代码或模型，请引用以下论文：

@misc{zheng2024trajectory,
      title={Trajectory Consistency Distillation}, 
      author={Jianbin Zheng and Minghui Hu and Zhongyi Fan and Chaoyue Wang and Changxing Ding and Dacheng Tao and Tat-Jen Cham},
      year={2024},
      eprint={2402.19159},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

致谢

本代码库严重依赖 🤗Diffusers 库和 LCM。

📄 许可证

本项目采用 MIT 许可证。

其他信息

本项目是论文 Trajectory Consistency Distillation 的官方模型仓库。更多信息请查看 GitHub 仓库 和 项目页面。也欢迎尝试在 🤗 Space 上的演示。