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ブロック"""
    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: パッチトークン (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) パッチトークンシーケンス"""
        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 2段階パイプライン"""
        # 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. 該当ワーカープロセスのgraceful shutdown
# 2. GPUメモリ解放を確認
# 3. バッチサイズ/同時リクエスト数を調整
# 4. ワーカープロセス再起動
# 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")
            # ウォームアップ実行
            _ = 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で構図を制御'