（源码）Sora技术深度解析：从理论到落地的视频生成实践

qq_16226939

668人浏览 · 2025-04-25 00:49:26

qq_16226939 · 2025-04-25 00:49:26 发布

本文是关于另一篇文章Sora技术深度解析：从理论到落地的视频生成实践-CSDN博客的后续，下面是一个完整的、可执行的简化版Sora模型实现，包含训练、保存和加载模型的完整流程。这个实现基于PyTorch框架，包含了视频分词、扩散Transformer和训练流程等核心组件：

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from einops import rearrange, repeat
import numpy as np
from tqdm import tqdm
import os

# 1. 视频分词器 (ViViT)
class VideoTokenizer(nn.Module):
    def __init__(self, dim=128, patch_size=8, temporal_patch_size=2):
        super().__init__()
        self.spatial_patch = nn.Conv2d(3, dim, kernel_size=patch_size, stride=patch_size)
        self.temporal_patch = nn.Conv3d(dim, dim, kernel_size=(temporal_patch_size, 1, 1), 
                                     stride=(temporal_patch_size, 1, 1))
        self.patch_size = patch_size
        self.temporal_size = temporal_patch_size
        
    def forward(self, x):
        # x: (batch, frames, C, H, W)
        b, t, c, h, w = x.shape
        
        # 空间分块
        x = rearrange(x, 'b t c h w -> (b t) c h w')
        x = self.spatial_patch(x)  # (b*t, dim, h', w')
        h_p, w_p = x.shape[-2:]
        x = rearrange(x, '(b t) d h w -> b t d h w', b=b)
        
        # 时间分块
        x = rearrange(x, 'b t d h w -> b d t h w')
        x = self.temporal_patch(x)  # (b, dim, t', h, w)
        x = rearrange(x, 'b d t h w -> b (t h w) d')
        
        return x, (h_p, w_p)
    
    def reconstruct(self, x, hw_shape):
        # 从token重建视频
        h_p, w_p = hw_shape
        b, n, d = x.shape
        t_p = n // (h_p * w_p)
        
        x = rearrange(x, 'b (t h w) d -> b d t h w', h=h_p, w=w_p)
        x = F.interpolate(x, scale_factor=(self.temporal_size, 1, 1), mode='nearest')
        x = rearrange(x, 'b d t h w -> (b t) d h w')
        x = F.interpolate(x, scale_factor=self.patch_size, mode='nearest')
        x = rearrange(x, '(b t) d h w -> b t d h w', b=b)
        return x

# 2. 扩散Transformer
class DiffusionTransformer(nn.Module):
    def __init__(self, dim=128, depth=4, heads=4):
        super().__init__()
        self.time_embed = nn.Sequential(
            nn.Linear(1, dim*4),
            nn.SiLU(),
            nn.Linear(dim*4, dim)
        )
        
        self.blocks = nn.ModuleList([
            TransformerBlock(dim, heads) for _ in range(depth)
        ])
        
        self.norm = nn.LayerNorm(dim)
        
    def forward(self, x, t):
        # x: (batch, seq_len, dim)
        # t: (batch,)
        
        temb = self.time_embed(t.float().unsqueeze(1))
        x = x + temb.unsqueeze(1)
        
        for block in self.blocks:
            x = block(x)
            
        return self.norm(x)

class TransformerBlock(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.attn = SpatioTemporalAttention(dim, heads)
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim*4),
            nn.GELU(),
            nn.Linear(dim*4, dim)
        )
        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        
    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x

class SpatioTemporalAttention(nn.Module):
    def __init__(self, dim, heads):
        super().__init__()
        self.heads = heads
        self.scale = (dim // heads) ** -0.5
        self.to_qkv = nn.Linear(dim, dim*3)
        self.to_out = nn.Linear(dim, dim)
        
    def forward(self, x):
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h=self.heads), qkv)
        
        dots = torch.einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
        attn = dots.softmax(dim=-1)
        
        out = torch.einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        return self.to_out(out)

# 3. 完整MiniSora模型
class MiniSora(nn.Module):
    def __init__(self, dim=128, patch_size=8, temporal_patch_size=2):
        super().__init__()
        self.tokenizer = VideoTokenizer(dim, patch_size, temporal_patch_size)
        self.diffusion = DiffusionTransformer(dim)
        self.to_pixels = nn.Linear(dim, patch_size**2 * 3)
        self.patch_size = patch_size
        
    def forward(self, x, t):
        tokens, hw_shape = self.tokenizer(x)
        tokens = self.diffusion(tokens, t)
        pred = self.to_pixels(tokens)
        return pred, hw_shape
    
    def loss(self, x, noise, t):
        pred, hw_shape = self(x, t)
        target = rearrange(noise, 'b t c (h p1) (w p2) -> b (t h w) (p1 p2 c)', 
                         p1=self.patch_size, p2=self.patch_size)
        return F.mse_loss(pred, target)

# 4. 训练工具函数
def linear_beta_schedule(timesteps):
    scale = 1000 / timesteps
    beta_start = scale * 0.0001
    beta_end = scale * 0.02
    return torch.linspace(beta_start, beta_end, timesteps)

def extract(a, t, x_shape):
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))

# 5. 简单数据集
class VideoDataset(Dataset):
    def __init__(self, num_videos=100, frames=8, size=64):
        # 生成随机视频数据作为示例
        self.data = torch.randn(num_videos, frames, 3, size, size)
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]

# 6. 训练函数（包含模型保存）
def train_model():
    # 参数设置
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    batch_size = 8
    timesteps = 1000
    epochs = 10  # 示例中减少epoch数以便快速运行
    
    # 初始化模型和优化器
    model = MiniSora().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
    
    # 噪声调度
    betas = linear_beta_schedule(timesteps)
    alphas = 1. - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    
    # 数据集和数据加载器
    dataset = VideoDataset()
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
    
    # 训练循环
    for epoch in range(epochs):
        pbar = tqdm(dataloader)
        for batch in pbar:
            batch = batch.to(device)
            
            # 采样时间步和噪声
            t = torch.randint(0, timesteps, (batch.shape[0],), device=device)
            noise = torch.randn_like(batch)
            
            # 添加噪声
            sqrt_alphas_cumprod_t = extract(alphas_cumprod.sqrt(), t, batch.shape)
            sqrt_one_minus_alphas_cumprod_t = extract((1. - alphas_cumprod).sqrt(), t, batch.shape)
            noisy = sqrt_alphas_cumprod_t * batch + sqrt_one_minus_alphas_cumprod_t * noise
            
            # 计算损失
            loss = model.loss(noisy, noise, t)
            
            # 反向传播
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            pbar.set_description(f'Epoch {epoch} Loss: {loss.item():.4f}')
    
    # 训练完成后保存模型
    torch.save(model.state_dict(), 'minisora.pth')
    print("模型已保存为 minisora.pth")
    return model

# 7. 采样生成函数
@torch.no_grad()
def generate_video(model, shape=(1, 8, 3, 64, 64), timesteps=1000):
    device = next(model.parameters()).device
    b, t, c, h, w = shape
    
    # 初始化随机噪声
    x = torch.randn(shape, device=device)
    
    # 噪声调度参数
    betas = linear_beta_schedule(timesteps)
    alphas = 1. - betas
    alphas_cumprod = torch.cumprod(alphas, dim=0)
    
    # 逐步去噪
    for i in reversed(range(timesteps)):
        time = torch.full((b,), i, device=device, dtype=torch.long)
        
        # 预测噪声
        pred, hw_shape = model(x, time)
        pred_noise = rearrange(pred, 'b (t h w) (p1 p2 c) -> b t c (h p1) (w p2)',
                             p1=model.patch_size, p2=model.patch_size,
                             h=hw_shape[0], w=hw_shape[1], t=t)
        
        # 更新x
        sqrt_recip_alphas = torch.sqrt(1.0 / alphas[i])
        sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod[i])
        x = sqrt_recip_alphas * (x - pred_noise * (1 - alphas[i]) / sqrt_one_minus_alphas_cumprod)
        
        if i > 0:
            noise = torch.randn_like(x)
            x += torch.sqrt(betas[i]) * noise
            
    return x

# 8. 主程序
def main():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"使用设备: {device}")
    
    # 检查是否有已保存的模型
    if os.path.exists('minisora.pth'):
        print("发现已保存的模型，加载中...")
        model = MiniSora().to(device)
        model.load_state_dict(torch.load('minisora.pth', map_location=device))
        print("模型加载成功！")
    else:
        print("未找到已保存的模型，开始训练...")
        model = train_model()
    
    # 生成示例视频
    print("生成视频中...")
    video = generate_video(model)
    print(f"生成的视频形状: {video.shape}")  # (1, 8, 3, 64, 64)

if __name__ == '__main__':
    main()

代码使用说明

首次运行：
- 代码会自动检测没有模型文件(minisora.pth)
- 开始训练模型（示例中设置为10个epoch以便快速运行）
- 训练完成后自动保存模型
后续运行：
- 检测到已有模型文件后自动加载
- 跳过训练直接生成视频
关键参数调整：
- 修改MiniSora类的参数可以改变模型结构
- 修改train_model()中的epochs可以改变训练轮数
- 修改generate_video()中的shape可以改变生成视频的尺寸
输出说明：
- 生成的视频是一个5维张量：(batch, frames, channels, height, width)
- 示例输出为(1, 8, 3, 64, 64)，表示1个视频，8帧，3通道颜色，64x64分辨率