Cosmos3架构可视化工具

Cosmos3架构可视化工具——展示Latent World Model和DiT生成原理文章目录Cosmos3架构可视化工具前言一、Cosmos3-mini的基本程序1.__init__.py2.config.py3.model.py4.vae.py5.diffusion.py6.transformer.py二、enviroment和运行说明三、代码总结前言2026年5月Nvidia推出开源世界模型——Cosmos3Cosmos3 是一套全模态世界模型能够根据文本、图像、视频和动作轨迹输入的组合生成动态、高质量的视频、图像、音频和动作指令。涵盖世界理解、世界生成、模拟和具身策略学习等领域。其网络架构采用混合 Transformer (MoT)该架构由两个互补的 Transformer 塔组成一个用于离散token生成的自回归 Transformer 和一个用于连续多模态生成的扩散 Transformer。在推理过程中文本通过标准的下一token自回归解码生成而非文本模态例如图像、视频、音频和动作则通过迭代去噪合成。这种统一的架构使 Cosmos3 能够在单一框架内对异构模态进行建模同时保留最适合每种模态的生成机制。“本文用华为云码道智能体——Codearts生成一个展示Cosmos3架构的可视化页面Cosmos3-mini其中只展示Cosmos3 非文本模态通过迭代去噪合成也就是从噪声生成视频的推理过程即只包含用于连续多模态生成的扩散 Transformer。”一、Cosmos3-mini的基本程序将VAE DiT Diffusion 组装为端到端程序数据流如下Pixel Space → [VAE Encoder] → Latent Space → [Rectified Flow DiT] → Latent Space → [VAE Decoder] → Pixel Space1.init.py负责在程序中导入配置类编码器和解码器 DiT和Transformer模块调度器和整流器以及端到端模型定义了可以从Cosmos3_mini导出的8个类。2.config.py定义ToyCosmos3Config配置类用一个 class 集中管理所有超参数各模块后续可通过 config.xxx读取VAE压缩后的潜空间通道数16VAE中间层的通道数32时间维度总压缩比48→2空间维度总压缩比1664×64→4×4DiTTransformer 隐藏层维度128Transformer 层数2注意力头数4每个头的维度32hidden_size4头×32128MoT的模态数2Diffusion推理时只跑 2 步去噪mRoPE编码位置信息tuple (12, 10, 10)拆成3段分给 T/H/W3.model.py定义ToyCosmos3Model把 VAE 编码器、DiT 去噪网络、扩散调度器拼成一条流水线准备 def __init__(self,configNone)def encode(self, video)def diffuse(self, latent_shape)def decode(self, z)前向传播 def forward(self, video)zself.encode(video)Step1:用 VAE 把视频压缩到潜在空间 z_genself.diffuse(z.shape)Step2: 扩散 video_outself.decode(z_gen)Step3: 用 VAE 把去噪后的 latent 解码回像素空间def count_parameters(self)统计模型参数量4.vae.py输入数据格式为B,3,T,H,W含义batchsize这一批里有几段视频、通道数、时间维度即视频有多少帧、高度即每帧画面的垂直像素数、宽度即每帧画面的水平像素数实现了视频的压缩-解压压缩Encoder 用 4 级 3D 下采样把像素视频压成 16 通道的 latentDecoder则镜像还原成像素视频Encoder网络结构: Norm归一化→ SiLU用sigmoid激活→ Conv → Norm → SiLU → Conv shortcut残差连接一、准备class ToyCausalConv3d(nn.Module) 将Cosmos3中的因果 3D 卷积简化为普通卷积完成初始化和前向卷积 class ToyResBlock3d(nn.Module) 标准残差连接[如果输入输出通道不同用1*1卷积对齐]前向传播过程 Norm → SiLU → Conv → Norm → SiLU → Conv 包含了归一化卷积×2 class ToyDownsample3d(nn.Module): 下采样 class ToyUpsample3d(nn.Module): 上采样二、调用class ToyVAEEncoder(nn.Module):像素空间视频编码到潜在空间数据流 (Pixel Space → Latent Space):输入: (B, 3, T, H, W) — 像素空间输出: (B, 16, T/4, H/16, W/16) — 潜在空间具体过程1.初始化隐藏层潜空间残差卷积时间空间下采样空间下采样×3中间层2.前向传播3.VAE采样class ToyVAEDecoder(nn.Module):将潜在空间解码回像素空间具体过程1.初始化隐藏层潜空间残差卷积中间层空间上采样×3时间空间上采样2.前向传播三、Web UI 可视化说明统计给张量做体检确认数据分布正常所有元素的平均大小、元素的分散程度、最小值和最大值图示为压缩后潜在空间的一个切片就是高维张量的一个横截面 热力图就是直观看到模型在潜在空间里看到了什么。 切片切了一个通道比如轮廓特征。 由于是随机初始化权重现在看到的基本是杂色——如果模型训练好了就会出现有意义的结构。 颜色相近的区域 → 这些位置的特征相似比如可能是同一个物体 颜色突变的边界 → 可能是物体的边缘 整体偏蓝/偏红 → 这个通道在抑制/激活某些特征5.diffusion.py一、准备class ToyFlowScheduler(nn.Module): 初始化 设置推理时使用的时间步 添加噪声[Rectified Flow 的加噪公式:z_t(1- t)× z_0 t × noise]单步去噪[Euler方法 z_{t-dt}z_t - dt ×v]二、调用class ToyRectifiedFlow(nn.Module):完整的扩散推理流程组装调度器和 DiT Transformer用到上面定义的函数纯高斯噪声设置时间步逐步去噪得到干净的latent从纯高斯噪声出发通过 Rectified Flow 多步去噪 逐步雕刻出符合物理规律的潜在空间表示。每一步去噪中Transformer 预测一个速度场 v(z_t, t) 指引噪声向干净数据方向移动。这就像物理中的退火过程 从混沌 (纯噪声) 到有序 (物理一致的表示)。diffusion.py 的ToyRectifiedFlow.forward() 在去噪中调用了transformer.py得到预测的速度场。三、Web UI 可视化说明去噪统计给去噪后的 latent 张量做体检确认数据分布正常 特征值散点对比切片4*4扩散去噪前后值的对比 特征差异热力图对应于散点图中的点扩散去噪前后值的差 重点展示了用ToyRectifiedFlow.forward()去噪过程 生成了纯噪声均值, 标准差, 范围 逐步去噪 DiT 预测速度场: v_predself.transformer(z_t, t_batch)去噪时当shift1时t_shifted 大于 t把时间步分布往高噪声端挤得到高分辨率的输出 Euler 更新: 得到干净的latent但是由于预测速度场是随意进行的并没有用真实视频训练在本toy的推理过程中这个干净latent即z_0是没有意义的本扩散程序只有推理 简单介绍一下扩散程序中未包含的训练过程1. 取一张真实图片 → 编码器 → 得到干净 z_02. 随机采样 t ∈[0,1]3. 随机采样噪声 ε4. 构造 z_t(1-t)*z_0 t*ε ← 这就是 add_noise 做的事5. 计算真实速度 v_truez_0 - ε ← 从 z_t 到 z_0 的方向6. 让 DiT 预测 v_predDiT(z_t, t)7. 算损失 loss||v_pred - v_true||²8. 反向传播更新权重6.transformer.py一、准备class ToyMRoPE(nn.Module): 多维旋转位置编码是给每个位置贴3D坐标标签head_dim32被切成3段时间12维高度10维宽度10维 class ToyAttention(nn.Module): 多头注意力机制看全局关系谁和谁有关联 class ToyFeedForward(nn.Module): SwiGLU 激活逐个加工每个Token前馈网络 class ToyAdaLNModulation(nn.Module): Adaptive Layer Norm自适应归一化用timestep嵌入生成6个调制参数(shift1/scale1/gate1/shift2/scale2/gate2)来调制归一化后的特征 class ToyDiTBlock(nn.Module): DiT基本块结构为AdaLN→Attention→Gate残差→AdaLN→FFN→Gate残差 class ToyTimestepEmbedding(nn.Module): 时间步嵌入将标量timestep通过正弦位置编码MLP映射为隐藏层维度的向量 class ToyPatchEmbed3d(nn.Module): 用3D卷积将VAE latent切成空间patch并投影为token序列二、调用class ToyDiTTransformer(nn.Module):完整 DiT 骨干输入: VAE latent timestep输出: 去噪后的 latent 预测具体过程1.一段加噪的视频latent进来2.patch_embed“切碎”—— 用 3D 卷积把VAE latent 切成2×2小块转为token序列用kernel大小来实现3.t_embed“报时”——把标量时间步 t 编码为128维向量告诉模型现在是第几步去噪4.ToyDiTBlock×2“思考”一个Transformer层AdaLN听时间 → Attention看关系 → FFN做变换——重复两遍理解视频内容5.Unpatchify“拼回去”将预测的 patch 值重新排列回 5D 张量将 patch 内的 2×2 空间维度和 patch 网格维度正确交错恢复为 (1, 16, 2, 4, 4) 的完整 latent6.输出预测的速度场v其中前向传播过程patch_embedt_embedbuild_position_ids2 层 DiT Block 依次处理每层都注入时间步嵌入和 3D 位置编码调用关系如下图三、Web UI 可视化说明上文描述的ToyDiTTransformer及其它类如ToyAttention、ToyFeedForward在配套的 Web UI 中进行了可视化展示帮助直观理解计算过程。重点展示了ToyDiTTransformer.forward()计算过程 输入[1,16,4,4,4]每个2×2 块有16×1×2×264个值 输出[1,128,4,2,2]用kernel大小来实现patch_embed“切块”再由64维用线性变换“投影”到128维(1,4×2×2,128)(1,16,128)展平成token序列128维16个 3D位置编码16个token来自T4×pH2×pW2共128维按时间帧分组4个空间token共享同一个时间位置 AdaLN1自适应层归一化层的6个调制参数 shift1/scale1/gate1 和 shift2/scale2/gate2 是 AdaLN 中进行特征变换平移特征分布、缩放拉伸压缩特征幅度、门控控制注意力信息注入量的参数可以根据绝对值判断哪些参数更活跃 具体流程 调制正则化、平移、缩放用到shift1/scale1这两个参数与扩散时间有关 计算注意力 注意力注入注入attention_out前的x的均值和注入后的均值用到gate1 调制用到shift2/scale2)FFN用到gate2)多头注意力Attention和前馈网络FFN的实现流程图如下attention的实现过程 x_norm(16tokens ×128hidden)↓ ┌─────────────────────────────────────┐ │1. Q/K/V 投影 │ │ Qx_norm W_Q →(B, 4heads,16,32)│ │ Kx_norm W_K →(B, 4heads,16,32)│ │ Vx_norm W_V →(B, 4heads,16,32)│ ├─────────────────────────────────────┤ │2. QK Norm(RMSNorm)│ │ QRMSNorm(Q)│ │ KRMSNorm(K)│ ├─────────────────────────────────────┤ │3. mRoPE 旋转位置编码 │ │ Q, K 按(T,H,W)三轴施加旋转编码 │ │head_dim32分为(12,10,10)三段 │ ├─────────────────────────────────────┤ │4. 注意力计算 │ │ attnsoftmax(Q K^T / √32)│ │ →(B, 4heads,16,16)注意力权重 │ ├─────────────────────────────────────┤ │5. 加权求和 │ │ outattn V │ │ →(B, 4heads,16,32)│ ├─────────────────────────────────────┤ │6. 输出投影 │ │ outout W_O →(B,16,128)│ └─────────────────────────────────────┘ ↓ attn_out热力图(16tokens ×128hidden)FFN的实现过程 x_norm2(16tokens ×128hidden)↓ ┌──────────────────────────────────────────────┐ │1. gate_proj: Linear(128→256)│ │ gatex_norm2 W_gate │ │ →(16,256)│ ├──────────────────────────────────────────────┤ │2. SwiGLU 激活 │ │ silu(gate)gate * σ(gate)│ ├──────────────────────────────────────────────┤ │3. up_proj: Linear(128→256)│ │ upx_norm2 W_up │ │ →(16,256)│ ├──────────────────────────────────────────────┤ │4. 逐元素相乘 │ │ midsilu(gate)* up │ │ →(16,256)│ ├──────────────────────────────────────────────┤ │5. down_proj: Linear(256→128)│ │ outmid W_down │ │ →(16,128)← 热力图 │ └──────────────────────────────────────────────┘残差连接二、enviroment和运行说明环境要求Python3.9 PyTorch2.0.0 Gradio4.0.0 Matplotlib3.7.0 NumPy1.24.0启动cd D:\huaweicloud\mini_cosmos\demo\UI_show换成自己的UI_show保存路径 py start.py启动后浏览器访问 http://127.0.0.1:7860/三、代码完整代码在这一篇代码的单独博文总结有些细节我也不太明白重点关注了多头注意力具体怎样实现数据流以及DiT架构中扩散和transformer的调用之后会再来修改和补充。欢迎大家来交流讨论。