필사 모드: Diffusion Model 论文综述：从 DDPM 到 Stable Diffusion、DiT、SDXL，图像生成模型的演进

import torch
import torch.nn as nn
import numpy as np

class DDPMScheduler:
    """DDPM Forward Process 调度器"""
    def __init__(self, num_timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.num_timesteps = num_timesteps
        # 线性噪声调度
        self.betas = torch.linspace(beta_start, beta_end, num_timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)

    def add_noise(self, x_0, t, noise=None):
        """生成任意时间步 t 上的噪声图像"""
        if noise is None:
            noise = torch.randn_like(x_0)

        sqrt_alpha_bar = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
        sqrt_one_minus_alpha_bar = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)

        # x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
        x_t = sqrt_alpha_bar * x_0 + sqrt_one_minus_alpha_bar * noise
        return x_t

    def sample_timesteps(self, batch_size):
        """用于训练的随机时间步采样"""
        return torch.randint(0, self.num_timesteps, (batch_size,))

class DDPMSampler:
    """DDPM Reverse Process 采样器"""
    def __init__(self, scheduler):
        self.scheduler = scheduler

    @torch.no_grad()
    def sample(self, model, shape, device):
        """DDPM 逆扩散采样"""
        # 从纯噪声开始
        x = torch.randn(shape, device=device)

        for t in reversed(range(self.scheduler.num_timesteps)):
            t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)

            # 预测噪声
            predicted_noise = model(x, t_batch)

            # 计算均值
            alpha = self.scheduler.alphas[t]
            alpha_bar = self.scheduler.alphas_cumprod[t]
            beta = self.scheduler.betas[t]

            mean = (1 / torch.sqrt(alpha)) * (
                x - (beta / torch.sqrt(1 - alpha_bar)) * predicted_noise
            )

            # 仅当 t > 0 时添加噪声
            if t > 0:
                noise = torch.randn_like(x)
                sigma = torch.sqrt(beta)
                x = mean + sigma * noise
            else:
                x = mean

        return x

def ddpm_training_step(model, x_0, scheduler, optimizer):
    """DDPM 单步训练"""
    batch_size = x_0.shape[0]
    device = x_0.device

    # 1. 随机采样时间步
    t = scheduler.sample_timesteps(batch_size).to(device)

    # 2. 生成噪声并生成带噪图像
    noise = torch.randn_like(x_0)
    x_t = scheduler.add_noise(x_0, t, noise)

    # 3. 模型预测噪声
    predicted_noise = model(x_t, t)

    # 4. 计算 Simple Loss
    loss = nn.functional.mse_loss(predicted_noise, noise)

    # 5. 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss.item()

class DDIMSampler:
    """DDIM 加速采样器"""
    def __init__(self, scheduler, ddim_steps=50, eta=0.0):
        self.scheduler = scheduler
        self.ddim_steps = ddim_steps
        self.eta = eta
        # 生成时间步子集(例如：1000 -> 50)
        self.timesteps = np.linspace(
            0, scheduler.num_timesteps - 1, ddim_steps, dtype=int
        )[::-1]

    @torch.no_grad()
    def sample(self, model, shape, device):
        """DDIM 加速采样 - 用 50 步生成高质量结果"""
        x = torch.randn(shape, device=device)

        for i in range(len(self.timesteps)):
            t = self.timesteps[i]
            t_prev = self.timesteps[i + 1] if i + 1 < len(self.timesteps) else 0

            t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)
            predicted_noise = model(x, t_batch)

            alpha_bar_t = self.scheduler.alphas_cumprod[t]
            alpha_bar_prev = self.scheduler.alphas_cumprod[t_prev]

            # 预测 x_0
            x_0_pred = (x - torch.sqrt(1 - alpha_bar_t) * predicted_noise) / torch.sqrt(alpha_bar_t)
            x_0_pred = torch.clamp(x_0_pred, -1, 1)

            # 计算方向
            sigma = self.eta * torch.sqrt(
                (1 - alpha_bar_prev) / (1 - alpha_bar_t) * (1 - alpha_bar_t / alpha_bar_prev)
            )
            direction = torch.sqrt(1 - alpha_bar_prev - sigma**2) * predicted_noise

            # 计算 x_{t-1}
            x = torch.sqrt(alpha_bar_prev) * x_0_pred + direction

            if self.eta > 0 and t > 0:
                x = x + sigma * torch.randn_like(x)

        return x

import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler

class LatentDiffusionInference:
    """Stable Diffusion 推理流水线(简化版)"""

    def __init__(self, model_id="stable-diffusion-v1-5/stable-diffusion-v1-5"):
        self.pipe = StableDiffusionPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.float16,
            safety_checker=None
        ).to("cuda")

        # 替换为 DDIM 调度器(加速到 50 步)
        self.pipe.scheduler = DDIMScheduler.from_config(
            self.pipe.scheduler.config
        )

    def generate(self, prompt, negative_prompt="", num_steps=50, guidance_scale=7.5):
        """文本-图像生成"""
        image = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
        ).images[0]
        return image

    def generate_with_latent_control(self, prompt, seed=42):
        """直接控制潜在空间"""
        generator = torch.Generator(device="cuda").manual_seed(seed)

        # 直接生成潜在向量
        latents = torch.randn(
            (1, 4, 64, 64),
            generator=generator,
            device="cuda",
            dtype=torch.float16
        )

        image = self.pipe(
            prompt=prompt,
            latents=latents,
            num_inference_steps=50,
            guidance_scale=7.5,
        ).images[0]
        return image

class CrossAttention(nn.Module):
    """Stable Diffusion U-Net 的 Cross-Attention 层"""
    def __init__(self, d_model=320, d_context=768, n_heads=8):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads

        self.to_q = nn.Linear(d_model, d_model, bias=False)
        self.to_k = nn.Linear(d_context, d_model, bias=False)
        self.to_v = nn.Linear(d_context, d_model, bias=False)
        self.to_out = nn.Linear(d_model, d_model)

    def forward(self, x, context):
        """
        x: 图像潜在表示 (B, H*W, d_model)
        context: 文本嵌入 (B, seq_len, d_context)
        """
        B, N, C = x.shape

        q = self.to_q(x).view(B, N, self.n_heads, self.d_head).transpose(1, 2)
        k = self.to_k(context).view(B, -1, self.n_heads, self.d_head).transpose(1, 2)
        v = self.to_v(context).view(B, -1, self.n_heads, self.d_head).transpose(1, 2)

        # Scaled Dot-Product Attention
        scale = self.d_head ** -0.5
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
        attn = torch.softmax(attn, dim=-1)
        out = torch.matmul(attn, v)

        out = out.transpose(1, 2).contiguous().view(B, N, C)
        return self.to_out(out)

def classifier_free_guidance_step(model, x_t, t, text_embedding, null_embedding, guidance_scale=7.5):
    """Classifier-free Guidance 单步"""

    # 把条件/无条件预测合并为一个批次一次性处理
    x_in = torch.cat([x_t, x_t], dim=0)
    t_in = torch.cat([t, t], dim=0)
    c_in = torch.cat([null_embedding, text_embedding], dim=0)

    # 用一次 forward pass 同时生成两个预测
    noise_pred = model(x_in, t_in, encoder_hidden_states=c_in)
    noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)

    # 应用 CFG
    noise_pred_guided = noise_pred_uncond + guidance_scale * (
        noise_pred_cond - noise_pred_uncond
    )
    return noise_pred_guided

class DiTBlock(nn.Module):
    """DiT 的 adaLN-Zero Transformer Block"""
    def __init__(self, d_model, n_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(d_model, elementwise_affine=False)
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.norm2 = nn.LayerNorm(d_model, elementwise_affine=False)
        self.mlp = nn.Sequential(
            nn.Linear(d_model, d_model * 4),
            nn.GELU(),
            nn.Linear(d_model * 4, d_model),
        )
        # adaLN modulation：6 个参数(gamma1, beta1, alpha1, gamma2, beta2, alpha2)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(d_model, 6 * d_model),
        )
        # Zero 初始化 - 使模型在训练初期表现为 identity
        nn.init.zeros_(self.adaLN_modulation[-1].weight)
        nn.init.zeros_(self.adaLN_modulation[-1].bias)

    def forward(self, x, c):
        """
        x: patch token (B, N, D)
        c: 条件嵌入 - 时间步 + 类别 (B, D)
        """
        # 生成 adaLN 参数
        shift1, scale1, gate1, shift2, scale2, gate2 = (
            self.adaLN_modulation(c).chunk(6, dim=-1)
        )

        # Self-Attention with adaLN
        h = self.norm1(x)
        h = h * (1 + scale1.unsqueeze(1)) + shift1.unsqueeze(1)
        h, _ = self.attn(h, h, h)
        x = x + gate1.unsqueeze(1) * h

        # FFN with adaLN
        h = self.norm2(x)
        h = h * (1 + scale2.unsqueeze(1)) + shift2.unsqueeze(1)
        h = self.mlp(h)
        x = x + gate2.unsqueeze(1) * h

        return x

class PatchEmbed(nn.Module):
    """DiT 的 Patchify 层"""
    def __init__(self, patch_size=2, in_channels=4, embed_dim=1152):
        super().__init__()
        self.patch_size = patch_size
        self.proj = nn.Conv2d(
            in_channels, embed_dim,
            kernel_size=patch_size, stride=patch_size
        )

    def forward(self, x):
        """(B, C, H, W) -> (B, N, D) patch token 序列"""
        x = self.proj(x)  # (B, D, H/p, W/p)
        x = x.flatten(2).transpose(1, 2)  # (B, N, D)
        return x

from diffusers import StableDiffusionXLPipeline
import torch

class SDXLInference:
    """SDXL 推理流水线"""

    def __init__(self):
        self.pipe = StableDiffusionXLPipeline.from_pretrained(
            "stabilityai/stable-diffusion-xl-base-1.0",
            torch_dtype=torch.float16,
            variant="fp16",
            use_safetensors=True,
        ).to("cuda")

        # 内存优化
        self.pipe.enable_model_cpu_offload()
        self.pipe.enable_vae_tiling()

    def generate(self, prompt, negative_prompt="", steps=30):
        """SDXL 基础生成"""
        image = self.pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=steps,
            guidance_scale=7.5,
            height=1024,
            width=1024,
        ).images[0]
        return image

    def generate_with_refiner(self, prompt, base_pipe, refiner_pipe):
        """Base + Refiner 两阶段流水线"""
        # Base 模型：占全部步数的 80%
        high_noise_frac = 0.8
        image = base_pipe(
            prompt=prompt,
            num_inference_steps=40,
            denoising_end=high_noise_frac,
            output_type="latent",
        ).images

        # Refiner：剩余的 20%(提升细节质量)
        image = refiner_pipe(
            prompt=prompt,
            num_inference_steps=40,
            denoising_start=high_noise_frac,
            image=image,
        ).images[0]
        return image

def get_sdxl_conditioning(original_size, crop_coords, target_size):
    """生成 SDXL 的尺寸/裁剪条件"""
    # 原始尺寸 (height, width)
    original_size = torch.tensor(original_size, dtype=torch.float32)
    # 裁剪坐标 (top, left)
    crop_coords = torch.tensor(crop_coords, dtype=torch.float32)
    # 目标尺寸 (height, width)
    target_size = torch.tensor(target_size, dtype=torch.float32)

    # Fourier Feature Encoding
    conditioning = torch.cat([original_size, crop_coords, target_size])

    # Sinusoidal embedding
    freqs = torch.exp(
        -torch.arange(0, 128) * np.log(10000) / 128
    )
    emb = conditioning.unsqueeze(-1) * freqs.unsqueeze(0)
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)

    return emb.flatten()

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from controlnet_aux import CannyDetector
from PIL import Image
import torch

def controlnet_canny_generation(input_image_path, prompt):
    """基于 ControlNet Canny Edge 的图像生成"""
    # 加载 ControlNet 模型
    controlnet = ControlNetModel.from_pretrained(
        "lllyasviel/control_v11p_sd15_canny",
        torch_dtype=torch.float16,
    )

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

    # 提取 Canny Edge
    canny_detector = CannyDetector()
    input_image = Image.open(input_image_path)
    canny_image = canny_detector(input_image, low_threshold=100, high_threshold=200)

    # 基于 ControlNet 生成
    output = pipe(
        prompt=prompt,
        image=canny_image,
        num_inference_steps=30,
        guidance_scale=7.5,
        controlnet_conditioning_scale=1.0,
    ).images[0]

    return output

# LoRA 微调(Stable Diffusion)
accelerate launch train_text_to_image_lora.py \
    --pretrained_model_name_or_path="stable-diffusion-v1-5/stable-diffusion-v1-5" \
    --dataset_name="custom_dataset" \
    --resolution=512 \
    --train_batch_size=4 \
    --gradient_accumulation_steps=4 \
    --learning_rate=1e-4 \
    --lr_scheduler="cosine" \
    --lr_warmup_steps=500 \
    --max_train_steps=10000 \
    --rank=64 \
    --output_dir="./lora_output" \
    --mixed_precision="fp16" \
    --enable_xformers_memory_efficient_attention

# DreamBooth 微调(特定对象/风格学习)
accelerate launch train_dreambooth.py \
    --pretrained_model_name_or_path="stable-diffusion-v1-5/stable-diffusion-v1-5" \
    --instance_data_dir="./my_images" \
    --instance_prompt="a photo of sks dog" \
    --class_data_dir="./class_images" \
    --class_prompt="a photo of dog" \
    --with_prior_preservation \
    --prior_loss_weight=1.0 \
    --num_class_images=200 \
    --resolution=512 \
    --train_batch_size=1 \
    --learning_rate=5e-6 \
    --max_train_steps=800

import torch
from diffusers import StableDiffusionXLPipeline, DPMSolverMultistepScheduler

def optimized_sdxl_pipeline():
    """经过生产环境优化的 SDXL 流水线"""
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "stabilityai/stable-diffusion-xl-base-1.0",
        torch_dtype=torch.float16,
        variant="fp16",
        use_safetensors=True,
    ).to("cuda")

    # 1. 应用高速调度器
    pipe.scheduler = DPMSolverMultistepScheduler.from_config(
        pipe.scheduler.config,
        algorithm_type="dpmsolver++",
        use_karras_sigmas=True,
    )

    # 2. VAE Tiling(高分辨率生成时节省内存)
    pipe.enable_vae_tiling()

    # 3. Attention Slicing(VRAM 不足时)
    pipe.enable_attention_slicing()

    # 4. torch.compile(PyTorch 2.0+)
    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)

    return pipe

# GPU 内存监控
def monitor_gpu_memory():
    """监控 GPU 内存使用量"""
    allocated = torch.cuda.memory_allocated() / 1024**3
    reserved = torch.cuda.memory_reserved() / 1024**3
    max_allocated = torch.cuda.max_memory_allocated() / 1024**3
    print(f"Allocated: {allocated:.2f} GB")
    print(f"Reserved:  {reserved:.2f} GB")
    print(f"Peak:      {max_allocated:.2f} GB")

# 检查 GPU 内存状态
nvidia-smi --query-gpu=memory.used,memory.total --format=csv

# 检查 Python 进程的 GPU 内存泄漏
fuser -v /dev/nvidia*

# 强制释放 GPU 内存(无需重启进程)
python -c "
import torch
import gc
gc.collect()
torch.cuda.empty_cache()
print('GPU memory cleared')
print(f'Allocated: {torch.cuda.memory_allocated()/1024**3:.2f} GB')
"

# 发生 OOM 时的服务恢复流程
# 1. 对相应 worker 进程执行 graceful shutdown
# 2. 确认 GPU 内存已释放
# 3. 调整批次大小/并发请求数
# 4. 重启 worker 进程
# 5. 确认健康检查通过后恢复流量

# Safety Checker 设置(生产环境必需)
pipe = StableDiffusionPipeline.from_pretrained(
    "stable-diffusion-v1-5/stable-diffusion-v1-5",
    safety_checker=None,  # 仅在开发环境中禁用
)

# 生产环境中必须启用
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
from transformers import CLIPImageProcessor

safety_checker = StableDiffusionSafetyChecker.from_pretrained(
    "CompVis/stable-diffusion-safety-checker"
)
feature_extractor = CLIPImageProcessor.from_pretrained(
    "openai/clip-vit-base-patch32"
)

import os
from diffusers import StableDiffusionXLPipeline

def robust_model_loading(model_id, max_retries=3):
    """稳健的模型加载(含重试)"""
    for attempt in range(max_retries):
        try:
            pipe = StableDiffusionXLPipeline.from_pretrained(
                model_id,
                torch_dtype=torch.float16,
                use_safetensors=True,
                local_files_only=os.path.exists(
                    os.path.join(model_id, "model_index.json")
                ),
            )
            pipe = pipe.to("cuda")
            # 执行 warmup
            _ = pipe("test", num_inference_steps=1)
            return pipe
        except Exception as e:
            print(f"Attempt {attempt + 1} failed: {e}")
            if attempt < max_retries - 1:
                import time
                time.sleep(10)
                # 清除缓存后重试
                torch.cuda.empty_cache()
            else:
                raise RuntimeError(f"Model loading failed after {max_retries} attempts")

# CFG Scale 指南
guidance_scale_guidelines:
  1.0: '几乎忽略条件 - 接近随机生成'
  3.0-5.0: '富有创意且多样的生成'
  7.0-8.5: '常规推荐范围 - 质量/多样性平衡'
  10.0-15.0: '文本一致性高 - 存在过饱和风险'
  20.0+: '过度引导 - 产生伪影'

# 问题诊断清单
troubleshooting:
  blurry_output:
    - 'num_inference_steps 增加(至少 30 以上)'
    - '将调度器切换为 DPM-Solver++'
  oversaturated:
    - '将 guidance_scale 降低到 7.0 以下'
    - "在 negative_prompt 中添加 'oversaturated, vivid'"
  wrong_composition:
    - '改进提示词结构(明确主语-动词-宾语)'
    - '用 ControlNet 控制构图'