4. 文本到图像生成模型微调

任务说明

本案例基于预训练的 Stable Diffusion 模型，在特定风格的图像数据集上进行微调，使模型能够根据文本提示生成对应风格的图像。核心目标是让模型学会目标风格（例如“水彩插画风格”），并保持对文本内容的语义跟随能力。

环境与依赖

需要 Python 3.10+、PyTorch 2.0+、CUDA 11.8+，以及以下核心库：

pip install diffusers[training]==0.27.2 accelerate==0.27.2 peft==0.9.0 transformers==4.36.2 datasets==2.16.1 pillow

数据集准备

使用 30 张水彩插画风格图像（建议分辨率 512×512），每张图像配一个简短文本描述（如“a watercolor painting of a mountain landscape”）。将图像与描述组织为如下格式的 JSON 文件 data.json：

[
  {"image": "images/001.jpg", "caption": "a watercolor painting of a red fox in the forest"},
  {"image": "images/002.jpg", "caption": "a watercolor painting of a sailing boat on the sea"}
]

完整微调代码

以下代码使用 LoRA 对 Stable Diffusion 的 UNet 进行参数高效微调，训练后保存 LoRA 权重：

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import json, os
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 transformers import CLIPTextModel, CLIPTokenizer
from tqdm import tqdm

# ---------- 配置 ----------
model_id = "runwayml/stable-diffusion-v1-5"
data_file = "data.json"
output_dir = "./lora_weights"
batch_size = 2
learning_rate = 1e-4
num_epochs = 10
resolution = 512
lora_rank = 8

# ---------- 数据集 ----------
class StyleDataset(Dataset):
    def __init__(self, json_file):
        with open(json_file, 'r') as f:
            self.samples = json.load(f)
        self.transform = transforms.Compose([
            transforms.Resize((resolution, resolution)),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])
        ])

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        sample = self.samples[idx]
        image = Image.open(sample["image"]).convert("RGB")
        image = self.transform(image)
        return {"pixel_values": image, "caption": sample["caption"]}

# ---------- 加载预训练组件 ----------
tokenizer = CLIPTokenizer.from_pretrained(model_id, subfolder="tokenizer")
text_encoder = CLIPTextModel.from_pretrained(model_id, subfolder="text_encoder")
vae = AutoencoderKL.from_pretrained(model_id, subfolder="vae")
unet = torch.compile(__import__('diffusers').UNet2DConditionModel.from_pretrained(
    model_id, subfolder="unet", torch_dtype=torch.float32))

# 冻结 VAE 和文本编码器
vae.requires_grad_(False)
text_encoder.requires_grad_(False)

# ---------- 配置 LoRA ----------
lora_config = LoraConfig(
    r=lora_rank,
    lora_alpha=lora_rank,
    target_modules=["to_q", "to_k", "to_v", "to_out.0"],
    lora_dropout=0.1
)
unet = get_peft_model(unet, lora_config)
unet.train()

# ---------- 优化器与调度器 ----------
optimizer = torch.optim.AdamW(unet.parameters(), lr=learning_rate)
dataset = StyleDataset(data_file)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
scheduler = DDPMScheduler.from_pretrained(model_id, subfolder="scheduler")
lr_scheduler = get_scheduler(
    "cosine", optimizer=optimizer,
    num_warmup_steps=100, num_training_steps=len(dataloader) * num_epochs
)

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

for epoch in range(num_epochs):
    for step, batch in enumerate(tqdm(dataloader)):
        pixel_values = batch["pixel_values"].to(device)
        captions = batch["caption"]

        # 编码文本
        text_inputs = tokenizer(captions, padding="max_length", max_length=77,
                                truncation=True, return_tensors="pt").to(device)
        encoder_hidden_states = text_encoder(text_inputs.input_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()

    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item():.4f}")

# ---------- 保存 LoRA 权重 ----------
os.makedirs(output_dir, exist_ok=True)
unet.save_pretrained(output_dir)
print(f"LoRA weights saved to {output_dir}")

推理生成图像

训练完成后，使用以下代码加载 LoRA 权重并生成图像：

import torch
from diffusers import StableDiffusionPipeline
from peft import PeftModel

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

# 加载 LoRA 权重
pipe.unet = PeftModel.from_pretrained(pipe.unet, "./lora_weights")

prompt = "a watercolor painting of a castle on a hill at sunset"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.5).images[0]
image.save("output.png")
print("Generated image saved to output.png")

验证方法与通过标准

训练收敛验证：训练过程中 loss 应逐步下降，最终稳定在 0.1 以下（预期值，非实测）。
风格一致性验证：用训练集之外的文本提示（如“a watercolor painting of a cat sleeping on a windowsill”）生成图像，观察输出是否呈现水彩风格（笔触、色彩晕染、纸张纹理）。
文本跟随验证：对比不同文本提示（如“a watercolor painting of a dog” vs “a watercolor painting of a boat”）生成的图像，主体内容应随文本变化。
定量对比（可选）：计算生成图像与训练集图像的 CLIP 特征余弦相似度，若风格迁移成功，相似度应高于未微调模型的结果。

常见问题

显存不足：将 batch_size 降为 1，或使用 torch.cuda.amp 混合精度训练。
风格过拟合：若生成图像内容与训练集过于雷同，可增加训练图像数量或降低 LoRA rank。
文本跟随丢失：若模型只学风格而忽略文本，可适当增加学习率或训练轮数，并确保训练描述文本多样化。

📚 课程目录