第9章：3D扩散模型

def train_step(x_0, class_label, dropout_prob=0.1):
    # 随机时间步
    t = random.randint(1, T)

    # 添加噪声
    epsilon = randn_like(x_0)
    x_t = sqrt(alpha_bar[t]) * x_0 + sqrt(1 - alpha_bar[t]) * epsilon

    # 条件dropout
    if random.random() < dropout_prob:
        condition = null_token  # 使用特殊的空条件标记
    else:
        condition = class_embedding[class_label]

    # 预测噪声
    epsilon_pred = model(x_t, t, condition)

    # 计算损失
    loss = MSE(epsilon_pred, epsilon)
    return loss

def sample_with_cfg(class_label, guidance_scale=7.5):
    x_t = randn(shape)

    for t in reversed(range(T)):
        # 有条件预测
        epsilon_cond = model(x_t, t, class_embedding[class_label])

        # 无条件预测
        epsilon_uncond = model(x_t, t, null_token)

        # CFG组合
        epsilon = epsilon_uncond + guidance_scale * (epsilon_cond - epsilon_uncond)

        # 去噪步骤
        x_t = denoise_step(x_t, epsilon, t)

    return x_t

def check_consistency(p1_list, p2_list, F, threshold=1.0):
    consistent = []
    inconsistent = []

    for p1, p2 in zip(p1_list, p2_list):
        # 计算极线误差
        error = abs(p2.T @ F @ p1)

        if error < threshold:
            consistent.append((p1, p2))
        else:
            # 修正：投影到极线
            l2 = F @ p1  # p1在图2中的极线
            p2_corrected = project_to_line(p2, l2)

            l1 = F.T @ p2  # p2在图1中的极线
            p1_corrected = project_to_line(p1, l1)

            # 选择移动较小的修正
            if norm(p1 - p1_corrected) < norm(p2 - p2_corrected):
                inconsistent.append((p1_corrected, p2))
            else:
                inconsistent.append((p1, p2_corrected))

    return consistent, inconsistent

def project_to_line(point, line):
    # 点到直线的投影
    # line: ax + by + c = 0
    a, b, c = line
    x, y, _ = point
    t = -(a*x + b*y + c) / (a*a + b*b)
    return [x + t*a, y + t*b, 1]

def vsd_gradient(theta, phi):
    x = render(theta)
    t = sample_timestep()
    epsilon = randn_like(x)
    x_t = forward_diffusion(x, t, epsilon)

    # 优化一个独立的分布q_phi
    epsilon_q = q_network(x_t, t, phi)
    epsilon_p = pretrained_model(x_t, t)

    # VSD梯度
    grad = (epsilon_p - epsilon_q) * dx_dtheta

    # 同时更新q网络
    update_q_network(phi, epsilon_q, epsilon)

    return grad

def annealed_timestep(iteration, max_iter):
    progress = iteration / max_iter
    t_max = T * (1 - progress) + 1 * progress
    t_min = T * 0.5 * (1 - progress) + 1 * progress
    return random.uniform(t_min, t_max)

START
├─ 实时要求？
│  ├─ 是 → 内存受限？
│  │       ├─ 是 → 三平面/点云
│  │       └─ 否 → Mesh/体素
│  └─ 否 → 质量优先？
│          ├─ 是 → NeRF/高分辨率SDF
│          └─ 否 → 标准Mesh

class VariableLengthPointDiffusion(nn.Module):
    def __init__(self):
        # 1. 点编码器（处理变长）
        self.point_encoder = nn.Sequential(
            nn.Linear(3, 128),  # 坐标 → 特征
            nn.ReLU(),
            nn.Linear(128, 256)
        )

        # 2. Set Transformer（置换不变+变长）
        self.set_attention = nn.ModuleList([
            InducedSetAttention(256, num_inducing_points=32),
            SelfAttention(256),
            InducedSetAttention(256, num_inducing_points=32)
        ])

        # 3. 时间嵌入
        self.time_mlp = nn.Sequential(
            SinusoidalEmbedding(256),
            nn.Linear(256, 512),
            nn.GELU(),
            nn.Linear(512, 256)
        )

        # 4. 去噪网络
        self.denoiser = nn.ModuleList([
            CrossAttentionBlock(256),  # 全局信息
            GraphConvBlock(256, k=20),  # 局部结构
            PointwiseMLP(256)  # 逐点处理
        ])

    def forward(self, x, t, mask=None):
        # x: [B, N, 3], 变长通过mask处理
        B, N, _ = x.shape

        # 编码点
        h = self.point_encoder(x)  # [B, N, 256]

        # 时间编码
        t_emb = self.time_mlp(t)  # [B, 256]

        # Set Transformer处理变长
        for layer in self.set_attention:
            h = layer(h, mask=mask)

        # 加入时间信息
        h = h + t_emb.unsqueeze(1)

        # 去噪
        for layer in self.denoiser:
            h = layer(h, mask=mask)

        # 预测噪声
        return h[..., :3]  # [B, N, 3]

class InducedSetAttention(nn.Module):
    """使用诱导点处理变长序列"""
    def __init__(self, dim, num_inducing_points):
        super().__init__()
        self.inducing_points = nn.Parameter(
            torch.randn(num_inducing_points, dim)
        )
        self.attention = nn.MultiheadAttention(dim, 8)

    def forward(self, x, mask=None):
        # 通过固定数量的诱导点聚合信息
        induced = self.inducing_points.unsqueeze(0).expand(x.size(0), -1, -1)

        # 诱导点attend to输入
        induced, _ = self.attention(induced, x, x, key_padding_mask=mask)

        # 输入attend to诱导点
        x, _ = self.attention(x, induced, induced)

        return x

class MultiModalConditioner(nn.Module):
    def __init__(self):
        # 各模态编码器
        self.text_encoder = CLIPTextEncoder()
        self.image_encoder = DINOv2Encoder()
        self.sketch_encoder = SketchCNN()
        self.audio_encoder = AudioTransformer()  # 用于描述声音的3D源

        # 模态对齐
        self.projectors = nn.ModuleDict({
            'text': nn.Linear(768, 512),
            'image': nn.Linear(1024, 512),
            'sketch': nn.Linear(256, 512),
            'audio': nn.Linear(512, 512)
        })

        # 跨模态注意力
        self.cross_modal_attention = nn.ModuleList([
            CrossModalTransformer(512, num_heads=8)
            for _ in range(3)
        ])

        # 自适应融合
        self.fusion_gate = nn.Sequential(
            nn.Linear(512 * 4, 512),
            nn.Sigmoid()
        )

    def forward(self, conditions):
        embeddings = {}

        # 编码各模态
        for modality, data in conditions.items():
            if data is not None:
                feat = getattr(self, f'{modality}_encoder')(data)
                embeddings[modality] = self.projectors[modality](feat)

        # 填充缺失模态
        all_modalities = ['text', 'image', 'sketch', 'audio']
        for mod in all_modalities:
            if mod not in embeddings:
                embeddings[mod] = torch.zeros(B, 512)  # 学习的null embedding

        # 堆叠所有嵌入
        stacked = torch.stack(list(embeddings.values()), dim=1)  # [B, 4, 512]

        # 跨模态交互
        for layer in self.cross_modal_attention:
            stacked = layer(stacked)

        # 自适应融合
        concat = stacked.flatten(1, 2)  # [B, 4*512]
        gate = self.fusion_gate(concat)  # [B, 512]

        # 加权平均
        fused = (stacked * gate.unsqueeze(1)).mean(dim=1)  # [B, 512]

        return fused, embeddings  # 返回融合结果和各模态嵌入

class CrossModalTransformer(nn.Module):
    """跨模态注意力层"""
    def __init__(self, dim, num_heads):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(dim, num_heads)
        self.cross_attn = nn.MultiheadAttention(dim, num_heads)
        self.ffn = nn.Sequential(
            nn.Linear(dim, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim)
        )
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)

    def forward(self, x):
        # x: [B, num_modalities, dim]
        x = x + self.self_attn(x, x, x)[0]
        x = self.norm1(x)

        # 循环cross attention
        for i in range(x.size(1)):
            others = torch.cat([x[:, :i], x[:, i+1:]], dim=1)
            x[:, i:i+1] = x[:, i:i+1] + self.cross_attn(
                x[:, i:i+1], others, others
            )[0]

        x = self.norm2(x)
        x = x + self.ffn(x)
        x = self.norm3(x)

        return x

class FastDiffusionModel(nn.Module):
    def __init__(self, teacher_model):
        super().__init__()

        # 1. 知识蒸馏 - 学生网络
        self.student = nn.Sequential(
            # 更少的层
            ResBlock(256, 256),
            ResBlock(256, 256),
            ResBlock(256, 256)
        )

        # 2. 渐进式蒸馏
        self.progressive_steps = [1000, 500, 100, 50, 10, 4, 1]
        self.step_networks = nn.ModuleList([
            self.create_step_network(s) for s in self.progressive_steps
        ])

        # 3. 缓存机制
        self.cache = {}
        self.cache_size = 1000

        # 4. 量化准备
        self.quant = torch.quantization.QuantStub()
        self.dequant = torch.quantization.DeQuantStub()

    def create_step_network(self, num_steps):
        """为特定步数创建优化网络"""
        if num_steps > 100:
            return self.teacher_model  # 使用完整模型
        elif num_steps > 10:
            return self.student  # 使用学生模型
        else:
            # 极简网络
            return nn.Sequential(
                nn.Linear(256, 128),
                nn.ReLU(),
                nn.Linear(128, 256)
            )

    def fast_sample(self, condition, num_steps=50):
        """快速采样"""
        x = torch.randn(shape)

        # 选择合适的网络
        network_idx = min(range(len(self.progressive_steps)),
                         key=lambda i: abs(self.progressive_steps[i] - num_steps))
        network = self.step_networks[network_idx]

        # DDIM采样
        timesteps = torch.linspace(1000, 0, num_steps)

        for i, t in enumerate(timesteps[:-1]):
            t_next = timesteps[i + 1]

            # 缓存检查
            cache_key = (x.sum().item(), t.item(), condition.sum().item())
            if cache_key in self.cache:
                pred_noise = self.cache[cache_key]
            else:
                # 预测噪声
                with torch.cuda.amp.autocast():  # 混合精度
                    pred_noise = network(x, t, condition)

                # 更新缓存
                if len(self.cache) < self.cache_size:
                    self.cache[cache_key] = pred_noise.detach()

            # DDIM更新
            x = self.ddim_step(x, pred_noise, t, t_next)

        return x

    def compile_for_inference(self):
        """编译优化"""
        # 1. TorchScript
        self.student = torch.jit.script(self.student)

        # 2. 图优化
        self.student = torch.jit.optimize_for_inference(self.student)

        # 3. 量化
        self.student = torch.quantization.quantize_dynamic(
            self.student,
            {nn.Linear, nn.Conv2d},
            dtype=torch.qint8
        )

        # 4. ONNX导出（可选）
        dummy_input = torch.randn(1, 256)
        torch.onnx.export(self.student, dummy_input, "fast_model.onnx",
                         opset_version=11,
                         do_constant_folding=True)

# 使用示例
teacher = load_pretrained_model()
fast_model = FastDiffusionModel(teacher)
fast_model.compile_for_inference()

# 基准测试
@torch.no_grad()
def benchmark():
    times = []
    for _ in range(100):
        start = time.time()
        output = fast_model.fast_sample(condition, num_steps=4)
        times.append(time.time() - start)

    print(f"Average inference time: {np.mean(times):.3f}s")
    print(f"FPS: {1/np.mean(times):.1f}")

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt

# 1. 前向扩散过程
class ForwardDiffusion:
    def __init__(self, num_steps=1000, beta_start=0.0001, beta_end=0.02):
        self.num_steps = num_steps

        # 线性噪声调度
        self.betas = torch.linspace(beta_start, beta_end, num_steps)
        self.alphas = 1 - self.betas
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)

        # 预计算常数
        self.sqrt_alpha_bars = torch.sqrt(self.alpha_bars)
        self.sqrt_one_minus_alpha_bars = torch.sqrt(1 - self.alpha_bars)

    def add_noise(self, x0, t, noise=None):
        """添加噪声到时间步t"""
        if noise is None:
            noise = torch.randn_like(x0)

        sqrt_alpha_bar_t = self.sqrt_alpha_bars[t].view(-1, 1, 1)
        sqrt_one_minus_alpha_bar_t = self.sqrt_one_minus_alpha_bars[t].view(-1, 1, 1)

        return sqrt_alpha_bar_t * x0 + sqrt_one_minus_alpha_bar_t * noise, noise

# 2. 简单的U-Net去噪网络
class SimpleUNet2D(nn.Module):
    def __init__(self, in_channels=2, time_dim=32):
        super().__init__()

        # 时间嵌入
        self.time_mlp = nn.Sequential(
            nn.Linear(1, time_dim),
            nn.SiLU(),
            nn.Linear(time_dim, time_dim)
        )

        # 编码器
        self.enc1 = self.conv_block(in_channels, 64)
        self.enc2 = self.conv_block(64, 128)
        self.enc3 = self.conv_block(128, 256)

        # 中间层
        self.bottleneck = self.conv_block(256, 512)

        # 解码器
        self.dec3 = self.conv_block(512 + 256, 256)
        self.dec2 = self.conv_block(256 + 128, 128)
        self.dec1 = self.conv_block(128 + 64, 64)

        # 输出层
        self.output = nn.Conv2d(64, in_channels, 1)

        # 时间注入层
        self.time_layers = nn.ModuleList([
            nn.Linear(time_dim, 64),
            nn.Linear(time_dim, 128),
            nn.Linear(time_dim, 256),
            nn.Linear(time_dim, 512)
        ])

    def conv_block(self, in_c, out_c):
        return nn.Sequential(
            nn.Conv2d(in_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.SiLU(),
            nn.Conv2d(out_c, out_c, 3, padding=1),
            nn.BatchNorm2d(out_c),
            nn.SiLU()
        )

    def forward(self, x, t):
        # 时间嵌入
        t_emb = self.time_mlp(t.float().unsqueeze(-1))

        # 编码
        e1 = self.enc1(x)
        e1 = e1 + self.time_layers[0](t_emb).unsqueeze(-1).unsqueeze(-1)

        e2 = self.enc2(F.max_pool2d(e1, 2))
        e2 = e2 + self.time_layers[1](t_emb).unsqueeze(-1).unsqueeze(-1)

        e3 = self.enc3(F.max_pool2d(e2, 2))
        e3 = e3 + self.time_layers[2](t_emb).unsqueeze(-1).unsqueeze(-1)

        # 瓶颈层
        b = self.bottleneck(F.max_pool2d(e3, 2))
        b = b + self.time_layers[3](t_emb).unsqueeze(-1).unsqueeze(-1)

        # 解码
        d3 = self.dec3(torch.cat([F.interpolate(b, e3.shape[-2:]), e3], 1))
        d2 = self.dec2(torch.cat([F.interpolate(d3, e2.shape[-2:]), e2], 1))
        d1 = self.dec1(torch.cat([F.interpolate(d2, e1.shape[-2:]), e1], 1))

        return self.output(d1)

# 3. 训练循环
def train_diffusion_model(model, data_loader, num_epochs=100):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)

    diffusion = ForwardDiffusion()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(num_epochs):
        total_loss = 0
        for batch in data_loader:
            batch = batch.to(device)
            batch_size = batch.shape[0]

            # 随机时间步
            t = torch.randint(0, diffusion.num_steps, (batch_size,), device=device)

            # 添加噪声
            noisy_batch, noise = diffusion.add_noise(batch, t)

            # 预测噪声
            predicted_noise = model(noisy_batch, t)

            # 计算损失
            loss = F.mse_loss(predicted_noise, noise)

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

            total_loss += loss.item()

        if (epoch + 1) % 10 == 0:
            print(f'Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(data_loader):.4f}')

    return model

# 4. 采样函数
@torch.no_grad()
def sample(model, diffusion, num_samples=16, image_size=32, device='cuda'):
    model.eval()

    # 从纯噪声开始
    x = torch.randn(num_samples, 2, image_size, image_size).to(device)

    for t in reversed(range(diffusion.num_steps)):
        t_batch = torch.full((num_samples,), t, device=device)

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

        # 去噪步骤
        alpha = diffusion.alphas[t]
        alpha_bar = diffusion.alpha_bars[t]

        if t > 0:
            noise = torch.randn_like(x)
            sigma = torch.sqrt(diffusion.betas[t])
        else:
            noise = 0
            sigma = 0

        x = (1 / torch.sqrt(alpha)) * (x - (1 - alpha) / torch.sqrt(1 - alpha_bar) * predicted_noise)
        x = x + sigma * noise

    return x

# 5. 生成2D形状数据集
def generate_2d_shapes(num_samples=1000, image_size=32):
    """生成简单的2D形状（圆形、方形、三角形）"""
    data = []

    for _ in range(num_samples):
        img = np.zeros((2, image_size, image_size))
        shape_type = np.random.choice(['circle', 'square', 'triangle'])

        # 随机位置和大小
        center = np.random.randint(8, 24, 2)
        size = np.random.randint(4, 12)

        if shape_type == 'circle':
            # 生成圆形
            y, x = np.ogrid[:image_size, :image_size]
            mask = (x - center[0])**2 + (y - center[1])**2 <= size**2
            img[0, mask] = 1

        elif shape_type == 'square':
            # 生成方形
            x1, y1 = center - size
            x2, y2 = center + size
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(image_size, x2), min(image_size, y2)
            img[0, y1:y2, x1:x2] = 1

        else:  # triangle
            # 生成三角形
            pts = np.array([
                [center[0], center[1] - size],
                [center[0] - size, center[1] + size],
                [center[0] + size, center[1] + size]
            ])
            # 简化：使用方形近似
            mask = np.zeros((image_size, image_size))
            for i in range(image_size):
                for j in range(image_size):
                    if point_in_triangle([j, i], pts):
                        mask[i, j] = 1
            img[0] = mask

        # 添加边缘信息到第二个通道
        img[1] = edge_detect(img[0])

        data.append(torch.FloatTensor(img))

    return torch.stack(data)

def point_in_triangle(p, triangle):
    """检查点是否在三角形内"""
    def sign(p1, p2, p3):
        return (p1[0] - p3[0]) * (p2[1] - p3[1]) - (p2[0] - p3[0]) * (p1[1] - p3[1])

    d1 = sign(p, triangle[0], triangle[1])
    d2 = sign(p, triangle[1], triangle[2])
    d3 = sign(p, triangle[2], triangle[0])

    has_neg = (d1 < 0) or (d2 < 0) or (d3 < 0)
    has_pos = (d1 > 0) or (d2 > 0) or (d3 > 0)

    return not (has_neg and has_pos)

def edge_detect(img):
    """简单的边缘检测"""
    kernel = np.array([[-1, -1, -1],
                       [-1,  8, -1],
                       [-1, -1, -1]])
    from scipy.signal import convolve2d
    return convolve2d(img, kernel, mode='same', boundary='fill')

# 6. 可视化函数
def visualize_samples(samples, title="Generated Samples"):
    fig, axes = plt.subplots(4, 4, figsize=(10, 10))

    for i, ax in enumerate(axes.flat):
        if i < len(samples):
            # 显示第一个通道（形状）
            img = samples[i, 0].cpu().numpy()
            ax.imshow(img, cmap='gray')
            ax.axis('off')

    plt.suptitle(title)
    plt.tight_layout()
    plt.show()

# 7. 主训练脚本
if __name__ == "__main__":
    # 生成数据
    print("Generating dataset...")
    dataset = generate_2d_shapes(5000, 32)
    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=32, shuffle=True
    )

    # 创建模型
    model = SimpleUNet2D()

    # 训练
    print("Training model...")
    model = train_diffusion_model(model, data_loader, num_epochs=50)

    # 生成样本
    print("Generating samples...")
    diffusion = ForwardDiffusion()
    samples = sample(model, diffusion, num_samples=16)

    # 可视化
    visualize_samples(samples)

    # 保存模型
    torch.save(model.state_dict(), 'diffusion_2d.pth')
    print("Model saved!")

# 8. 条件生成扩展
class ConditionalUNet2D(SimpleUNet2D):
    def __init__(self, in_channels=2, time_dim=32, num_classes=3):
        super().__init__(in_channels, time_dim)

        # 类别嵌入
        self.class_emb = nn.Embedding(num_classes, time_dim)

    def forward(self, x, t, class_label=None):
        # 时间嵌入
        t_emb = self.time_mlp(t.float().unsqueeze(-1))

        # 添加类别条件
        if class_label is not None:
            class_emb = self.class_emb(class_label)
            t_emb = t_emb + class_emb

        # 继续原始forward流程...
        return super().forward(x, t)