Implementation:Zai org CogVideo MoVQ Decoder3D Dev

Knowledge Sources	Zai_org_CogVideo
Domains	Video_Generation, Autoencoding, 3D_Decoding
Last Updated	2026-02-10 00:00 GMT

Overview

MoVQ_Decoder3D_Dev is the development variant of the 3D MoVQ decoder that adds beartype runtime type checking and preserves commented-out implementation history documenting the transition from standard 3D convolutions to CausalConv3d operations.

Description

This module is structurally identical to the production movq_dec_3d.py decoder, implementing the same MOVQDecoder3D and NewDecoder3D classes with the same quantized-code-conditioned spatial normalization, causal 3D convolution architecture, and multi-resolution temporal upsampling. The key differences that distinguish this development variant are:

Runtime type checking: Imports beartype and its typing module (Union, Tuple, Optional, List) for runtime type validation during development and testing. This provides additional safety guarantees when experimenting with new configurations.

Implementation history: Contains extensive commented-out code blocks showing the original nn.Conv3d implementations that were iteratively replaced with CausalConv3d. These comments document the design evolution and serve as reference for understanding why causal convolutions were chosen over standard ones.

Explicit F import: Imports torch.nn.functional as F explicitly, whereas the production version does not include this import (though both use torch.nn.functional through other paths).

The SpatialNorm3D class in this variant includes commented-out lines showing the original nn.Conv3d-based implementation for the conditioning convolutions (conv, conv_y, conv_b) alongside the active CausalConv3d replacements. Similarly, ResnetBlock3D preserves commented nn.Conv3d code for conv1, conv2, conv_shortcut, and the conv_in/conv_out in the decoder classes.

All functional behavior -- the forward pass, normalization conditioning, temporal upsampling logic, and the get_last_layer method -- is identical to the production version.

Usage

Use MoVQ_Decoder3D_Dev during development, experimentation, and testing phases when runtime type validation is desired and when the commented implementation history is useful as a reference. For production deployments, prefer the standard MoVQ_Decoder3D from movq_dec_3d.py, which has a cleaner codebase without the development annotations.

Code Reference

Source Location

Repository: Zai_org_CogVideo
File: sat/sgm/modules/autoencoding/vqvae/movq_dec_3d_dev.py
Lines: 229-384 (MOVQDecoder3D), 387-549 (NewDecoder3D), 53-92 (SpatialNorm3D), 108-180 (ResnetBlock3D), 183-226 (AttnBlock2D)

Signature

class MOVQDecoder3D(nn.Module):
    def __init__(
        self,
        *,
        ch,
        out_ch,
        ch_mult=(1, 2, 4, 8),
        num_res_blocks,
        attn_resolutions,
        dropout=0.0,
        resamp_with_conv=True,
        in_channels,
        resolution,
        z_channels,
        give_pre_end=False,
        zq_ch=None,
        add_conv=False,
        pad_mode="first",
        temporal_compress_times=4,
        **ignorekwargs,
    ):
        ...

    def forward(self, z, use_cp=False):
        ...

    def get_last_layer(self):
        ...

class NewDecoder3D(nn.Module):
    def __init__(
        self,
        *,
        ch,
        out_ch,
        ch_mult=(1, 2, 4, 8),
        num_res_blocks,
        attn_resolutions,
        dropout=0.0,
        resamp_with_conv=True,
        in_channels,
        resolution,
        z_channels,
        give_pre_end=False,
        zq_ch=None,
        add_conv=False,
        pad_mode="first",
        temporal_compress_times=4,
        post_quant_conv=False,
        **ignorekwargs,
    ):
        ...

    def forward(self, z):
        ...

Import

from sat.sgm.modules.autoencoding.vqvae.movq_dec_3d_dev import MOVQDecoder3D, NewDecoder3D, SpatialNorm3D

I/O Contract

MOVQDecoder3D Constructor Inputs

Name	Type	Required	Description
ch	int	Yes	Base channel count for the decoder
out_ch	int	Yes	Number of output channels (e.g., 3 for RGB)
ch_mult	tuple	No (default (1,2,4,8))	Channel multipliers per resolution level
num_res_blocks	int	Yes	Number of residual blocks per resolution level
attn_resolutions	list	Yes	Resolutions at which to apply spatial attention
dropout	float	No (default 0.0)	Dropout rate in residual blocks
resamp_with_conv	bool	No (default True)	Whether to use convolution-based upsampling
in_channels	int	Yes	Number of input channels to the model
resolution	int	Yes	Target spatial resolution of the output
z_channels	int	Yes	Number of channels in the latent representation
give_pre_end	bool	No (default False)	If True, return features before final normalization and convolution
zq_ch	int	No (default None)	Channel count for the quantized conditioning tensor; defaults to z_channels
add_conv	bool	No (default False)	Whether to add an extra convolution in SpatialNorm3D for zq conditioning
pad_mode	str	No (default "first")	Padding mode for CausalConv3d: "first" or "constant"
temporal_compress_times	int	No (default 4)	Temporal compression factor; log2 determines the number of temporal upsampling levels

NewDecoder3D Additional Inputs

Name	Type	Required	Description
post_quant_conv	bool	No (default False)	Whether to apply a CausalConv3d layer to the latent before decoding

Forward Inputs

Name	Type	Required	Description
z	torch.Tensor	Yes	Quantized latent tensor of shape (B, C, T, H, W)
use_cp	bool	No (default False)	Checkpoint flag (MOVQDecoder3D only)

Outputs

Name	Type	Description
output	torch.Tensor	Reconstructed video tensor of shape (B, out_ch, T_out, H_out, W_out)
get_last_layer	torch.Tensor	Returns self.conv_out.conv.weight for discriminator-based training losses

Differences from Production Version

Feature	movq_dec_3d.py (Production)	movq_dec_3d_dev.py (Development)
beartype import	No	Yes (beartype, Union, Tuple, Optional, List)
torch.nn.functional import	No explicit F import	Yes (import torch.nn.functional as F)
Commented-out Conv3d code	Minimal comments	Extensive commented Conv3d originals preserved
Functional behavior	Identical	Identical
Line count	505	549

Usage Examples

from sat.sgm.modules.autoencoding.vqvae.movq_dec_3d_dev import MOVQDecoder3D

# Create a development 3D MoVQ decoder (identical API to production)
decoder = MOVQDecoder3D(
    ch=128,
    out_ch=3,
    ch_mult=(1, 2, 4, 4),
    num_res_blocks=2,
    attn_resolutions=[16],
    dropout=0.0,
    in_channels=256,
    resolution=256,
    z_channels=256,
    zq_ch=256,
    add_conv=False,
    pad_mode="first",
    temporal_compress_times=4,
)

# Decode from a quantized latent (B=2, C=256, T=4, H=16, W=16)
z_q = torch.randn(2, 256, 4, 16, 16)
video = decoder(z_q)

# NewDecoder3D with post-quantization convolution
from sat.sgm.modules.autoencoding.vqvae.movq_dec_3d_dev import NewDecoder3D

decoder_v2 = NewDecoder3D(
    ch=128,
    out_ch=3,
    ch_mult=(1, 2, 4, 4),
    num_res_blocks=2,
    attn_resolutions=[16],
    in_channels=256,
    resolution=256,
    z_channels=256,
    post_quant_conv=True,
)
video = decoder_v2(z_q)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment