Stable Diffusion（二）：从原理到实战

# Stable Diffusion 三大组件概览
from diffusers import StableDiffusionPipeline
import torch

# 加载预训练模型（自动下载 VAE + U-Net + Text Encoder）
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
    safety_checker=None
)
pipe = pipe.to("cuda")

# 一键文生图
image = pipe(
    prompt="a futuristic city at sunset, cyberpunk style",
    negative_prompt="blurry, low quality",
    num_inference_steps=50,
    guidance_scale=7.5,
).images[0]
image.save("output.png")

# 模型参数量与显存需求对比
models = {
    "DDPM (256x256)":  {"params": "350M",  "vram": "~8GB",  "space": "像素空间"},
    "DALL-E 2":        {"params": "3.5B",  "vram": "~40GB", "space": "CLIP 空间"},
    "SD 1.5 (512x512)": {"params": "860M",  "vram": "~4GB",  "space": "潜在空间"},
    "SDXL (1024x1024)": {"params": "2.6B",  "vram": "~8GB",  "space": "潜在空间"},
}

for name, spec in models.items():
    print(f"{name:20s} | {spec['params']:>6s} | {spec['vram']:>7s} | {spec['space']}")

# VAE 编码与解码过程
import torch
import torch.nn as nn
from diffusers import AutoencoderKL

# 加载预训练 VAE
vae = AutoencoderKL.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    subfolder="vae",
    torch_dtype=torch.float16
)
vae = vae.to("cuda")

# 编码：像素空间 -> 潜在空间
# 输入: (1, 3, 512, 512) -> 输出: (1, 4, 64, 64)
with torch.no_grad():
    latent = vae.encode(pixels).latent_dist.sample()
    latent = latent * vae.config.scaling_factor  # 缩放因子 0.18215

# 解码：潜在空间 -> 像素空间
with torch.no_grad():
    latent = latent / vae.config.scaling_factor
    decoded = vae.decode(latent).sample
    # decoded shape: (1, 3, 512, 512)

# VAE 重构质量评估
import torch
from torch.nn.functional import mse_loss
from lpips import LPIPS

lpips_model = LPIPS(net="vgg").to("cuda")

def evaluate_vae(vae, images):
    """评估 VAE 的重构质量"""
    with torch.no_grad():
        # 编码
        posterior = vae.encode(images).latent_dist
        latents = posterior.sample() * vae.config.scaling_factor
        
        # 解码
        latents = latents / vae.config.scaling_factor
        reconstructed = vae.decode(latents).sample
    
    # 计算重构误差
    mse = mse_loss(images, reconstructed).item()
    lpips = lpips_model(images, reconstructed).mean().item()
    
    print(f"MSE: {mse:.6f}")
    print(f"LPIPS: {lpips:.4f} (越低越好, <0.05 为优秀)")
    return reconstructed

# CLIP 文本编码器解析
from transformers import CLIPTokenizer, CLIPTextModel
import torch

tokenizer = CLIPTokenizer.from_pretrained(
    "openai/clip-vit-large-patch14")
text_encoder = CLIPTextModel.from_pretrained(
    "openai/clip-vit-large-patch14")

# 文本 -> token -> 嵌入
prompt = "a beautiful sunset over the ocean, golden hour"
inputs = tokenizer(prompt, return_tensors="pt", padding="max_length",
                   max_length=77, truncation=True)

with torch.no_grad():
    outputs = text_encoder(input_ids=inputs.input_ids,
                           attention_mask=inputs.attention_mask)
    text_embeddings = outputs.last_hidden_state

print(f"Token count: {inputs.input_ids.shape[-1]}")  # 77
print(f"Embedding shape: {text_embeddings.shape}")     # (1, 77, 768)
print(f"Tokens: {tokenizer.convert_ids_to_tokens(inputs.input_ids[0][:10])}")

# 提示词嵌入相似度分析
import torch
import torch.nn.functional as F

def compute_text_similarity(texts):
    """计算提示词之间的语义相似度"""
    tokens = tokenizer(texts, return_tensors="pt", padding=True,
                       max_length=77, truncation=True)
    with torch.no_grad():
        embeddings = text_encoder(**tokens).last_hidden_state
        # 使用 [CLS] token 的嵌入
        cls_embeddings = embeddings[:, 0, :]
    
    similarities = F.cosine_similarity(
        cls_embeddings.unsqueeze(1),
        cls_embeddings.unsqueeze(0),
        dim=-1
    )
    
    for i, t1 in enumerate(texts):
        for j, t2 in enumerate(texts):
            if i < j:
                print(f"{t1[:30]:30s} <-> {t2[:30]:30s}: " 
                      f"{similarities[i,j]:.4f}")

texts = ["a dog running in a park",
         "a cat sleeping on a sofa",
         "a beautiful dog playing outdoors"]
compute_text_similarity(texts)

# Stable Diffusion U-Net 结构解析
from diffusers import UNet2DConditionModel
import torch

# 加载 U-Net
unet = UNet2DConditionModel.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    subfolder="unet",
    torch_dtype=torch.float16
)

# U-Net 参数统计
total_params = sum(p.numel() for p in unet.parameters())
print(f"Total params: {total_params / 1e6:.1f}M")  # ~860M

# 前向传播
cross_attention_dim = 768
latent = torch.randn(1, 4, 64, 64, dtype=torch.float16, device="cuda")
timestep = torch.tensor([500], device="cuda")
encoder_hidden = torch.randn(1, 77, cross_attention_dim, dtype=torch.float16, device="cuda")

with torch.no_grad():
    noise_pred = unet(latent, timestep, encoder_hidden_states=encoder_hidden).sample

print(f"Input:  {latent.shape}")     # (1, 4, 64, 64)
print(f"Output: {noise_pred.shape}")  # (1, 4, 64, 64)

# 分析 U-Net 各层结构
from diffusers import UNet2DConditionModel

unet = UNet2DConditionModel.from_pretrained(
    "runwayml/stable-diffusion-v1-5", subfolder="unet")

# 打印块结构
print("=== Down Blocks ===")
for i, block in enumerate(unet.down_blocks):
    has_attn = hasattr(block, "attentions") and block.attentions[0] is not None
    print(f"  Down {i}: channels={block.resnets[0].in_channels}, " 
          f"attn={has_attn}, downsample={block.downsamplers is not None}")

print("=== Middle Block ===")
mid = unet.mid_block
print(f"  Attention: {len(mid.attentions)} blocks")
print(f"  Resnets: {len(mid.resnets)} blocks")

print("=== Up Blocks ===")
for i, block in enumerate(unet.up_blocks):
    has_attn = hasattr(block, "attentions") and block.attentions[0] is not None
    print(f"  Up {i}: channels={block.resnets[0].in_channels}, " 
          f"attn={has_attn}, upsample={block.upsamplers is not None}")

# 不同采样器对比实验
from diffusers import StableDiffusionPipeline, \
    DDIMScheduler, EulerDiscreteScheduler, \
    DPMSolverMultistepScheduler, LMSDiscreteScheduler
import torch, time

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

samplers = {
    "DDIM": DDIMScheduler.from_config(pipe.scheduler.config),
    "Euler": EulerDiscreteScheduler.from_config(pipe.scheduler.config),
    "DPM++ 2M": DPMSolverMultistepScheduler.from_config(pipe.scheduler.config),
    "PLMS": LMSDiscreteScheduler.from_config(pipe.scheduler.config),
}

prompt = "a portrait of a warrior queen, fantasy art"
results = []
for name, scheduler in samplers.items():
    pipe.scheduler = scheduler
    start = time.time()
    img = pipe(prompt, num_inference_steps=25, generator=torch.manual_seed(42)).images[0]
    elapsed = time.time() - start
    results.append((name, f"{elapsed:.2f}s"))
    img.save(f"output_{name}.png")
    print(f"{name:12s} | {elapsed:.2f}s")

# 采样步数收敛曲线
import torch
from diffusers import DPMSolverMultistepScheduler

def convergence_test(pipe, prompt, steps_list, seed=42):
    """测试不同步数下的生成稳定性"""
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config)
    
    # 基准：100 步
    pipe.scheduler.set_timesteps(100)
    gen = torch.Generator("cuda").manual_seed(seed)
    ref = pipe(prompt, num_inference_steps=100, generator=gen).images[0]
    
    for steps in steps_list:
        gen = torch.Generator("cuda").manual_seed(seed)
        img = pipe(prompt, num_inference_steps=steps, generator=gen).images[0]
        # 计算与基准的 LPIPS 距离
        lpips_val = compute_lpips(ref, img)
        print(f"Steps: {steps:3d} | LPIPS vs 100-step: {lpips_val:.4f}")

convergence_test(pipe, "a cute cat", steps_list=[4, 8, 12, 16, 20, 25, 50])

# ControlNet：边缘图控制生成
from diffusers import StableDiffusionControlNetPipeline, \
    ControlNetModel, UniPCMultistepScheduler
from controlnet_aux import CannyDetector
import torch
from PIL import Image

# 加载 Canny ControlNet
canny = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-canny", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=canny, torch_dtype=torch.float16)
pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

# 提取边缘图并生成
input_img = Image.open("input.png").convert("RGB")
canny_img = CannyDetector()(input_img, low_threshold=100, high_threshold=200)

result = pipe(
    prompt="a beautiful landscape, photorealistic, 4k",
    image=canny_img,
    num_inference_steps=20,
    controlnet_conditioning_scale=0.8,
).images[0]

# LoRA 训练与推理
from diffusers import StableDiffusionPipeline
from peft import LoraConfig, get_peft_model
import torch

# 方法一：加载预训练 LoRA 权重
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16)
pipe = pipe.to("cuda")

# 加载 LoRA（仅 ~200MB，远小于全模型 4GB）
pipe.load_lora_weights("path/to/lora_weights", weight_name="pytorch_lora.bin")
pipe.fuse_lora(lora_scale=0.8)  # 融合到主模型

# 生成
image = pipe("a warrior in anime style").images[0]

# 方法二：使用 diffusers + peft 训练
lora_config = LoraConfig(
    r=16,
    lora_alpha=27,
    target_modules=["to_q", "to_k", "to_v", "to_out.0"],
    lora_dropout=0.05,
)
print(f"LoRA params: {sum(p.numel() for p in lora_params) / 1e6:.2f}M")

# 高级 diffusers 管道配置
from diffusers import (
    StableDiffusionImg2ImgPipeline,
    StableDiffusionInpaintPipeline,
    DPMSolverMultistepScheduler,
)
import torch
from PIL import Image

# 优化配置：显存效率 + 推理加速
def setup_pipeline(model_id="runwayml/stable-diffusion-v1-5"):
    pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
        model_id,
        torch_dtype=torch.float16,
        use_safetensors=True,
    )
    # 显存优化
    pipe.enable_attention_slicing()
    pipe.enable_vae_slicing()
    # 推理加速（需要安装 xformers）
    try:
        pipe.enable_xformers_memory_efficient_attention()
    except Exception:
        print("xformers not available, using default attention")
    
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config)
    return pipe.to("cuda")

pipe = setup_pipeline()
print(f"Pipeline ready, device: {pipe.device}")

# 完整工作流：文生图 -> 图生图 -> 局部重绘
import torch
from PIL import Image
import numpy as np

pipe = setup_pipeline()
generator = torch.Generator("cuda").manual_seed(42)

# Step 1: 文生图
prompt = "a medieval castle on a cliff, epic lighting"
result = pipe(
    prompt=prompt,
    negative_prompt="ugly, blurry, watermark",
    num_inference_steps=30,
    guidance_scale=7.5,
    generator=generator,
).images[0]
result.save("step1_castle.png")

# Step 2: 图生图（添加季节变化）
autumn_result = pipe(
    prompt=prompt + ", autumn foliage, golden leaves",
    image=result,
    strength=0.4,  # 0=完全保留, 1=完全重绘
    num_inference_steps=30,
    guidance_scale=7.5,
    generator=generator,
).images[0]
autumn_result.save("step2_autumn.png")

# Step 3: 局部重绘（添加飞鸟）
mask = create_bird_mask(result.size)  # 自定义掩码
bird_result = inpaint_pipe(
    prompt="birds flying in the sky",
    image=autumn_result,
    mask_image=mask,
    num_inference_steps=30,
).images[0]
bird_result.save("step3_final.png")