图像生成：StyleGAN, Diffusion

"""生成模型三大范式对比"""
from abc import ABC, abstractmethod
import torch
import torch.nn as nn

class GenerativeModel(ABC):
    """生成模型抽象基类"""
    @abstractmethod
    def sample(self, n: int) -> torch.Tensor:
        """从模型采样 n 个样本"""
        pass

    @abstractmethod
    def likelihood(self, x: torch.Tensor) -> torch.Tensor:
        """计算样本对数似然（部分模型可计算）"""
        pass

    @abstractmethod
    def encode(self, x: torch.Tensor) -> torch.Tensor:
        """编码到潜在空间（如适用）"""
        pass

class VAE(GenerativeModel):
    """变分自编码器：显式密度模型"""
    def __init__(self, latent_dim: int = 128):
        self.latent_dim = latent_dim
        # 编码器 q(z|x) 和解码器 p(x|z)
        self.encoder = nn.Sequential(
            nn.Conv2d(3, 64, 4, 2, 1), nn.ReLU(),
            nn.Conv2d(64, 128, 4, 2, 1), nn.ReLU(),
            nn.Flatten(),
        )
        self.fc_mu = nn.Linear(128 * 8 * 8, latent_dim)
        self.fc_logvar = nn.Linear(128 * 8 * 8, latent_dim)
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, 128 * 8 * 8),
            nn.Unflatten(1, (128, 8, 8)),
            nn.ConvTranspose2d(128, 64, 4, 2, 1), nn.ReLU(),
            nn.ConvTranspose2d(64, 3, 4, 2, 1), nn.Sigmoid(),
        )

    def reparameterize(self, mu: torch.Tensor, logvar: torch.Tensor) -> torch.Tensor:
        """重参数化技巧：z = mu + sigma * eps"""
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def sample(self, n: int) -> torch.Tensor:
        z = torch.randn(n, self.latent_dim)
        return self.decoder(z)

    def likelihood(self, x: torch.Tensor) -> torch.Tensor:
        mu, logvar = self._encode(x)
        recon = self.decoder(self.reparameterize(mu, logvar))
        # 重构损失（高斯假设下的负对数似然）
        return -nn.MSELoss(reduction='none')(recon, x).sum(dim=[1, 2, 3])

"""生成模型能力矩阵"""
model_comparison = {
    "VAE": {
        "显式密度": "是（可计算 ELBO）",
        "采样速度": "极快（一次前向传播）",
        "图像质量": "中等（模糊倾向）",
        "训练稳定性": "非常稳定",
        "模式覆盖": "好（不易模式崩溃）",
        "典型应用": "表示学习、异常检测",
    },
    "GAN": {
        "显式密度": "否（隐式模型）",
        "采样速度": "快（一次前向传播）",
        "图像质量": "极高（清晰锐利）",
        "训练稳定性": "不稳定（需精心设计）",
        "模式覆盖": "差（易模式崩溃）",
        "典型应用": "超分辨率、风格迁移",
    },
    "Diffusion": {
        "显式密度": "近似（变分下界）",
        "采样速度": "慢（多步迭代去噪）",
        "图像质量": "极高（超越 GAN）",
        "训练稳定性": "稳定（MLE 目标）",
        "模式覆盖": "极好",
        "典型应用": "文本生成图像、视频生成",
    },
    "Flow": {
        "显式密度": "是（精确似然）",
        "采样速度": "中等（可逆网络）",
        "图像质量": "中高",
        "训练稳定性": "稳定",
        "模式覆盖": "好",
        "典型应用": "密度估计、可逆编辑",
    },
}

print(f"{'模型':<12} | {'质量':<6} | {'速度':<6} | {'稳定性':<8}")
print("-" * 40)
for name, props in model_comparison.items():
    print(f"{name:<12} | {props['图像质量']:<14} | {props['采样速度']:<8} | {props['训练稳定性']}")

import torch
import torch.nn as nn
import math

class MappingNetwork(nn.Module):
    """StyleGAN 映射网络：z -> w"""
    def __init__(self, z_dim: int = 512, w_dim: int = 512, num_layers: int = 8):
        super().__init__()
        layers = []
        for _ in range(num_layers):
            layers.extend([
                nn.Linear(z_dim if _ == 0 else w_dim, w_dim),
                nn.LeakyReLU(0.2, inplace=True),
            ])
        self.mapping = nn.Sequential(*layers)
        # 初始化：缩小方差
        for layer in self.mapping:
            if isinstance(layer, nn.Linear):
                fan_in = layer.in_features
                layer.weight.data *= 1.0 / math.sqrt(fan_in)

    def forward(self, z: torch.Tensor) -> torch.Tensor:
        """归一化输入 + 映射"""
        z = z / torch.linalg.norm(z, dim=-1, keepdim=True)  # 归一化
        return self.mapping(z)

class AdaIN(nn.Module):
    """自适应实例归一化"""
    def __init__(self, channels: int, w_dim: int = 512):
        super().__init__()
        self.norm = nn.InstanceNorm2d(channels, affine=False)
        self.style_scale = nn.Linear(w_dim, channels)
        self.style_bias = nn.Linear(w_dim, channels)
        # 初始化：输出接近 1 和 0
        nn.init.zeros_(self.style_scale.weight)
        nn.init.ones_(self.style_scale.bias)
        nn.init.zeros_(self.style_bias.weight)
        nn.init.zeros_(self.style_bias.bias)

    def forward(self, x: torch.Tensor, w: torch.Tensor) -> torch.Tensor:
        """AdaIN: norm(x) * scale(w) + bias(w)"""
        normalized = self.norm(x)
        scale = self.style_scale(w).view(x.size(0), -1, 1, 1)
        bias = self.style_bias(w).view(x.size(0), -1, 1, 1)
        return normalized * scale + bias

class StyleGAN2Generator(nn.Module):
    """StyleGAN2 生成器简化版"""
    def __init__(self, z_dim=512, w_dim=512, img_size=128):
        super().__init__()
        self.mapping = MappingNetwork(z_dim, w_dim)
        self.w_dim = w_dim

        # 计算合成网络层数（从 4x4 到 img_size）
        self.num_layers = int(math.log2(img_size)) * 2 - 2

        # 常数输入 4x4
        self.const_input = nn.Parameter(torch.randn(1, 512, 4, 4))

        # 每层的 AdaIN + 卷积
        self.synthesis = nn.ModuleList()
        channels = 512
        for i in range(self.num_layers):
            if i > 0 and i % 2 == 0:
                channels = max(channels // 2, 64)
            self.synthesis.extend([
                nn.Conv2d(channels, channels, 3, padding=1),
                AdaIN(channels, w_dim),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Conv2d(channels, channels, 3, padding=1),
                AdaIN(channels, w_dim),
                nn.LeakyReLU(0.2, inplace=True),
                nn.Upsample(scale_factor=2) if i < self.num_layers - 2 else nn.Identity(),
            ])

        self.to_rgb = nn.Conv2d(channels, 3, 1)

    def forward(self, z: torch.Tensor, truncation: float = 0.7) -> torch.Tensor:
        w = self.mapping(z)
        # 截断技巧（Truncation Trick）：向均值收缩，提高质量、降低多样性
        w_avg = torch.zeros(1, self.w_dim, device=z.device)  # 预计算的 w 均值
        w = w_avg + truncation * (w - w_avg)

        x = self.const_input.expand(z.size(0), -1, -1, -1)
        for layer in self.synthesis:
            if isinstance(layer, AdaIN):
                x = layer(x, w)
            else:
                x = layer(x)
        return torch.tanh(self.to_rgb(x))

    def style_mixing(self, z1: torch.Tensor, z2: torch.Tensor,
                      cutoff_layer: int = 4) -> torch.Tensor:
        """风格混合：z1 提供 coarse 风格，z2 提供 fine 风格"""
        w1 = self.mapping(z1)
        w2 = self.mapping(z2)
        # 在 cutoff_layer 处切换风格
        w_mixed = torch.cat([w1[:, :cutoff_layer], w2[:, cutoff_layer:]], dim=1)
        return self._forward_with_w(w_mixed)

import torch
import torch.nn as nn
import math

class LinearNoiseScheduler:
    """线性噪声调度器"""
    def __init__(self, timesteps: int = 1000, beta_start: float = 1e-4,
                 beta_end: float = 0.02):
        self.timesteps = timesteps
        self.betas = torch.linspace(beta_start, beta_end, timesteps)
        self.alphas = 1.0 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

    def add_noise(self, x0: torch.Tensor, t: torch.Tensor) -> tuple:
        """前向加噪：一步得到 x_t（不需要逐步模拟）"""
        noise = torch.randn_like(x0)
        alpha_bar = self.alpha_bars[t].view(-1, 1, 1, 1)
        x_t = torch.sqrt(alpha_bar) * x0 + torch.sqrt(1 - alpha_bar) * noise
        return x_t, noise

    def get_alpha_bar(self, t: torch.Tensor) -> torch.Tensor:
        return self.alpha_bars[t].view(-1, 1, 1, 1)

class SimpleUNet(nn.Module):
    """简化版 U-Net（扩散模型的去噪网络）"""
    def __init__(self, in_channels: int = 3, base_channels: int = 64):
        super().__init__()
        self.time_embed = nn.Sequential(
            nn.Linear(1, base_channels * 4),
            nn.SiLU(),
            nn.Linear(base_channels * 4, base_channels * 4),
        )
        # 编码器
        self.enc1 = nn.Conv2d(in_channels, base_channels, 3, padding=1)
        self.enc2 = nn.Conv2d(base_channels, base_channels * 2, 3, padding=1)
        self.enc3 = nn.Conv2d(base_channels * 2, base_channels * 4, 3, padding=1)
        # 解码器
        self.dec3 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, 3, padding=1)
        self.dec2 = nn.ConvTranspose2d(base_channels * 2, base_channels, 3, padding=1)
        self.dec1 = nn.ConvTranspose2d(base_channels, in_channels, 3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        t_emb = self.time_embed(t.float().unsqueeze(-1)).unsqueeze(-1).unsqueeze(-1)
        e1 = torch.relu(self.enc1(x) + t_emb)
        e2 = torch.relu(self.enc2(self.pool(e1)) + t_emb)
        e3 = torch.relu(self.enc3(self.pool(e2)) + t_emb)
        d3 = torch.relu(self.dec3(e3) + e2 + t_emb)
        d2 = torch.relu(self.dec2(self.up(d3)) + e1 + t_emb)
        d1 = self.dec1(self.up(d2))
        return d1  # 预测的噪声

"""DDPM 训练与采样循环"""
import torch
from torch.utils.data import DataLoader

class DDPMSampler:
    """DDPM 采样器"""
    def __init__(self, scheduler: LinearNoiseScheduler, model: nn.Module):
        self.scheduler = scheduler
        self.model = model

    @torch.no_grad()
    def sample(self, n_samples: int, shape: tuple, device: str = 'cpu') -> torch.Tensor:
        """从纯噪声逐步去噪生成样本"""
        x = torch.randn(n_samples, *shape, device=device)
        for t in reversed(range(self.scheduler.timesteps)):
            t_batch = torch.full((n_samples,), t, device=device, dtype=torch.long)

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

            # 提取调度参数
            alpha = self.scheduler.alphas[t]
            alpha_bar = self.scheduler.alpha_bars[t]
            beta = self.scheduler.betas[t]

            # 反向步（DDPM 更新公式）
            mean = (1.0 / torch.sqrt(alpha)) * (
                x - (beta / torch.sqrt(1 - alpha_bar)) * predicted_noise
            )
            variance = beta

            # 除了最后一步，都添加噪声
            if t > 0:
                noise = torch.randn_like(x)
                x = mean + torch.sqrt(variance) * noise
            else:
                x = mean

        return torch.clamp(x, -1.0, 1.0)

# 训练循环示例
def train_diffusion(model, dataloader, scheduler, epochs=100):
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-4)
    for epoch in range(epochs):
        for x0 in dataloader:
            x0 = x0.cuda()
            # 随机采样时间步
            t = torch.randint(0, scheduler.timesteps, (x0.size(0),), device='cuda')
            # 前向加噪（一步到位）
            x_t, noise = scheduler.add_noise(x0, t)
            # 预测噪声
            predicted_noise = model(x_t, t)
            # MSE 损失
            loss = nn.MSELoss()(predicted_noise, noise)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
        if epoch % 10 == 0:
            print(f"Epoch {epoch}, Loss: {loss.item():.6f}")

class DDIMSampler:
    """DDIM 确定性采样器"""
    def __init__(self, scheduler: LinearNoiseScheduler, model: nn.Module):
        self.scheduler = scheduler
        self.model = model

    @torch.no_grad()
    def sample(self, n_samples: int, shape: tuple, num_steps: int = 50,
               eta: float = 0.0, device: str = 'cpu') -> torch.Tensor:
        """
        DDIM 采样（支持加速）
        num_steps: 采样步数（可以远小于 scheduler.timesteps）
        eta: 随机性系数，0=确定性(DDIM)，1=随机(DDPM)
        """
        # 选择子序列（均匀间隔的时间步）
        timesteps = torch.linspace(
            self.scheduler.timesteps - 1, 0, num_steps, dtype=torch.long
        )
        timesteps_next = torch.cat([timesteps[1:], torch.tensor([-1])])

        x = torch.randn(n_samples, *shape, device=device)

        for i, (t, t_next) in enumerate(zip(timesteps, timesteps_next)):
            t_batch = torch.full((n_samples,), t, device=device, dtype=torch.long)

            # 预测噪声
            eps_pred = self.model(x, t_batch)

            # 从 x_t 和 eps_pred 恢复 x_0 预测
            alpha_bar_t = self.scheduler.alpha_bars[t]
            alpha_bar_next = self.scheduler.alpha_bars[t_next] if t_next >= 0 else torch.ones_like(alpha_bar_t)

            x0_pred = (x - torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_bar_t)
            x0_pred = torch.clamp(x0_pred, -1.0, 1.0)  # 裁剪防止数值溢出

            if t_next < 0:
                # 最后一步：直接返回 x0_pred
                x = x0_pred
            else:
                # 方向指向 x0 的分量
                dir_xt = torch.sqrt(1 - alpha_bar_next - (eta ** 2) * (
                    (1 - alpha_bar_t) / alpha_bar_t
                )) * eps_pred
                # 随机分量（当 eta > 0 时）
                sigma = eta * torch.sqrt((1 - alpha_bar_t) / alpha_bar_t
                                         * (1 - alpha_bar_next) / (1 - alpha_bar_t))
                noise = torch.randn_like(x) if eta > 0 else 0

                x = torch.sqrt(alpha_bar_next) * x0_pred + dir_xt + sigma * noise

        return torch.clamp(x, -1.0, 1.0)

# 速度对比
for steps in [1000, 500, 100, 50, 25]:
    print(f"  DDIM {steps:>4} 步: {'~' * max(1, steps // 10)} ({1000/steps:.0f}x 加速)")

"""DDIM 图像编码与编辑"""
class DDIMInversion:
    """DDIM 反向编码：从图像找到对应的初始噪声"""
    def __init__(self, scheduler: LinearNoiseScheduler, model: nn.Module):
        self.scheduler = scheduler
        self.model = model

    @torch.no_grad()
    def invert(self, x0: torch.Tensor, num_steps: int = 50) -> torch.Tensor:
        """将真实图像 x0 编码为初始噪声 x_T"""
        timesteps = torch.linspace(0, self.scheduler.timesteps - 1, num_steps, dtype=torch.long)
        timesteps_next = torch.cat([timesteps[:1] - 1, timesteps[:-1]])

        x = x0.clone()
        for t, t_next in zip(timesteps, timesteps_next):
            t_batch = torch.full((x0.size(0),), t, device=x0.device, dtype=torch.long)

            eps_pred = self.model(x, t_batch)
            alpha_bar_t = self.scheduler.alpha_bars[t]
            alpha_bar_next = self.scheduler.alpha_bars[t_next] if t_next >= 0 else torch.zeros_like(alpha_bar_t)

            # 反向 DDIM 步
            x0_pred = (x - torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_bar_t)
            dir_xt = torch.sqrt(1 - alpha_bar_next) * eps_pred
            x = torch.sqrt(alpha_bar_next) * x0_pred + dir_xt

        return x

class ImageEditor:
    """基于 DDIM 的图像编辑器"""
    def __init__(self, sampler: DDIMSampler, inverter: DDIMInversion):
        self.sampler = sampler
        self.inverter = inverter

    def edit(self, image: torch.Tensor, edit_fn, num_steps: int = 50) -> torch.Tensor:
        """
        编辑图像：编码 → 编辑潜在空间 → 解码
        edit_fn: 对初始噪声 x_T 进行操作的函数
        """
        # 1. 编码到噪声空间
        x_T = self.inverter.invert(image, num_steps)
        # 2. 编辑噪声
        x_T_edited = edit_fn(x_T)
        # 3. 从编辑后的噪声重新采样
        # （注意：这里需要修改 sampler 从指定 x_T 开始）
        return self.sampler.sample_from_xt(x_T_edited, num_steps)

import torch
import numpy as np
from scipy import linalg
from torchvision.models import inception_v3, Inception_V3_Weights
import torch.nn.functional as F

class FIDEvaluator:
    """Fréchet Inception Distance 评估器"""
    def __init__(self, device: str = 'cuda'):
        self.device = device
        # 使用 Inception-v3 的 pool3 层（2048 维特征）
        weights = Inception_V3_Weights.IMAGENET1K_V1
        model = inception_v3(weights=weights, transform_input=False)
        model.fc = nn.Identity()  # 去掉分类头
        model.aux_logits = False
        self.model = model.eval().to(device)

    @torch.no_grad()
    def extract_features(self, images: torch.Tensor) -> np.ndarray:
        """提取 Inception 特征"""
        # images: [N, 3, 299, 299], 范围 [0, 1]
        images = images.to(self.device)
        features = self.model(images)  # [N, 2048]
        return features.cpu().numpy()

    @staticmethod
    def calculate_fid(mu1: np.ndarray, sigma1: np.ndarray,
                      mu2: np.ndarray, sigma2: np.ndarray) -> float:
        """计算 Fréchet 距离"""
        diff = mu1 - mu2
        covmean, _ = linalg.sqrtm(sigma1 @ sigma2, disp=False)
        if np.iscomplexobj(covmean):
            covmean = covmean.real
        fid = diff.dot(diff) + np.trace(sigma1 + sigma2 - 2 * covmean)
        return max(0, fid)  # FID 不应为负

    def evaluate(self, real_images: torch.Tensor,
                 fake_images: torch.Tensor) -> float:
        """评估真实集和生成集的 FID"""
        real_feats = self.extract_features(real_images)
        fake_feats = self.extract_features(fake_images)

        mu_real, sigma_real = real_feats.mean(0), np.cov(real_feats, rowvar=False)
        mu_fake, sigma_fake = fake_feats.mean(0), np.cov(fake_feats, rowvar=False)

        return self.calculate_fid(mu_real, sigma_real, mu_fake, sigma_fake)

"""Precision & Recall 评估"""
class PrecisionRecallEvaluator:
    """评估生成质量（Precision）和多样性（Recall）"""
    def __init__(self, device: str = 'cuda', k: int = 5):
        self.device = device
        self.k = k  # 最近邻数量

    @torch.no_grad()
    def extract_features(self, images: torch.Tensor) -> np.ndarray:
        """提取特征（简化为 Flattened 图像或预训练特征）"""
        return images.view(images.size(0), -1).cpu().numpy()

    def compute_precision_recall(self, real_feats: np.ndarray,
                                  fake_feats: np.ndarray) -> tuple:
        """
        Precision: 生成样本中落在真实数据流形内的比例
        Recall: 真实数据流形中被生成样本覆盖的比例
        """
        # 计算真实特征的 k-NN 半径
        from sklearn.neighbors import NearestNeighbors

        # 真实数据流形边界（用 k-NN 半径估计）
        nn_real = NearestNeighbors(n_neighbors=self.k).fit(real_feats)
        dists_real, _ = nn_real.kneighbors(real_feats)
        radii_real = dists_real[:, -1]  # 每个真实样本的 k-NN 半径

        # 生成数据流形边界
        nn_fake = NearestNeighbors(n_neighbors=self.k).fit(fake_feats)
        dists_fake, _ = nn_fake.kneighbors(fake_feats)
        radii_fake = dists_fake[:, -1]

        # Precision: 每个生成样本是否在真实流形内
        dists_fake_to_real, _ = nn_real.kneighbors(fake_feats)
        precision = np.mean(dists_fake_to_real[:, 0] < radii_real.min())

        # Recall: 每个真实样本是否在生成了流形内
        dists_real_to_fake, _ = nn_fake.kneighbors(real_feats)
        recall = np.mean(dists_real_to_fake[:, 0] < radii_fake.min())

        return float(precision), float(recall)

# 典型 FID 参考值
fid_benchmarks = {
    "StyleGAN2 (FFHQ 1024x1024)": 2.84,
    "StyleGAN3 (FFHQ 1024x1024)": 2.43,
    "Guided Diffusion (ImageNet 256x256)": 2.07,
    "DALL-E 2": 10.39,
    "Stable Diffusion v1.4": 9.63,
    "Stable Diffusion v2.1": 7.71,
    "SDXL 1.0": 6.64,
}

for model, fid in sorted(fid_benchmarks.items(), key=lambda x: x[1]):
    print(f"  {model:<45} FID: {fid:.2f}")

"""Classifier-Free Guidance (CFG)：Diffusion 的条件控制"""
class CFGDiffusionSampler:
    """带分类器自由引导的扩散模型采样器"""
    def __init__(self, model: nn.Module, scheduler: LinearNoiseScheduler):
        self.model = model  # 联合训练的条件/无条件模型
        self.scheduler = scheduler

    @torch.no_grad()
    def sample_cfg(self, n: int, shape: tuple, cond: torch.Tensor,
                   uncond: torch.Tensor, guidance_scale: float = 7.5,
                   num_steps: int = 50) -> torch.Tensor:
        """
        CFG 采样：同时运行条件和无条件前向传播
        guidance_scale: 越大 → 越忠于条件但质量可能下降
        典型值: 3-10（Stable Diffusion 默认 7.5）
        """
        timesteps = torch.linspace(
            self.scheduler.timesteps - 1, 0, num_steps, dtype=torch.long
        )
        timesteps_next = torch.cat([timesteps[1:], torch.tensor([-1])])

        x = torch.randn(n, *shape)

        for t, t_next in zip(timesteps, timesteps_next):
            t_batch = torch.full((n,), t, dtype=torch.long)

            # 条件预测
            eps_cond = self.model(x, t_batch, cond)
            # 无条件预测
            eps_uncond = self.model(x, t_batch, uncond)

            # CFG: 引导方向 = 无条件 + scale * (条件 - 无条件)
            eps_guided = eps_uncond + guidance_scale * (eps_cond - eps_uncond)

            # 使用 DDIM 更新
            alpha_bar_t = self.scheduler.alpha_bars[t]
            alpha_bar_next = self.scheduler.alpha_bars[t_next] if t_next >= 0 else torch.ones_like(alpha_bar_t)

            x0_pred = (x - torch.sqrt(1 - alpha_bar_t) * eps_guided) / torch.sqrt(alpha_bar_t)
            x0_pred = torch.clamp(x0_pred, -1.0, 1.0)

            if t_next >= 0:
                dir_xt = torch.sqrt(1 - alpha_bar_next) * eps_guided
                x = torch.sqrt(alpha_bar_next) * x0_pred + dir_xt
            else:
                x = x0_pred

        return torch.clamp(x, -1.0, 1.0)

# CFG scale 的影响
for scale in [1.0, 3.0, 5.0, 7.5, 10.0, 15.0]:
    quality = "低" if scale < 3 else "中" if scale < 7 else "高" if scale < 12 else "过饱和"
    diversity = "高" if scale < 5 else "中" if scale < 10 else "低"
    print(f"  CFG={scale:>4.1f}: 质量={quality:<6} 多样性={diversity}")

"""GAN vs Diffusion 决策树"""
def choose_generative_model(requirements: dict) -> str:
    """根据需求选择生成模型"""
    if requirements.get("real_time", False):
        return "StyleGAN2/3 (GAN) —— 需要实时生成"

    if requirements.get("max_quality", False):
        if requirements.get("text_conditioned", False):
            return "Stable Diffusion / DALL-E 3 (Diffusion)"
        return "Guided Diffusion (Diffusion)"

    if requirements.get("mobile_deploy", False):
        return "MobileStyleGAN (轻量 GAN) —— 端侧部署"

    if requirements.get("image_editing", False):
        return "DDIM Inversion + Diffusion —— 支持编码编辑"

    if requirements.get("high_diversity", False):
        return "Diffusion —— 更好的模式覆盖"

    if requirements.get("training_ease", False):
        return "Diffusion —— 训练稳定，几乎无需调参"

    if requirements.get("deterministic_output", False):
        return "DDIM (eta=0) —— 确定性采样"

    return "StyleGAN3 (GAN) —— 均衡选择"

# 典型场景决策
scenarios = [
    {"name": "AI 绘画 App", "requirements": {"max_quality": True, "text_conditioned": True}},
    {"name": "游戏 NPC 头像", "requirements": {"real_time": True, "max_quality": False}},
    {"name": "移动端滤镜", "requirements": {"mobile_deploy": True}},
    {"name": "产品图修改", "requirements": {"image_editing": True}},
    {"name": "数据集增强", "requirements": {"high_diversity": True}},
]

for scenario in scenarios:
    choice = choose_generative_model(scenario["requirements"])
    print(f"  {scenario['name']:<15} → {choice}")

"""完整扩散模型训练 Pipeline"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import save_image
import math

class SinusoidalTimeEmbedding(nn.Module):
    """正弦位置编码时间嵌入"""
    def __init__(self, dim: int):
        super().__init__()
        self.dim = dim

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        half_dim = self.dim // 2
        embeddings = math.log(10000.0) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=t.device) * -embeddings)
        embeddings = t.float().unsqueeze(-1) * embeddings.unsqueeze(0)
        embeddings = torch.cat([embeddings.sin(), embeddings.cos()], dim=-1)
        return embeddings

class SimpleDiffusionModel(nn.Module):
    """简化扩散模型 U-Net"""
    def __init__(self, img_size: int = 28, channels: int = 1, hidden_dim: int = 128):
        super().__init__()
        self.time_embed = SinusoidalTimeEmbedding(hidden_dim * 4)
        self.time_mlp = nn.Sequential(
            nn.Linear(hidden_dim * 4, hidden_dim * 4),
            nn.SiLU(),
            nn.Linear(hidden_dim * 4, hidden_dim * 4),
        )

        # 编码器
        self.down1 = nn.Conv2d(channels, hidden_dim, 3, padding=1)
        self.down2 = nn.Conv2d(hidden_dim, hidden_dim * 2, 4, 2, 1)
        self.down3 = nn.Conv2d(hidden_dim * 2, hidden_dim * 4, 4, 2, 1)

        # 中间层
        self.mid = nn.Sequential(
            nn.Conv2d(hidden_dim * 4, hidden_dim * 4, 3, padding=1),
            nn.SiLU(),
            nn.Conv2d(hidden_dim * 4, hidden_dim * 4, 3, padding=1),
        )

        # 解码器
        self.up1 = nn.ConvTranspose2d(hidden_dim * 4, hidden_dim * 2, 4, 2, 1)
        self.up2 = nn.ConvTranspose2d(hidden_dim * 2, hidden_dim, 4, 2, 1)
        self.up3 = nn.Conv2d(hidden_dim, channels, 3, padding=1)

    def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        t_emb = self.time_mlp(self.time_embed(t))
        t_emb = t_emb.unsqueeze(-1).unsqueeze(-1)

        d1 = F.silu(self.down1(x) + t_emb)
        d2 = F.silu(self.down2(d1) + t_emb)
        d3 = F.silu(self.down3(d2) + t_emb)
        m = self.mid(d3)
        u1 = F.silu(self.up1(m) + d2 + t_emb)
        u2 = F.silu(self.up2(u1) + d1 + t_emb)
        return self.up3(u2)

"""训练与采样主循环"""
def train_diffusion_mnist(epochs: int = 50, batch_size: int = 128):
    """在 MNIST 上训练扩散模型"""
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    scheduler = LinearNoiseScheduler(timesteps=1000)
    model = SimpleDiffusionModel().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

    # EMA 权重（用于采样）
    ema_model = SimpleDiffusionModel().to(device)
    ema_model.load_state_dict(model.state_dict())
    ema_decay = 0.995

    # 数据
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.5,), (0.5,)),  # 归一化到 [-1, 1]
    ])
    dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    for epoch in range(epochs):
        model.train()
        total_loss = 0
        for x0, _ in loader:
            x0 = x0.to(device)
            t = torch.randint(0, scheduler.timesteps, (x0.size(0),), device=device)

            # 前向加噪
            x_t, noise = scheduler.add_noise(x0, t)

            # 预测噪声
            pred_noise = model(x_t, t)
            loss = F.mse_loss(pred_noise, noise)

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()

            # EMA 更新
            for ema_p, p in zip(ema_model.parameters(), model.parameters()):
                ema_p.data = ema_decay * ema_p.data + (1 - ema_decay) * p.data

            total_loss += loss.item()

        print(f"Epoch {epoch+1}/{epochs}, Loss: {total_loss/len(loader):.6f}")

        # 每个 epoch 采样一张图看看进展
        if (epoch + 1) % 5 == 0:
            sampler = DDIMSampler(scheduler, ema_model)
            samples = sampler.sample(16, (1, 28, 28), num_steps=50, device=device)
            save_image((samples + 1) / 2, f'samples_epoch_{epoch+1}.png', nrow=4)

    torch.save(ema_model.state_dict(), 'diffusion_mnist_ema.pt')
    print("训练完成！模型已保存。")

# 推理：加载训练好的模型生成图像
def generate_samples(model_path: str = 'diffusion_mnist_ema.pt', n: int = 16):
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    scheduler = LinearNoiseScheduler(timesteps=1000)
    model = SimpleDiffusionModel().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    sampler = DDIMSampler(scheduler, model)
    samples = sampler.sample(n, (1, 28, 28), num_steps=50, device=device)
    save_image((samples + 1) / 2, 'generated.png', nrow=4)
    print(f"已生成 {n} 张样本 → generated.png")