Stable Diffusion 完整训练与推理流程详解（附伪代码）

1.2 数据集要求

本文以「图像 - 文本配对数据集」为例。

二、数据预处理（核心：从原始数据到潜空间张量）

数据预处理是 SD 训练的基础，核心目标是：将原始 2K 图像缩放归一化、文本编码，最终转换为模型可直接输入的潜空间张量和文本嵌入，分为 3 个关键步骤。

2.1 基础数据集封装（图像 + 文本配对）

首先读取原始图像和文本，对图像进行缩放、归一化等基础预处理，将两者封装为 {image, text} 的配对格式，适配后续数据增强和 VAE 编码。

关键注意点：图像需缩放到 SD 标准训练尺寸（512×512），归一化到 [-1, 1]（匹配 VAE 输入要求）；文本暂不编码，仅做基础清洗。

class ImageTextDataset(Dataset):
    def __init__(self, image_dir, caption_csv, transform=None):
        """
        Args:
            image_dir: 图像文件夹路径
            caption_csv: 文本描述 csv 文件路径
            transform: 图像预处理 transform
        """
        self.image_dir = image_dir
        self.captions = pd.read_csv(caption_csv) # 读取文本描述
        self.transform = transform

    def __len__(self):
        return len(self.captions) # 数据集总样本数

    def __getitem__(self, idx):
        # 1. 读取图像
        image_name = self.captions.iloc[idx]['image_name']
        image_path = os.path.join(self.image_dir, image_name)
        image = Image.open(image_path).convert("RGB") # 转为 RGB 三通道
        
        # 2. 读取文本（基础清洗）
        text = self.captions.iloc[idx]['text'].strip()
        
        # 3. 图像预处理（缩放、归一化）
        if self.transform is not None:
            image = self.transform(image)
            
        # 返回配对数据（image: [3,512,512], text: 字符串）
        return {"image": image, "text": text}

# ------------------- 伪代码调用 -------------------
# 定义图像预处理 transform（核心：缩放 + 归一化）
image_transform = transforms.Compose([
    transforms.Resize((512, 512), interpolation=transforms.InterpolationMode.BILINEAR), # 缩放到 512×512
    transforms.ToTensor(), # 转为张量 [3,512,512]，像素值 [0,1]
    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) # 归一化到 [-1,1]
])

# 初始化基础数据集
base_dataset = ImageTextDataset(
    image_dir="dataset/images",
    caption_csv="dataset/captions.csv",
    transform=image_transform
)

# 查看数据集输出维度（BS=4 时，后续 dataloader 输出参考）
sample = base_dataset[0]
print("预处理后图像维度:", sample["image"].shape) # torch.Size([3, 512, 512])
print("文本示例:", sample["text"]) # "a red cat sitting on a chair, high resolution"

2.2 增强型潜空间数据集（AugmentedLatentDataset）

核心作用：在像素空间做数据增强（提升模型泛化性），再将增强后的图像通过 VAE 编码为潜空间张量（64×64×4）——数据增强仅在像素空间进行，潜空间不做增强（避免破坏 VAE 的压缩特征）。

常见数据增强：随机水平翻转、随机裁切、亮度/对比度调整等，增强后需保持 512×512 尺寸，再送入 VAE 编码。

class AugmentedLatentDataset(Dataset):
    def __init__(self, base_dataset, vae, augment_transform=None):
        """
        Args:
            base_dataset: 基础 ImageTextDataset
            vae: VAE 编码器（用于将像素空间转为潜空间）
            augment_transform: 像素空间的数据增强 transform
        """
        self.base_dataset = base_dataset
        self.vae = vae
        self.augment_transform = augment_transform
        # VAE 设置为评估模式（不训练 VAE，仅用于编码）
        self.vae.eval()

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

    def __getitem__(self, idx):
        # 1. 获取基础数据（预处理后的图像 + 文本）
        data = self.base_dataset[idx]
        image = data["image"] # [3,512,512]
        text = data["text"]
        
        # 2. 像素空间数据增强（可选，提升泛化性）
        if self.augment_transform is not None:
            image = self.augment_transform(image)
            
        # 3. VAE 编码：将像素空间图像转为潜空间张量（64×64×4）
        # 注意：VAE 输入需加 batch 维度，编码后去除 batch 维度，缩放潜空间（SD 标准操作）
        with torch.no_grad():
            # 编码时不计算梯度，节省显存
            latent = self.vae.encode(image.unsqueeze(0)).latent_dist.sample() # [1,4,64,64]
            latent = latent * 0.18215 # SD 固定缩放系数，匹配 VAE 训练时的归一化
            
        # 返回潜空间张量 + 文本（latent: [4,64,64], text: 字符串）
        return {"latent": latent.squeeze(0), "text": text}

# ------------------- 伪代码调用 -------------------
# 初始化 VAE（使用 SD 预训练 VAE，冻结参数）
vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
vae.requires_grad_(False) # 冻结 VAE，不参与训练

# 定义像素空间数据增强（仅在训练时使用）
augment_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转（概率 50%）
    transforms.RandomAdjustSharpness(sharpness_factor=1.5, p=0.3), # 随机调整锐度
])

# 初始化增强型潜空间数据集
latent_dataset = AugmentedLatentDataset(
    base_dataset=base_dataset,
    vae=vae,
    augment_transform=augment_transform
)

# 查看潜空间数据维度
sample = latent_dataset[0]
print("VAE 编码后潜空间维度:", sample["latent"].shape) # torch.Size([4, 64, 64])

2.3 DataLoader 封装（批量处理）

将潜空间数据集封装为 DataLoader，完成批量读取、打乱、丢弃最后不足一个 batch 的样本等操作，适配模型训练的批量输入需求，核心参数：batch_size=4（本文示例）、shuffle=True（训练时打乱数据）、drop_last=True（避免最后一个不完整 batch 影响训练）。

def create_dataloader(latent_dataset, batch_size=4, shuffle=True, drop_last=True):
    """创建 DataLoader，批量输出潜空间张量和文本"""
    dataloader = DataLoader(
        dataset=latent_dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        drop_last=drop_last,
        pin_memory=True, # 加速数据读取，适配 GPU 训练
        num_workers=4 # 多线程读取，根据 CPU 核心数调整
    )
    return dataloader

# ------------------- 伪代码调用 -------------------
# 训练集 DataLoader（shuffle=True）
train_dataloader = create_dataloader(
    latent_dataset=latent_dataset,
    batch_size=4,
    shuffle=True,
    drop_last=True
)

# 验证集 DataLoader（shuffle=False，仅用于评估）
# val_dataloader = create_dataloader(latent_dataset=val_latent_dataset, batch_size=4, shuffle=False, drop_last=True)

# 查看 DataLoader 输出维度（BS=4）
for batch in train_dataloader:
    print("Batch 潜空间维度:", batch["latent"].shape) # torch.Size([4, 4, 64, 64])
    print("Batch 文本数量:", len(batch["text"])) # 4（每个样本对应 1 条文本）
    break

三、SD 模型训练流程（核心：潜空间 DDPM 训练）

SD 的训练核心是「在潜空间内训练 DDPM」，模型输入为：加噪后的潜空间张量（noisy_latents）、时间步（timesteps）、文本嵌入（text_embeddings），目标是让 UNet 精准预测加进去的噪声，全程不涉及像素空间，仅在潜空间操作。

训练流程分为：时间步采样与加噪、文本编码、UNet 前向传播、损失计算、反向传播与参数更新，共 5 个关键环节。

3.1 初始化核心模型与优化器

SD 训练需初始化 3 个核心模型：CLIP Text Encoder（文本编码）、UNet（扩散模型核心，预测噪声）、VAE（已在数据预处理时初始化，冻结），以及优化器、学习率调度器。

关键注意点：训练时仅更新 UNet 参数，CLIP 和 VAE 预训练后冻结，大幅降低计算量和显存消耗。

def init_training_components():
    # 1. 初始化 CLIP Text Encoder 和 Tokenizer（文本编码）
    tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
    text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
    text_encoder.requires_grad_(False) # 冻结 CLIP，不参与训练
    
    # 2. 初始化 UNet（扩散模型核心，预测噪声）
    unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
    unet.train() # UNet 设为训练模式
    
    # 3. 初始化优化器（AdamW 是 SD 训练的标准优化器）
    optimizer = optim.AdamW(
        unet.parameters(),
        lr=1e-4, # 基础学习率，可根据 batchsize 调整
        betas=(0.9, 0.999),
        weight_decay=0.01
    )
    
    # 4. 初始化学习率调度器（线性衰减，适配 SD 训练）
    num_epochs = 10 # 训练总轮次
    num_training_steps = num_epochs * len(train_dataloader)
    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=num_training_steps * 0.1, # 预热步数（10%）
        num_training_steps=num_training_steps
    )
    return tokenizer, text_encoder, unet, optimizer, lr_scheduler

# ------------------- 伪代码调用 -------------------
tokenizer, text_encoder, unet, optimizer, lr_scheduler = init_training_components()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 优先使用 GPU
unet.to(device)
text_encoder.to(device)
vae.to(device)
print("核心模型初始化完成，设备:", device)

3.2 时间步采样与潜空间加噪（训练的核心前提）

DDPM 的训练核心是「加噪 - 去噪」的迭代学习，这里的加噪操作仅在潜空间进行（VAE 编码后的张量），步骤如下：

1. 为每个 batch 的样本，随机采样时间步 t（范围 1~1000，t=0 为无噪声，不采样）；
2. 生成与潜空间张量形状一致的标准正态噪声ε（训练必需的噪声信号）；
3. 根据 DDPM 前向公式，计算加噪后的潜空间张量 noisy_latents。

DDPM 前向加噪公式：

其中：

，

，是预定义的固定噪声调度序列（1e-4~0.02 线性分布）。

def add_noise_to_latents(latents, timesteps, noise_scheduler):
    """
    对潜空间张量加噪，生成 noisy_latents
    Args:
        latents: 原始潜空间张量 [BS,4,64,64]
        timesteps: 随机采样的时间步 [BS]
        noise_scheduler: DDPM 噪声调度器（预定义β序列）
    Returns:
        noisy_latents: 加噪后的潜空间张量 [BS,4,64,64]
        noise: 真实加噪的噪声 [BS,4,64,64]
    """
    # 1. 生成标准正态噪声（与潜空间张量形状一致）
    noise = torch.randn_like(latents, device=latents.device)
    
    # 2. 用噪声调度器计算加噪后的 latents（DDPM 前向公式）
    noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
    return noisy_latents, noise

# ------------------- 伪代码调用 -------------------
# 初始化 DDPM 噪声调度器（SD 标准配置：T=1000，β从 1e-4 到 0.02 线性分布）
from diffusers import DDPMScheduler
noise_scheduler = DDPMScheduler(
    num_train_timesteps=1000,
    beta_start=1e-4,
    beta_end=0.02,
    beta_schedule="linear"
)

# 从 dataloader 取一个 batch，进行加噪操作（BS=4）
for batch in train_dataloader:
    latents = batch["latent"].to(device) # [4,4,64,64]
    texts = batch["text"]
    
    # 1. 随机采样时间步 t（1~1000，每个样本的 t 不同）
    timesteps = torch.randint(1, noise_scheduler.num_train_timesteps, (latents.shape[0],), device=device)
    
    # 2. 潜空间加噪
    noisy_latents, real_noise = add_noise_to_latents(latents, timesteps, noise_scheduler)
    
    print("原始潜空间维度:", latents.shape) # [4,4,64,64]
    print("加噪后潜空间维度:", noisy_latents.shape) # [4,4,64,64]
    print("真实噪声维度:", real_noise.shape) # [4,4,64,64]
    print("随机时间步:", timesteps) # 示例：tensor([345, 890, 120, 780], device='cuda:0')
    break

3.3 文本编码（text→text_embeddings）

将 batch 中的文本字符串，通过 CLIP Tokenizer 转为 token 张量，再通过 CLIP Text Encoder 编码为文本嵌入（text_embeddings），用于后续 UNet 的 Cross-Attention 融合。

关键注意点：CLIP Tokenizer 默认将文本转为 77 维 token（不足 77 维补 0，超过 77 维截断），编码后得到 [BS, 77, 768] 的文本嵌入，需与 UNet 的注意力维度适配。

def encode_text(texts, tokenizer, text_encoder):
    """
    将文本转为 text_embeddings
    Args:
        texts: batch 文本列表（长度=BS）
        tokenizer: CLIP Tokenizer
        text_encoder: CLIP Text Encoder
    Returns:
        text_embeddings: 文本嵌入 [BS, 77, 768]
    """
    # 1. Tokenizer 编码：文本→token 张量 [BS, 77]
    inputs = tokenizer(
        texts, # 补全到 77 维
        max_length=tokenizer.model_max_length, # 77
        truncation=True, # 截断超过 77 维的文本
        return_tensors="pt" # 返回 PyTorch 张量
    ).to(text_encoder.device)
    
    # 2. Text Encoder 编码：token→文本嵌入 [BS, 77, 768]
    with torch.no_grad(): # CLIP 冻结，不计算梯度
        text_embeddings = text_encoder(**inputs).last_hidden_state
    return text_embeddings

# ------------------- 伪代码调用 -------------------
# 对当前 batch 的文本进行编码（BS=4）
text_embeddings = encode_text(texts, tokenizer, text_encoder)
print("文本嵌入维度:", text_embeddings.shape) # torch.Size([4, 77, 768])

3.4 UNet 前向传播（预测噪声）

UNet 是 SD 的核心，输入为 3 个部分：noisy_latents（加噪潜空间张量）、timesteps（时间步）、text_embeddings（文本嵌入），输出为与真实噪声形状一致的预测噪声（ε_θ）。

关键细节：

• timesteps：需先做位置编码→MLP 投影→广播，与 noisy_latents 的特征图相加，实现时间步信息的融入；
• text_embeddings：仅在 UNet 的 Cross-Attention 层融合，投影后作为 K/V，与图像特征（Q）做注意力计算，实现文本 - 图像的关联。

def unet_forward(noisy_latents, timesteps, text_embeddings, unet):
    """
    UNet 前向传播，预测噪声
    Args:
        noisy_latents: 加噪潜空间张量 [BS,4,64,64]
        timesteps: 时间步 [BS]
        text_embeddings: 文本嵌入 [BS,77,768]
        unet: UNet 模型
    Returns:
        noise_pred: 预测噪声 [BS,4,64,64]
    """
    # UNet 直接接收三个输入，内部自动完成 timesteps 和 text_embeddings 的维度适配
    # 1. timesteps：内部做位置编码→投影→广播，与 noisy_latents 特征相加
    # 2. text_embeddings：内部投影后，在 Cross-Attention 层作为 K/V 融合
    noise_pred = unet(
        sample=noisy_latents,
        timestep=timesteps,
        encoder_hidden_states=text_embeddings
    ).sample # sample 是 UNet 输出的预测噪声
    return noise_pred

# ------------------- 伪代码调用 -------------------
# UNet 前向传播，预测噪声
noise_pred = unet_forward(noisy_latents, timesteps, text_embeddings, unet)
print("预测噪声维度:", noise_pred.shape) # torch.Size([4,4,64,64])（与真实噪声维度一致）

3.5 损失计算与反向传播

SD 训练的核心损失是「预测噪声与真实噪声的 MSE Loss」——无需反推潜空间张量，直接对比 UNet 输出的 noise_pred 和加噪时的 real_noise，计算均方误差，再反向传播更新 UNet 参数。

关键注意点：Loss 仅计算噪声的差异，这是 DDPM 的核心简化设计，让模型专注于'猜中加进去的噪声'，后续推理时通过采样器反向去噪即可生成图像。

def train_one_batch(noisy_latents, timesteps, text_embeddings, real_noise, unet, optimizer, lr_scheduler):
    """训练一个 batch，完成前向、损失计算、反向传播、参数更新"""
    # 1. 前向传播，预测噪声
    noise_pred = unet_forward(noisy_latents, timesteps, text_embeddings, unet)
    
    # 2. 计算 MSE Loss（预测噪声 vs 真实噪声）
    loss_fn = nn.MSELoss()
    loss = loss_fn(noise_pred, real_noise)
    
    # 3. 反向传播（仅更新 UNet 参数）
    optimizer.zero_grad() # 清空梯度
    loss.backward() # 计算梯度
    optimizer.step() # 更新参数
    lr_scheduler.step() # 学习率调度
    return loss.item()

# ------------------- 伪代码调用 -------------------
# 训练一个 batch，查看 Loss
loss = train_one_batch(noisy_latents, timesteps, text_embeddings, real_noise, unet, optimizer, lr_scheduler)
print("当前 batch 的 Loss:", loss) # 示例：0.035（训练初期 Loss 较高，后期逐步下降）

3.6 完整训练循环（多 Epoch 迭代）

将上述环节整合，实现多 Epoch 的完整训练，定期保存模型权重（checkpoint），用于后续推理生成。

def full_training_loop(num_epochs, train_dataloader, noise_scheduler, tokenizer, text_encoder, unet, optimizer, lr_scheduler):
    """完整训练循环"""
    unet.train()
    for epoch in range(num_epochs):
        epoch_loss = 0.0
        for step, batch in enumerate(train_dataloader):
            # 1. 读取 batch 数据
            latents = batch["latent"].to(device)
            texts = batch["text"]
            
            # 2. 时间步采样与潜空间加噪
            timesteps = torch.randint(1, noise_scheduler.num_train_timesteps, (latents.shape[0],), device=device)
            noisy_latents, real_noise = add_noise_to_latents(latents, timesteps, noise_scheduler)
            
            # 3. 文本编码
            text_embeddings = encode_text(texts, tokenizer, text_encoder)
            
            # 4. 训练一个 batch，计算 Loss
            batch_loss = train_one_batch(noisy_latents, timesteps, text_embeddings, real_noise, unet, optimizer, lr_scheduler)
            epoch_loss += batch_loss
            
            # 打印日志（每 100 步打印一次）
            if (step + 1) % 100 == 0:
                print(f"Epoch [{epoch+1}/{num_epochs}], Step [{step+1}/{len(train_dataloader)}], Batch Loss: {batch_loss:.4f}")
        
        # 计算当前 Epoch 的平均 Loss
        avg_epoch_loss = epoch_loss / len(train_dataloader)
        print(f"Epoch [{epoch+1}/{num_epochs}] Finished, Average Loss: {avg_epoch_loss:.4f}")
        
        # 定期保存模型权重（每 1 个 Epoch 保存一次）
        torch.save(unet.state_dict(), f"unet_epoch_{epoch+1}.pth")
        print(f"Model saved to unet_epoch_{epoch+1}.pth")

# ------------------- 伪代码调用 -------------------
# 启动完整训练（10 个 Epoch）
num_epochs = 10
full_training_loop(
    num_epochs=num_epochs,
    train_dataloader=train_dataloader,
    noise_scheduler=noise_scheduler,
    tokenizer=tokenizer,
    text_encoder=text_encoder,
    unet=unet,
    optimizer=optimizer,
    lr_scheduler=lr_scheduler
)

四、SD 推理流程（核心：潜空间逐步去噪）

推理阶段的核心是「反向去噪」：从纯高斯噪声（t=1000）开始，按设定的步数逐步去噪，最终得到清晰的潜空间张量，再通过 VAE 解码为像素空间图像。

关键环节：逐步去噪（每步一次 UNet 前向）、CFG 增强（强化文本控制）、随机噪声（增加生成多样性）。

4.1 推理前准备（加载模型与参数）

def init_inference_components(unet_ckpt_path):
    """初始化推理所需组件，加载训练好的 UNet 权重"""
    # 1. 初始化 VAE（用于最终解码潜空间→像素空间）
    vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
    vae.eval()
    vae.requires_grad_(False)
    
    # 2. 初始化 CLIP（文本编码）
    tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
    text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
    text_encoder.eval()
    text_encoder.requires_grad_(False)
    
    # 3. 初始化 UNet，加载训练好的权重
    unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
    unet.load_state_dict(torch.load(unet_ckpt_path)) # 加载训练权重
    unet.eval()
    unet.requires_grad_(False)
    
    # 4. 初始化推理用噪声调度器（与训练时一致）
    noise_scheduler = DDPMScheduler(
        num_train_timesteps=1000,
        beta_start=1e-4,
        beta_end=0.02,
        beta_schedule="linear"
    )
    
    # 5. 初始化采样器（这里用 DDIM 采样器，加速推理，步数 20~50 步）
    from diffusers import DDIMScheduler
    sampler = DDIMScheduler.from_config(noise_scheduler.config)
    sampler.set_timesteps(num_inference_steps=50) # 推理步数（50 步，比训练时 1000 步快 20 倍）
    return vae, tokenizer, text_encoder, unet, sampler

# ------------------- 伪代码调用 -------------------
# 加载训练好的 UNet 权重（示例：第 10 个 Epoch 的权重）
unet_ckpt_path = "unet_epoch_10.pth"
vae, tokenizer, text_encoder, unet, sampler = init_inference_components(unet_ckpt_path)

# 移动到 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vae.to(device)
text_encoder.to(device)
unet.to(device)

4.2 文本编码（推理时与训练一致）

推理时的文本编码流程与训练完全一致，将输入的文本描述转为 text_embeddings，同时生成'空文本嵌入'（用于 CFG 增强）。

def encode_text_inference(prompt, tokenizer, text_encoder):
    """推理时的文本编码，同时生成有文本和无文本（空文本）的嵌入"""
    # 1. 有文本的嵌入（prompt 为输入文本）
    prompt_inputs = tokenizer(
        prompt, # 补全到 77 维
        max_length=tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt"
    ).to(text_encoder.device)
    with torch.no_grad():
        text_embeddings = text_encoder(**prompt_inputs).last_hidden_state
    
    # 2. 无文本的嵌入（空文本，用于 CFG 增强）
    null_inputs = tokenizer(
        null_prompt, # 补全到 77 维
        max_length=tokenizer.model_max_length,
        truncation=True,
        return_tensors="pt"
    ).to(text_encoder.device)
    with torch.no_grad():
        null_text_embeddings = text_encoder(**null_inputs).last_hidden_state
    return text_embeddings, null_text_embeddings

# ------------------- 伪代码调用 -------------------
# 输入推理文本（示例："a red cat sitting on a chair, high resolution"）
prompt = "a red cat sitting on a chair, high resolution"
text_embeddings, null_text_embeddings = encode_text_inference(prompt, tokenizer, text_encoder)
print("推理文本嵌入维度:", text_embeddings.shape) # [1,77,768]（推理时 BS=1，单张生成）

4.3 逐步去噪与 CFG 增强（推理核心）

推理时的去噪流程：从 t=1000 的纯高斯噪声开始，按采样器设定的步数（50 步）逐步从 t=1000→1 去噪，每步执行一次 UNet 前向传播，通过 CFG 增强文本控制，加入随机噪声增加多样性。

CFG 核心公式（强化文本引导）：

$$\epsilon_{cfg} = \epsilon_{null} + cfg_scale \times (\epsilon_{text} - \epsilon_{null})$$

其中：cfg_scale 默认 7.5，值越大，文本控制越强（过高会导致图像失真）。

def inference(prompt, vae, tokenizer, text_encoder, unet, sampler, cfg_scale=7.5):
    """
    SD 推理生成图像
    Args:
        prompt: 文本描述
        vae: VAE 解码器
        tokenizer: CLIP Tokenizer
        text_encoder: CLIP Text Encoder
        unet: 训练好的 UNet
        sampler: 采样器（DDIM）
        cfg_scale: CFG 系数，控制文本引导强度
    Returns:
        generated_image: 生成的像素空间图像 [3,512,512]
    """
    # 1. 文本编码，得到有文本/无文本嵌入
    text_embeddings, null_text_embeddings = encode_text_inference(prompt, tokenizer, text_encoder)
    # 拼接有文本和无文本嵌入（适配 CFG 计算）
    text_embeddings = torch.cat([null_text_embeddings, text_embeddings]) # [2,77,768]
    
    # 2. 初始化潜空间噪声（t=1000，纯高斯噪声）
    batch_size = 1
    latent_dim = 4
    latent_size = 64
    noise = torch.randn(
        (batch_size, latent_dim, latent_size, latent_size), device=unet.device
    )
    latents = noise # 初始潜空间噪声（t=1000）
    
    # 3. 逐步去噪（按采样器的时间步迭代）
    with torch.no_grad(): # 推理时不计算梯度
        for t in sampler.timesteps:
            # 3.1 扩展 latents 和 timesteps，适配 CFG 的双输入（有文本/无文本）
            latent_model_input = torch.cat([latents] * 2) # [2,4,64,64]
            timestep = torch.tensor([t] * batch_size * 2, device=unet.device)
            
            # 3.2 UNet 前向传播，预测噪声（一次预测有文本/无文本两种情况）
            noise_pred = unet(
                sample=latent_model_input,
                timestep=timestep,
                encoder_hidden_states=text_embeddings
            ).sample # [2,4,64,64]
            
            # 3.3 CFG 增强：分离无文本/有文本的噪声预测，计算最终噪声
            noise_pred_null, noise_pred_text = noise_pred.chunk(2) # 各 [1,4,64,64]
            noise_pred = noise_pred_null + cfg_scale * (noise_pred_text - noise_pred_null)
            
            # 3.4 采样器去噪，得到 t-1 的潜空间张量
            latents = sampler.step(noise_pred, t, latents).prev_sample
            
    # 4. VAE 解码：潜空间→像素空间（512×512）
    latents = latents / 0.18215 # 反缩放（与训练时的缩放对应）
    with torch.no_grad():
        generated_image = vae.decode(latents).sample # [1,3,512,512]
    
    # 5. 图像后处理：从 [-1,1] 转回 [0,255]，转为 PIL 图像
    generated_image = (generated_image / 2 + 0.5).clamp(0, 1) # 归一化到 [0,1]
    generated_image = generated_image.cpu().permute(0, 2, 3, 1).numpy()[0] # [512,512,3]
    generated_image = (generated_image * 255).astype(np.uint8)
    generated_image = Image.fromarray(generated_image)
    return generated_image

# ------------------- 伪代码调用 -------------------
# 执行推理，生成图像
generated_image = inference(
    prompt="a red cat sitting on a chair, high resolution",
    vae=vae,
    tokenizer=tokenizer,
    text_encoder=text_encoder,
    unet=unet,
    sampler=sampler,
    cfg_scale=7.5
)
# 保存生成的图像
generated_image.save("generated_image.jpg")
print("图像生成完成，已保存为 generated_image.jpg")

五、LoRA 轻量化训练（可选，核心：冻结主模型，训练适配器）

SD 的 UNet 参数量数十亿，全量训练显存消耗大（需 40GB 以上），LoRA（Low-Rank Adaptation）通过在 UNet 的注意力层挂载轻量化适配器，仅训练适配器参数（参数量仅百万级），大幅降低显存消耗，同时实现特定风格/内容的微调。

LoRA 训练流程分为两种方式：手动实现适配器、用 PEFT 库简化实现（推荐新手）。

5.1 PEFT 库简化实现 LoRA 训练（推荐）

PEFT 库已封装好 LoRA 逻辑，只需定义 LoRA 配置，挂载到 UNet，即可实现轻量化训练，无需手动编写适配器。

def init_lora_training():
    """初始化 LoRA 训练组件，冻结主模型，挂载 LoRA 适配器"""
    # 1. 初始化基础模型（与之前一致）
    tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
    text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
    vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
    unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
    
    # 2. 冻结主模型（仅训练 LoRA 适配器）
    text_encoder.requires_grad_(False)
    vae.requires_grad_(False)
    unet.requires_grad_(False)
    
    # 3. 定义 LoRA 配置（核心参数）
    lora_config = LoraConfig(
        r=8, # LoRA 秩，越小参数量越少，一般取 4~16
        lora_alpha=16, # 缩放系数，通常是 r 的 2 倍
        target_modules=["q_proj", "v_proj"], # 挂载到 UNet 的注意力层（Q/V 投影层）
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM"
    )
    
    # 4. 挂载 LoRA 适配器到 UNet
    unet = get_peft_model(unet, lora_config)
    unet.print_trainable_parameters() # 查看可训练参数（通常仅百万级）
    
    # 5. 初始化优化器和调度器（仅优化 LoRA 参数）
    optimizer = optim.AdamW(
        unet.parameters(),
        lr=5e-5, # LoRA 学习率可略低
        betas=(0.9, 0.999),
        weight_decay=0.01
    )
    num_epochs = 5
    num_training_steps = num_epochs * len(train_dataloader)
    lr_scheduler = get_scheduler(
        name="linear",
        optimizer=optimizer,
        num_warmup_steps=num_training_steps * 0.1,
        num_training_steps=num_training_steps
    )
    return tokenizer, text_encoder, vae, unet, optimizer, lr_scheduler

# ------------------- 伪代码调用 -------------------
# 初始化 LoRA 训练组件
tokenizer_lora, text_encoder_lora, vae_lora, unet_lora, optimizer_lora, lr_scheduler_lora = init_lora_training()

# 启动 LoRA 训练（训练流程与全量训练一致，仅训练 LoRA 参数）
full_training_loop(
    num_epochs=5,
    train_dataloader=train_dataloader,
    noise_scheduler=noise_scheduler,
    tokenizer=tokenizer_lora,
    text_encoder=text_encoder_lora,
    unet=unet_lora,
    optimizer=optimizer_lora,
    lr_scheduler=lr_scheduler_lora
)

# 保存 LoRA 权重（仅保存适配器参数，文件体积小，约几 MB）
unet_lora.save_pretrained("lora_weights")
print("LoRA 权重保存完成，路径：lora_weights")

5.2 LoRA 推理（挂载适配器）

LoRA 推理时，需加载预训练的 UNet 主模型，再挂载 LoRA 适配器，即可实现微调后的生成效果，无需加载完整的微调 UNet 权重。

def lora_inference(prompt, lora_path, cfg_scale=7.5):
    """LoRA 推理，挂载适配器"""
    # 1. 初始化基础模型（与推理时一致）
    vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
    tokenizer = CLIPTokenizer.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="tokenizer")
    text_encoder = CLIPTextModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="text_encoder")
    unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
    
    # 2. 挂载 LoRA 适配器
    unet = PeftModel.from_pretrained(unet, lora_path)
    unet.eval()
    unet.requires_grad_(False)
    
    # 3. 初始化采样器
    sampler = DDIMScheduler.from_config(DDPMScheduler(num_train_timesteps=1000).config)
    sampler.set_timesteps(num_inference_steps=50)
    
    # 4. 执行推理（与普通推理流程一致）
    generated_image = inference(
        prompt=prompt,
        vae=vae,
        tokenizer=tokenizer,
        text_encoder=text_encoder,
        unet=unet,
        sampler=sampler,
        cfg_scale=cfg_scale
    )
    return generated_image

# ------------------- 伪代码调用 -------------------
# LoRA 推理（示例：加载训练好的 LoRA 权重）
lora_path = "lora_weights"
lora_generated_image = lora_inference(
    prompt="a red cat sitting on a chair, high resolution, lora style",
    lora_path=lora_path,
    cfg_scale=7.5
)
# 保存 LoRA 生成的图像
lora_generated_image.save("lora_generated_image.jpg")
print("LoRA 图像生成完成，已保存")

六、常见问题与注意事项

• 显存不足：可降低分辨率、使用 LoRA、开启梯度检查点（gradient checkpointing）；
• 训练 Loss 不下降：检查时间步采样范围（需 1~1000）、噪声调度器配置、学习率是否过高；
• 生成图像不贴合文本：调大 CFG_scale（7~10）、增加推理步数（50 步）、检查文本编码是否正确；
• LoRA 训练无效：确认 target_modules 是否为 UNet 的注意力层（q_proj、v_proj）、LoRA 秩 r 是否合理。

七、总结

Stable Diffusion 的核心是「潜空间扩散」，全程围绕 VAE（潜空间映射）、CLIP（文本编码）、UNet（噪声预测）三大模型展开，训练时在潜空间加噪、让 UNet 预测噪声，推理时逐步去噪、用 CFG 强化文本控制，LoRA 则实现轻量化微调。