Implementation:Zai org CogVideo UNetModel
| Knowledge Sources | |
|---|---|
| Domains | Video_Generation, Diffusion_Models |
| Last Updated | 2026-02-10 00:00 GMT |
Overview
The primary UNet backbone for CogVideo's diffusion model, implementing a full encoder-decoder architecture with timestep embedding, cross-attention conditioning, skip connections, spatial transformers, and optional LoRA injection based on OpenAI's guided diffusion design.
Description
This module implements UNetModel, the core denoising neural network in CogVideo's diffusion pipeline. It is the most parameter-heavy component and is responsible for learning to predict noise (or the denoised signal) at each diffusion timestep.
The architecture is organized into three sections:
Input blocks (encoder path): A sequence of TimestepEmbedSequential modules, each containing ResBlock residual blocks with timestep embedding injection, optional SpatialTransformer blocks for cross-attention with text conditioning, and downsampling layers. The number of residual blocks per level is configurable per-level or globally. Downsampling can use either strided convolutions (Downsample) or residual blocks with built-in downsampling (when resblock_updown=True).
Middle block: A ResBlock, followed by either an AttentionBlock or SpatialTransformer, followed by another ResBlock. This processes features at the lowest resolution.
Output blocks (decoder path): Mirrors the input blocks with upsampling. Features from the encoder path are concatenated via skip connections before each block. Upsampling can use interpolation with convolution or residual blocks with built-in upsampling.
ResBlock implements timestep conditioning via two mechanisms: additive (default) where the timestep embedding is added to the hidden features, and scale-shift normalization (FiLM-like) where the embedding produces per-channel scale and shift parameters applied after GroupNorm. The block supports gradient checkpointing for memory efficiency and configurable kernel sizes.
TimestepEmbedSequential is a smart sequential container that routes inputs through different layer types appropriately: timestep blocks receive the timestep embedding, SpatialTransformers receive context for cross-attention, and SpatialVideoTransformers receive additional video-specific parameters.
The model supports class-conditional generation through label embeddings (discrete, continuous, timestep-based, or sequential), LoRA injection for efficient fine-tuning, and multiple attention variants including QKVAttention (split-then-heads) and QKVAttentionLegacy (heads-then-split).
Additional classes include EncoderUNetModel (encoder-only variant for classification), NoTimeUNetModel (zeroes out timestep input), and AttentionPool2d (CLIP-style attention pooling).
Usage
Use this model as the denoising backbone in a latent diffusion pipeline. It receives noisy latent tensors, timestep information, and optional text/class conditioning, and predicts the denoised output or noise estimate.
Code Reference
Source Location
- Repository: Zai_org_CogVideo
- File: sat/sgm/modules/diffusionmodules/openaimodel.py
- Lines: 1-1258
Signature
class UNetModel(nn.Module):
def __init__(
self,
in_channels,
model_channels,
out_channels,
num_res_blocks,
attention_resolutions,
dropout=0,
channel_mult=(1, 2, 4, 8),
conv_resample=True,
dims=2,
num_classes=None,
use_checkpoint=False,
use_fp16=False,
num_heads=-1,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=False,
resblock_updown=False,
use_new_attention_order=False,
use_spatial_transformer=False,
transformer_depth=1,
context_dim=None,
n_embed=None,
legacy=True,
disable_self_attentions=None,
num_attention_blocks=None,
disable_middle_self_attn=False,
use_linear_in_transformer=False,
spatial_transformer_attn_type="softmax",
adm_in_channels=None,
use_fairscale_checkpoint=False,
offload_to_cpu=False,
transformer_depth_middle=None,
dtype="fp32",
lora_init=False,
lora_rank=4,
lora_scale=1.0,
lora_weight_path=None,
):
Import
from sat.sgm.modules.diffusionmodules.openaimodel import UNetModel
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| in_channels | int | Yes | Channels in the input tensor |
| model_channels | int | Yes | Base channel count for the model |
| out_channels | int | Yes | Channels in the output tensor |
| num_res_blocks | int or list | Yes | Residual blocks per downsample level (int for global, list for per-level) |
| attention_resolutions | set/list/tuple | Yes | Downsample rates at which attention is applied |
| dropout | float | No | Dropout probability, default 0 |
| channel_mult | tuple of int | No | Channel multiplier per level, default (1, 2, 4, 8) |
| conv_resample | bool | No | Use learned convolutions for resampling, default True |
| dims | int | No | Signal dimensionality (1D, 2D, or 3D), default 2 |
| num_classes | int/str | No | Enables class-conditional generation. Supports int, "continuous", "timestep", "sequential" |
| use_checkpoint | bool | No | Enable gradient checkpointing, default False |
| num_heads | int | No | Attention heads, default -1 (use num_head_channels instead) |
| num_head_channels | int | No | Fixed channel width per attention head, default -1 |
| use_scale_shift_norm | bool | No | Use FiLM-style conditioning, default False |
| use_spatial_transformer | bool | No | Enable cross-attention with context, default False |
| transformer_depth | int or list | No | Depth of transformer blocks, default 1 |
| context_dim | int | No | Dimension of cross-attention context (required if use_spatial_transformer=True) |
| adm_in_channels | int | No | Input channels for sequential class embedding |
| dtype | str | No | Data type: "fp32", "fp16", or "bf16", default "fp32" |
| lora_init | bool | No | Initialize LoRA layers, default False |
| lora_rank | int | No | LoRA rank, default 4 |
Forward Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| x | torch.Tensor | Yes | Noisy input tensor (B, C, ...) |
| timesteps | torch.Tensor | Yes | 1-D batch of diffusion timesteps |
| context | torch.Tensor | No | Cross-attention conditioning (e.g. text embeddings) |
| y | torch.Tensor | No | Class labels (required if model is class-conditional) |
Outputs
| Name | Type | Description |
|---|---|---|
| output | torch.Tensor | Denoised prediction tensor of same spatial shape as input (B, out_channels, ...) |
Usage Examples
import torch
from sat.sgm.modules.diffusionmodules.openaimodel import UNetModel
model = UNetModel(
in_channels=4,
model_channels=320,
out_channels=4,
num_res_blocks=2,
attention_resolutions=[4, 2, 1],
channel_mult=(1, 2, 4, 4),
num_head_channels=64,
use_spatial_transformer=True,
transformer_depth=1,
context_dim=768,
use_checkpoint=True,
)
# Forward pass with noisy latent, timestep, and text context
x = torch.randn(2, 4, 64, 64)
t = torch.randint(0, 1000, (2,))
context = torch.randn(2, 77, 768) # text embeddings
output = model(x, timesteps=t, context=context)
# output shape: (2, 4, 64, 64)
# With LoRA fine-tuning
model_lora = UNetModel(
in_channels=4, model_channels=320, out_channels=4,
num_res_blocks=2, attention_resolutions=[4, 2, 1],
channel_mult=(1, 2, 4, 4), num_head_channels=64,
lora_init=True, lora_rank=4, lora_scale=1.0,
)