Skip to content

[[open-in-colab]]

轨迹一致性蒸馏-LoRA

轨迹一致性蒸馏(TCD)使模型能够在更少的步骤中生成更高质量和更详细图像。此外,由于在蒸馏过程中有效减少了误差,TCD 即使在大量推理步骤下也能表现出优越的性能。

TCD 的主要优势包括:

  • 超越教师模型:TCD 在小步长和大步长推理中都表现出优越的生成质量,并且在使用 Stable Diffusion XL (SDXL) 时超过了 DPM-Solver++(2S) 的性能。TCD 训练过程中没有包含额外的判别器或 LPIPS 监督。

  • 灵活的推理步长:TCD 采样的推理步长可以自由调整,而不会对图像质量产生负面影响。

  • 自由调整细节水平:在推理过程中,可以通过单个超参数 gamma 调整图像的细节水平。

TIP

有关 TCD 的更多技术细节,请参阅 论文 或官方 项目页面

对于像 SDXL 这样的大型模型,TCD 使用 LoRA 进行训练以减少内存使用。这也很有用,因为只要不同微调模型共享相同的基模型,你就可以在它们之间重用 LoRA,而无需进一步训练。

本指南将向你展示如何使用 TCD-LoRA 进行各种任务的推理,如文本到图像和修复,以及如何轻松地将 TCD-LoRA 与其他适配器结合使用。请选择下表中支持的基模型及其对应的 TCD-LoRA 检查点以开始。

Base modelTCD-LoRA checkpoint
stable-diffusion-v1-5TCD-SD15
stable-diffusion-2-1-baseTCD-SD21-base
stable-diffusion-xl-base-1.0TCD-SDXL

确保你已经安装了 PEFT,以获得更好的 LoRA 支持。

bash
pip install -U peft

通用任务

在本指南中,我们将使用 [StableDiffusionXLPipeline] 和 [TCDScheduler]。使用 [~StableDiffusionPipeline.load_lora_weights] 方法加载与 SDXL 兼容的 TCD-LoRA 权重。

在进行 TCD-LoRA 推理时,需要注意以下几点:

  • num_inference_steps 保持在 4 到 50 之间
  • eta(用于控制每一步的随机性)设置在 0 到 1 之间。当你增加推理步骤的数量时,应使用更高的 eta,但缺点是较大的 eta 会导致使用 [TCDScheduler] 生成的图像更模糊。建议使用 0.3 以获得良好的结果。

社区模型

TCD-LoRA 也支持许多社区微调的模型和插件。例如,加载 animagine-xl-3.0 检查点,这是一个社区微调的 SDXL 版本,用于生成动漫图像。

python
import torch
from diffusers import StableDiffusionXLPipeline, 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=0.3,
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]

TCD-LoRA 还支持其他在不同风格上训练的 LoRA。例如,让我们加载 TheLastBen/Papercut_SDXL LoRA,并使用 [~loaders.UNet2DConditionLoadersMixin.set_adapters] 方法将其与 TCD-LoRA 融合。

TIP

查看 合并 LoRA 指南,了解更多的高效合并方法。

python
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=0.3,
    generator=torch.Generator(device=device).manual_seed(0),
).images[0]

适配器

TCD-LoRA 非常灵活,可以与其他适配器类型(如 ControlNets、IP-Adapter 和 AnimateDiff)结合使用。