2. 文本到图像生成LoRA微调

任务说明

本案例基于 Stable Diffusion 与 LoRA 实现文本到图像生成模型的风格微调。目标是在消费级 GPU（显存 ≥ 8GB）上，使用约 20～30 张同一风格的图片，训练出能通过文本提示词稳定生成该风格图像的 LoRA 权重。训练完成后，得到一个 .safetensors 权重文件，可直接在 WebUI 或 ComfyUI 中加载使用。

环境与依赖

推荐使用 Python 3.10 + CUDA 11.8 环境，核心依赖如下：

pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install diffusers[training] accelerate peft datasets pillow
pip install safetensors

数据集准备

创建目录结构，将风格图片放入 train_images/，每张图片需配套一个同名 .txt 描述文件（如 img001.png 对应 img001.txt）。描述文件内容示例：

a watercolor painting of a castle, in the style of sks style

关键点：在所有描述中统一使用一个稀有标记 sks style 代表目标风格，避免与已有词汇冲突。

完整训练代码

以下代码完成数据加载、LoRA 注入与训练全流程，保存最终权重：

import torch
from diffusers import StableDiffusionPipeline, DDPMScheduler, AutoencoderKL
from diffusers.optimization import get_scheduler
from diffusers.training_utils import EMAModel
from peft import LoraConfig, get_peft_model
from torch.utils.data import Dataset
from PIL import Image
import os, glob
from torchvision import transforms

# ---------- 配置 ----------
model_id = "runwayml/stable-diffusion-v1-5"
train_dir = "./train_images"
output_dir = "./lora_weights"
epochs = 10
batch_size = 1
lr = 1e-4
lora_rank = 16
lora_alpha = 32
resolution = 512

# ---------- 数据集 ----------
class StyleDataset(Dataset):
    def __init__(self, img_dir, size=512):
        self.img_paths = sorted(glob.glob(os.path.join(img_dir, "*.png")) + 
                                glob.glob(os.path.join(img_dir, "*.jpg")))
        self.transform = transforms.Compose([
            transforms.Resize(size, interpolation=transforms.InterpolationMode.BILINEAR),
            transforms.CenterCrop(size),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5]),
        ])
        assert len(self.img_paths) > 0, "No training images found"
        
    def __len__(self):
        return len(self.img_paths)
    
    def __getitem__(self, idx):
        img_path = self.img_paths[idx]
        image = Image.open(img_path).convert("RGB")
        image = self.transform(image)
        txt_path = img_path.rsplit(".", 1)[0] + ".txt"
        with open(txt_path, "r") as f:
            prompt = f.read().strip()
        return {"pixel_values": image, "prompt": prompt}

# ---------- 模型与LoRA ----------
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
vae = pipe.vae
text_encoder = pipe.text_encoder
unet = pipe.unet
tokenizer = pipe.tokenizer
scheduler = DDPMScheduler.from_config(pipe.scheduler.config)

# 冻结所有参数
vae.requires_grad_(False)
text_encoder.requires_grad_(False)
unet.requires_grad_(False)

# 仅对 UNet 注入 LoRA
lora_config = LoraConfig(
    r=lora_rank,
    lora_alpha=lora_alpha,
    target_modules=["to_q", "to_k", "to_v", "to_out.0"],
    lora_dropout=0.1,
)
unet = get_peft_model(unet, lora_config)
unet.train()

# ---------- 优化器与数据加载 ----------
dataset = StyleDataset(train_dir, resolution)
train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)
optimizer = torch.optim.AdamW(unet.parameters(), lr=lr)
lr_scheduler = get_scheduler("cosine", optimizer=optimizer, 
                             num_warmup_steps=100, num_training_steps=len(train_loader)*epochs)

# ---------- 训练循环 ----------
device = "cuda" if torch.cuda.is_available() else "cpu"
unet.to(device)
vae.to(device)
text_encoder.to(device)

for epoch in range(epochs):
    for step, batch in enumerate(train_loader):
        pixel_values = batch["pixel_values"].to(device)
        prompts = batch["prompt"]
        
        # 编码文本
        text_inputs = tokenizer(prompts, padding="max_length", max_length=77, truncation=True, return_tensors="pt")
        text_ids = text_inputs.input_ids.to(device)
        encoder_hidden_states = text_encoder(text_ids)[0]
        
        # 编码图像并添加噪声
        latents = vae.encode(pixel_values).latent_dist.sample() * 0.18215
        noise = torch.randn_like(latents)
        timesteps = torch.randint(0, scheduler.config.num_train_timesteps, (latents.shape[0],), device=device).long()
        noisy_latents = scheduler.add_noise(latents, noise, timesteps)
        
        # 预测噪声
        noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
        loss = torch.nn.functional.mse_loss(noise_pred, noise)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        
        if step % 10 == 0:
            print(f"Epoch {epoch+1}/{epochs} | Step {step} | Loss: {loss.item():.4f}")

# ---------- 保存LoRA权重 ----------
unet.save_pretrained(output_dir)
print(f"LoRA weights saved to {output_dir}")

运行方式

# 1. 准备数据（目录结构如下）
# train_images/
#   img001.png + img001.txt
#   img002.png + img002.txt
#   ...

# 2. 执行训练
python train_lora.py

# 3. 训练完成后，output_dir 中包含 adapter_model.safetensors 等文件

验证方法

训练完成后，使用以下代码加载 LoRA 权重并生成图像验证风格迁移效果：

from diffusers import StableDiffusionPipeline
import torch

pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16
).to("cuda")
pipe.load_lora_weights("./lora_weights")

prompt = "a watercolor painting of a mountain landscape, in the style of sks style"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("output_test.png")
print("Generated image saved to output_test.png")

通过条件：生成图像应明显呈现训练集中的风格特征（如笔触、配色、构图习惯），且与未加载 LoRA 的原始模型输出有显著差异。建议用训练集中未出现过的内容（如不同物体或场景）进行测试，以确认模型学到的是风格而非特定物体。

常见问题

显存不足：将 resolution 降至 512 以下，或使用 --gradient_accumulation_steps 累积梯度；也可将 torch_dtype=torch.float16 加入训练。
过拟合：若生成图像与训练图几乎相同，降低 epochs 至 5～8，或增大 lora_dropout。
风格不显著：检查所有描述文件是否统一包含 sks style，并确保训练图片风格一致、数量不少于 15 张。

📚 课程目录