跳到主要内容Stable Diffusion 完整训练与推理流程详解(附伪代码) | 极客日志PythonAI算法
Stable Diffusion 完整训练与推理流程详解(附伪代码)
综述由AI生成Stable Diffusion 基于潜空间扩散模型,通过 VAE、CLIP 和 UNet 三大核心组件实现文本到图像的生成。详细拆解了从数据预处理、模型训练、推理生成到 LoRA 轻量化微调的完整工程流程,包含关键伪代码与实操细节。重点阐述了潜空间加噪去噪机制、Cross-Attention 融合原理及 CFG 增强策略,旨在帮助开发者理解 SD 底层逻辑并落地实践。
技术博主12 浏览 Stable Diffusion(SD)的核心理论基石源自论文《High-Resolution Image Synthesis with Latent Diffusion Models》(LDM)。其革命性创新在于将扩散模型从高维像素空间迁移至 VAE 预训练的低维潜空间,在大幅降低训练与推理的计算成本的同时,通过跨注意力机制实现文本、布局等多模态条件控制,兼顾了生成质量与灵活性。本文将基于这一核心思想,从数据预处理、模型训练、推理生成到 LoRA 轻量化训练,一步步拆解 SD 的完整技术流程,每个关键环节均搭配伪代码,结合实操场景,帮助理解 SD 的工程实现。

论文地址:https://arxiv.org/pdf/2112.10752
论文代码:https://github.com/CompVis/latent-diffusion
复现代码(基于非官方的复现,简化版):https://github.com/wenwenqqq/sd-demo
核心前提:SD 的核心设计是「潜空间扩散」——用 VAE 将图片映射到低维潜空间,在潜空间内完成 DDPM 的训练与推理,大幅降低计算量和显存消耗,这也是 SD 能高效训练大尺寸图片的关键。
一、前期准备与核心依赖
在开始流程前,需准备好核心依赖库和数据集。这里列出实操所需的基础依赖(基于 PyTorch 框架),以及数据集的基础要求。(以下伪代码仅供参考)
1.1 核心依赖库
SD 的训练/推理依赖 VAE、CLIP、UNet 三大核心模型,以及数据处理、扩散模型相关的工具库,伪代码如下:
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms
from diffusers import AutoencoderKL, CLIPTextModel, CLIPTokenizer, UNet2DConditionModel
from diffusers.optimization import get_scheduler
from diffusers.utils import logging
from peft import LoraConfig, get_peft_model, PeftModel
logging.set_verbosity_info()
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):
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")
text = self.captions.iloc[idx]['text'].strip()
if self.transform is not None:
image = self.transform(image)
return {"image": image, "text": text}
image_transform = transforms.Compose([
transforms.Resize((512, 512), interpolation=transforms.InterpolationMode.BILINEAR),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])
base_dataset = ImageTextDataset(
image_dir="dataset/images",
caption_csv="dataset/captions.csv",
transform=image_transform
)
sample = base_dataset[0]
print("预处理后图像维度:", sample["image"].shape)
print("文本示例:", sample["text"])
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
self.vae.eval()
def __len__(self):
return len(self.base_dataset)
def __getitem__(self, idx):
data = self.base_dataset[idx]
image = data["image"]
text = data["text"]
if self.augment_transform is not None:
image = self.augment_transform(image)
with torch.no_grad():
latent = self.vae.encode(image.unsqueeze(0)).latent_dist.sample()
latent = latent * 0.18215
return {"latent": latent.squeeze(0), "text": text}
vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
vae.requires_grad_(False)
augment_transform = transforms.Compose([
transforms.RandomHorizontalFlip(p=0.5),
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)
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,
num_workers=4
)
return dataloader
train_dataloader = create_dataloader(
latent_dataset=latent_dataset,
batch_size=4,
shuffle=True,
drop_last=True
)
for batch in train_dataloader:
print("Batch 潜空间维度:", batch["latent"].shape)
print("Batch 文本数量:", len(batch["text"]))
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():
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)
unet = UNet2DConditionModel.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="unet")
unet.train()
optimizer = optim.AdamW(
unet.parameters(),
lr=1e-4,
betas=(0.9, 0.999),
weight_decay=0.01
)
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,
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")
unet.to(device)
text_encoder.to(device)
vae.to(device)
print("核心模型初始化完成,设备:", device)
3.2 时间步采样与潜空间加噪(训练的核心前提)
DDPM 的训练核心是「加噪 - 去噪」的迭代学习,这里的加噪操作仅在潜空间进行(VAE 编码后的张量),步骤如下:
- 为每个 batch 的样本,随机采样时间步 t(范围 1~1000,t=0 为无噪声,不采样);
- 生成与潜空间张量形状一致的标准正态噪声ε(训练必需的噪声信号);
- 根据 DDPM 前向公式,计算加噪后的潜空间张量 noisy_latents。
,是预定义的固定噪声调度序列(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]
"""
noise = torch.randn_like(latents, device=latents.device)
noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
return noisy_latents, noise
from diffusers import DDPMScheduler
noise_scheduler = DDPMScheduler(
num_train_timesteps=1000,
beta_start=1e-4,
beta_end=0.02,
beta_schedule="linear"
)
for batch in train_dataloader:
latents = batch["latent"].to(device)
texts = batch["text"]
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)
print("原始潜空间维度:", latents.shape)
print("加噪后潜空间维度:", noisy_latents.shape)
print("真实噪声维度:", real_noise.shape)
print("随机时间步:", timesteps)
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]
"""
inputs = tokenizer(
texts,
max_length=tokenizer.model_max_length,
truncation=True,
return_tensors="pt"
).to(text_encoder.device)
with torch.no_grad():
text_embeddings = text_encoder(**inputs).last_hidden_state
return text_embeddings
text_embeddings = encode_text(texts, tokenizer, text_encoder)
print("文本嵌入维度:", text_embeddings.shape)
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]
"""
noise_pred = unet(
sample=noisy_latents,
timestep=timesteps,
encoder_hidden_states=text_embeddings
).sample
return noise_pred
noise_pred = unet_forward(noisy_latents, timesteps, text_embeddings, unet)
print("预测噪声维度:", noise_pred.shape)
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,完成前向、损失计算、反向传播、参数更新"""
noise_pred = unet_forward(noisy_latents, timesteps, text_embeddings, unet)
loss_fn = nn.MSELoss()
loss = loss_fn(noise_pred, real_noise)
optimizer.zero_grad()
loss.backward()
optimizer.step()
lr_scheduler.step()
return loss.item()
loss = train_one_batch(noisy_latents, timesteps, text_embeddings, real_noise, unet, optimizer, lr_scheduler)
print("当前 batch 的 Loss:", 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):
latents = batch["latent"].to(device)
texts = batch["text"]
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)
text_embeddings = encode_text(texts, tokenizer, text_encoder)
batch_loss = train_one_batch(noisy_latents, timesteps, text_embeddings, real_noise, unet, optimizer, lr_scheduler)
epoch_loss += batch_loss
if (step + 1) % 100 == 0:
print(f"Epoch [{epoch+1}/{num_epochs}], Step [{step+1}/{len(train_dataloader)}], Batch Loss: {batch_loss:.4f}")
avg_epoch_loss = epoch_loss / len(train_dataloader)
print(f"Epoch [{epoch+1}/{num_epochs}] Finished, Average Loss: {avg_epoch_loss:.4f}")
torch.save(unet.state_dict(), f"unet_epoch_{epoch+1}.pth")
print(f"Model saved to unet_epoch_{epoch+1}.pth")
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 权重"""
vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae")
vae.eval()
vae.requires_grad_(False)
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)
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)
noise_scheduler = DDPMScheduler(
num_train_timesteps=1000,
beta_start=1e-4,
beta_end=0.02,
beta_schedule="linear"
)
from diffusers import DDIMScheduler
sampler = DDIMScheduler.from_config(noise_scheduler.config)
sampler.set_timesteps(num_inference_steps=50)
return vae, tokenizer, text_encoder, unet, sampler
unet_ckpt_path = "unet_epoch_10.pth"
vae, tokenizer, text_encoder, unet, sampler = init_inference_components(unet_ckpt_path)
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):
"""推理时的文本编码,同时生成有文本和无文本(空文本)的嵌入"""
prompt_inputs = tokenizer(
prompt,
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
null_inputs = tokenizer(
null_prompt,
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
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)
4.3 逐步去噪与 CFG 增强(推理核心)
推理时的去噪流程:从 t=1000 的纯高斯噪声开始,按采样器设定的步数(50 步)逐步从 t=1000→1 去噪,每步执行一次 UNet 前向传播,通过 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]
"""
text_embeddings, null_text_embeddings = encode_text_inference(prompt, tokenizer, text_encoder)
text_embeddings = torch.cat([null_text_embeddings, text_embeddings])
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
with torch.no_grad():
for t in sampler.timesteps:
latent_model_input = torch.cat([latents] * 2)
timestep = torch.tensor([t] * batch_size * 2, device=unet.device)
noise_pred = unet(
sample=latent_model_input,
timestep=timestep,
encoder_hidden_states=text_embeddings
).sample
noise_pred_null, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_null + cfg_scale * (noise_pred_text - noise_pred_null)
latents = sampler.step(noise_pred, t, latents).prev_sample
latents = latents / 0.18215
with torch.no_grad():
generated_image = vae.decode(latents).sample
generated_image = (generated_image / 2 + 0.5).clamp(0, 1)
generated_image = generated_image.cpu().permute(0, 2, 3, 1).numpy()[0]
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 适配器"""
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")
text_encoder.requires_grad_(False)
vae.requires_grad_(False)
unet.requires_grad_(False)
lora_config = LoraConfig(
r=8,
lora_alpha=16,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
unet = get_peft_model(unet, lora_config)
unet.print_trainable_parameters()
optimizer = optim.AdamW(
unet.parameters(),
lr=5e-5,
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
tokenizer_lora, text_encoder_lora, vae_lora, unet_lora, optimizer_lora, lr_scheduler_lora = init_lora_training()
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
)
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 推理,挂载适配器"""
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")
unet = PeftModel.from_pretrained(unet, lora_path)
unet.eval()
unet.requires_grad_(False)
sampler = DDIMScheduler.from_config(DDPMScheduler(num_train_timesteps=1000).config)
sampler.set_timesteps(num_inference_steps=50)
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_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_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 则实现轻量化微调。
相关免费在线工具
- 加密/解密文本
使用加密算法(如AES、TripleDES、Rabbit或RC4)加密和解密文本明文。 在线工具,加密/解密文本在线工具,online
- RSA密钥对生成器
生成新的随机RSA私钥和公钥pem证书。 在线工具,RSA密钥对生成器在线工具,online
- Mermaid 预览与可视化编辑
基于 Mermaid.js 实时预览流程图、时序图等图表,支持源码编辑与即时渲染。 在线工具,Mermaid 预览与可视化编辑在线工具,online
- 随机西班牙地址生成器
随机生成西班牙地址(支持马德里、加泰罗尼亚、安达卢西亚、瓦伦西亚筛选),支持数量快捷选择、显示全部与下载。 在线工具,随机西班牙地址生成器在线工具,online
- Gemini 图片去水印
基于开源反向 Alpha 混合算法去除 Gemini/Nano Banana 图片水印,支持批量处理与下载。 在线工具,Gemini 图片去水印在线工具,online
- curl 转代码
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具,curl 转代码在线工具,online