LoRA-for-Diffusers

本仓库为 AIGC 研究者提供了最简明的教程代码，只需几行即可在 Diffusers 中使用 LoRA。借助本手册，您可以轻松试用来自活跃社区（如 Huggingface 和 civitai）的任何 LoRA 模型。

目前，我们还支持 ControlNet-for-Diffusers 和 T2I-Adapter-for-Diffusers。

背景

什么是 LoRA？

大型语言模型的低秩适应（LoRA）由微软开发，旨在通过学习一对秩分解矩阵并冻结原始权重，从而减少可训练参数的数量。LoRA 的核心思想是微调模型的“残差”部分，而非整个模型：即训练 $\Delta W$ 而不是 $W$。

$$ W' = W + \Delta W $$

其中，$\Delta W$ 可进一步分解为低秩矩阵：$\Delta W = A B^T$，这里 $A \in \mathbb{R}^{n \times d}$，$B \in \mathbb{R}^{m \times d}$，且 $d << n$。这就是 LoRA 的关键所在。我们只需微调 $A$ 和 $B$，而无需更新 $W$，最终得到一个规模极小的模型，因为 $A$ 和 $B$ 显然比 $W$ 小得多。

这种训练技巧对于在大型通用基础模型上微调定制化模型非常有用。许多文生图模型都是基于官方的 Stable Diffusion 构建的。现在，借助 LoRA，您可以用更少的资源高效地训练自己的模型。

什么是 Safetensors？

Safetensors 是 Hugging Face 推出的一种新型简单格式，用于安全地存储张量（与 pickle 不同），同时保持高效（零拷贝）。由于其高效性，许多稳定扩散模型，尤其是 LoRA 模型，都以 Safetensors 格式发布。您可以在 huggingface/safetensors 中了解更多优势，并通过 pip 安装：

pip install safetensors

如何加载 LoRA 权重？

在本教程中，我们将展示如何将预训练的 LoRA 加载或插入到 Diffusers 框架中。许多有趣的项目可以在 Huggingface 和 civitai 上找到，但它们大多基于 stable-diffusion-webui 框架，这对高级开发者来说并不方便。我们深受 cloneofsimo/lora 的启发，该库专注于加载、合并和插值训练好的 LoRA 模型。我们主要讨论的是 Safetensors 格式的模型，这些模型目前与 Diffusers 的兼容性尚不理想。

完整模型

完整模型包含所有必要的模块（带有或不带 LoRA 层的基础模型），通常以 .ckpt 或 .safetensors 格式存储。下面我们提供两个示例，演示如何使用。

• Hugging Face 上的 stabilityai/stable-diffusion-2-1。

• Civitai 上的 dreamshaper。

您可以仅下载 .ckpt 或 .safetensors 文件。尽管 Diffusers 目前不支持直接加载这些文件，但它们提供了转换脚本。首先将 Diffusers 克隆到本地：

git clone https://github.com/huggingface/diffusers

cd ./diffusers

# 假设您已下载 xxx.safetensors 文件，它会将其转换为 Diffusers 格式的 save_dir。
python ./scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path xxx.safetensors  --dump_path save_dir --from_safetensors

# 假设您已下载 xxx.ckpt 文件，它会将其转换为 Diffusers 格式的 save_dir。
python ./scripts/convert_original_stable_diffusion_to_diffusers.py --checkpoint_path xxx.ckpt  --dump_path save_dir

然后，您可以加载该模型：

from diffusers import StableDiffusionPipeline

pipeline = StableDiffusionPipeline.from_pretrained(save_dir,torch_dtype=torch.float32)

仅 LoRA 模型

目前，Diffusers 尚不支持直接加载 LoRA 权重（通常为 .safetensors 格式）。在此，我们以一种不太优雅的方式进行了尝试，并提供了一个示例。

• one-piece-wano-saga-style-lora

请注意，此类文件的大小远小于完整模型，因为它仅包含额外的 LoRA 权重。在这种情况下，我们需要先加载基础模型。虽然也可以直接使用 stable-diffusion 1.5 作为基础，但为了获得满意的效果，建议下载推荐的基础模型。

我们的方法非常直接：从 .safetensors 文件中提取权重，并将 LoRA 权重合并到 Diffusers 支持的权重中。我们并未将 .safetensors 转换为其他格式，而是直接更新基础模型的权重。

我们的脚本应能适用于大多数来自 Huggingface 和 civitai 的模型；如果遇到问题，您也可以自行修改代码。相信我，这真的很简单，您完全可以做到。

# 默认合并比例为 0.75，您也可以手动设置
python convert_lora_safetensor_to_diffusers.py

我们已针对此问题向 Diffusers 提交了 PR issue，并在其中进一步封装了转换函数，使其更加灵活。如果您等不及，可以直接查看该 PR。预计很快就会被合并到 Diffusers 中！

如何训练你的LoRA？

Diffusers 提供了一个简单的 train_text_to_image_lora.py 脚本，用于训练你自己的 LoRA 模型。请按照其说明安装所需的依赖项。

export MODEL_NAME="CompVis/stable-diffusion-v1-4"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"

accelerate launch --mixed_precision="fp16" train_text_to_image_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME \
  --dataset_name=$DATASET_NAME --caption_column="text" \
  --resolution=512 --random_flip \
  --train_batch_size=1 \
  --num_train_epochs=100 --checkpointing_steps=5000 \
  --learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
  --seed=42 \
  --output_dir="sd-pokemon-model-lora" \
  --validation_prompt="cute dragon creature" --report_to="wandb"

使用上述命令训练好模型后，只需加载训练好的 LoRA 权重，即可通过 StableDiffusionPipeline 进行推理。你需要指定 output_dir 来加载 LoRA 权重，在本例中为 sd-pokemon-model-lora。

import torch
from diffusers import StableDiffusionPipeline

model_path = "your_path/sd-model-finetuned-lora-t4"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe.unet.load_attn_procs(model_path)
pipe.to("cuda")

prompt = "A pokemon with green eyes and red legs."
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("pokemon.png")

目前，Diffusers 仅支持为 UNet 训练 LoRA。我们已经提交了一个 PR，如果需要的话，请查看该 PR 或者提交一个问题。

使用 ColossalAI 框架训练 LoRA

ColossalAI 已经支持 LoRA。我们只需要在 train_dreambooth_colossalai.py 的顶部修改几行代码即可。这个示例是针对 DreamBooth 的，但你可以轻松地将其应用于常规的文本到图像训练。生成的 LoRA 权重仅适用于 UNet 中的注意力层。如果你也希望支持文本编码器，则应使用 Diffusers 中的 accelerate 框架，因为 ColossalAI 目前还不支持多模型。

from diffusers.loaders import AttnProcsLayers
from diffusers.models.cross_attention import LoRACrossAttnProcessor

# 注意这里！必须在 ColoInitContext 下初始化 UNet，而不仅仅是 LoRA 层
with ColoInitContext(device=get_current_device()):
        
    unet = UNet2DConditionModel.from_pretrained(
        args.pretrained_model_name_or_path, 
        subfolder="unet", 
        revision=args.revision, 
        low_cpu_mem_usage=False
    )
    unet.requires_grad_(False)

    # 设置正确的 LoRA 层
    lora_attn_procs = {}
    for name in unet.attn_processors.keys():
        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
        if name.startswith("mid_block"):
            hidden_size = unet.config.block_out_channels[-1]
        elif name.startswith("up_blocks"):
            block_id = int(name[len("up_blocks.")])
            hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
        elif name.startswith("down_blocks"):
            block_id = int(name[len("down_blocks.")])
            hidden_size = unet.config.block_out_channels[block_id]

        lora_attn_procs[name] = LoRACrossAttnProcessor(
            hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
        )

    unet.set_attn_processor(lora_attn_procs)
    lora_layers = AttnProcsLayers(unet.attn_processors)

# DDP
unet = gemini_zero_dpp(unet, args.placement)

# 配置 ColossalAI Zero 的优化器，并将初始缩放系数设为较大值以避免下溢
optimizer = GeminiAdamOptimizer(unet, 
                                lr=args.learning_rate, 
                                betas=(args.adam_beta1, args.adam_beta2),
                                weight_decay=args.adam_weight_decay,
                                eps=args.adam_epsilon,
                                initial_scale=2**16, 
                                clipping_norm=args.max_grad_norm)

到这里就完成了，唯一需要注意的是 UNet 的初始化方式。为了只保存 LoRA 权重：

torch_unet = get_static_torch_model(unet)
if gpc.get_local_rank(ParallelMode.DATA) == 0:
    torch_unet = torch_unet.to(torch.float32)
    torch_unet.save_attn_procs(save_path)

然后进行推理：

from diffusers import StableDiffusionPipeline
import torch

model_path = "sd-model-finetuned-lora"
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
pipe.unet.load_attn_procs(model_path)
pipe.to("cuda")

prompt = "A pokemon with green eyes and red legs."
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("pokemon.png")

你可能会发现生成的 LoRA 权重只有约 3MB 大小，这是因为默认设置的原因。如果想增大权重文件的大小，可以手动为 LoRA 层设置秩（低秩分解的维度）。

lora_attn_procs[name] = LoRACrossAttnProcessor(hidden_size=hidden_size, cross_attention_dim=cross_attention_dim, rank=128)

这样生成的 LoRA 权重就会达到约 100-200MB。需要注意的是，LoRA 层很容易过拟合，通常情况下，在小型数据集（少于 1000 张图片）上以批量大小为 64 进行 100 到 2000 步的训练就足够了。

问答

(1) 在合并 LoRA 时，我可以手动调整它的权重吗？

可以，这里的 alpha 就是 LoRA 的权重。我们已经向 diffusers 提交了一个 PR，在其中提供了一个灵活的封装函数。

(2) 我是否只能将 LoRA (.safetensors) 转换为 diffusers 支持的其他格式？

理论上可以，但我们不建议这样做，详情请参阅 这个 issue。这样做存在许多限制。例如，我们的脚本无法通用处理所有 .safetensors 文件，因为有些文件的命名方式不同。此外，目前的 diffusers 框架仅支持将 LoRA 应用到 UNet 的注意力层中，而来自 civitai 的许多 .safetensors 文件包含针对文本编码器等其他模块的 LoRA 权重。不过，用于文本编码器的 LoRA 很快应该就会得到支持。

(3) 我可以混合使用多于一个 LoRA 模型吗？

可以，只需进行两次合并即可。但请注意仔细设置 alpha（LoRA 的权重），如果 alpha 设置得过大，模型性能可能会下降。

(4) 这个项目的初衷是什么？

我们发现 civitai 平台上有很多非常出色的模型，但大多数 LoRA 权重都采用 safetensors 格式，这对 diffusers 用户来说并不方便。因此，我们编写了一个转换脚本，以便您能够在 diffusers 中使用这些 LoRA 模型。需要注意的是，我们的目标并不是 stable-diffusion-webui，后者已经非常成熟，但其 API 与 diffusers 完全不同。

LoRA-for-Diffusers 快速上手指南

本指南旨在帮助开发者快速在 diffusers 框架中加载和使用 LoRA 模型，或训练自定义 LoRA 模型。

环境准备

系统要求

• 操作系统: Linux, macOS, Windows (推荐 Linux)
• Python: 3.8+
• GPU: 推荐使用 NVIDIA GPU (需安装 CUDA)，CPU 亦可运行但速度较慢

前置依赖

请确保已安装以下核心库：

• torch (建议版本 1.12+)
• diffusers
• transformers
• accelerate
• safetensors (用于高效加载模型权重)

国内加速建议： 在安装依赖时，推荐使用清华源或阿里源以加快下载速度：

pip install torch torchvision torchaudio --index-url https://pypi.tuna.tsinghua.edu.cn/simple
pip install diffusers transformers accelerate safetensors -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

1. 克隆项目代码 获取本教程所需的转换脚本：

   git clone https://github.com/haofanwang/Lora-for-Diffusers.git
   cd Lora-for-Diffusers
   

2. 安装 Safetensors 如果尚未安装，请执行：

   pip install safetensors
   

3. 准备 Diffusers 转换脚本 若需加载完整的 .ckpt 或 .safetensors 大模型，需借用 diffusers 官方的转换脚本：

   git clone https://github.com/huggingface/diffusers
   

基本使用

场景一：加载现有的 LoRA 模型 (.safetensors)

大多数社区模型（如 Civitai 上的模型）仅提供 LoRA 权重文件（通常较小），需要配合基础模型使用。

1. 下载基础模型与 LoRA 权重

   • 基础模型：推荐 stable-diffusion-v1-5 或模型作者指定的基底。
   • LoRA 文件：例如 one-piece-wano-saga-style-lora.safetensors。

2. 合并权重 使用提供的脚本将 LoRA 权重合并到 diffusers 格式中（默认融合比例为 0.75）：

   python convert_lora_safetensor_to_diffusers.py
   

   注：脚本会自动寻找当前目录下的 safetensors 文件并处理，生成可直接被 diffusers 加载的文件夹。

3. 推理生成 编写 Python 脚本加载模型并生成图像：

   import torch
   from diffusers import StableDiffusionPipeline
   
   # 替换为你的输出目录路径
   model_path = "sd-pokemon-model-lora" 
   
   # 加载基础模型
   pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float16)
   
   # 加载 LoRA 权重
   pipe.unet.load_attn_procs(model_path)
   pipe.to("cuda")
   
   prompt = "A pokemon with green eyes and red legs."
   image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
   image.save("pokemon.png")
   

场景二：训练自定义 LoRA 模型

使用 diffusers 官方示例脚本进行训练。

1. 设置环境变量并启动训练 以下示例使用宝可梦数据集训练 LoRA：

   export MODEL_NAME="CompVis/stable-diffusion-v1-4"
   export DATASET_NAME="lambdalabs/pokemon-blip-captions"
   
   accelerate launch --mixed_precision="fp16" train_text_to_image_lora.py \
     --pretrained_model_name_or_path=$MODEL_NAME \
     --dataset_name=$DATASET_NAME --caption_column="text" \
     --resolution=512 --random_flip \
     --train_batch_size=1 \
     --num_train_epochs=100 --checkpointing_steps=5000 \
     --learning_rate=1e-04 --lr_scheduler="constant" --lr_warmup_steps=0 \
     --seed=42 \
     --output_dir="sd-pokemon-model-lora" \
     --validation_prompt="cute dragon creature" --report_to="wandb"
   

2. 使用训练好的模型 训练完成后，使用与“场景一”相同的推理代码，将 model_path 指向你的输出目录（如 sd-pokemon-model-lora）即可。

进阶：使用 ColossalAI 训练

若需更高效的显存管理，可结合 ColossalAI 框架。只需在训练脚本初始化 UNet 时加入 LoRA 层设置：

from diffusers.loaders import AttnProcsLayers
from diffusers.models.cross_attention import LoRACrossAttnProcessor

# 在 ColoInitContext 下初始化 unet
with ColoInitContext(device=get_current_device()):
    unet = UNet2DConditionModel.from_pretrained(...)
    unet.requires_grad_(False)

    # 配置 LoRA 层
    lora_attn_procs = {}
    for name in unet.attn_processors.keys():
        cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
        # ... (根据 block 类型获取 hidden_size 的逻辑) ...
        lora_attn_procs[name] = LoRACrossAttnProcessor(
            hidden_size=hidden_size, cross_attention_dim=cross_attention_dim
        )

    unet.set_attn_processor(lora_attn_procs)
    lora_layers = AttnProcsLayers(unet.attn_processors)

训练结束后保存权重：

torch_unet = get_static_torch_model(unet)
if gpc.get_local_rank(ParallelMode.DATA) == 0:
    torch_unet = torch_unet.to(torch.float32)
    torch_unet.save_attn_procs(save_path)