Implementation:Zai org CogVideo CLI VAE Demo
| Knowledge Sources | |
|---|---|
| Domains | Video_Generation, Variational_Autoencoders |
| Last Updated | 2026-02-10 00:00 GMT |
Overview
A command-line tool for encoding video frames into latent representations and decoding latents back into video using the CogVideoX VAE, independent of the full diffusion pipeline.
Description
The CLI VAE Demo provides three operational modes for interacting with the CogVideoX Variational Autoencoder. In encode mode, it reads video frames using imageio, stacks them into a tensor of shape (1, C, T, H, W), and passes them through AutoencoderKLCogVideoX.encode() to produce latent tensors saved as .pt files. In decode mode, it loads a previously saved latent tensor and reconstructs video frames via model.decode(). In both mode, it runs encode followed by decode sequentially. The script enables VAE slicing and tiling for memory efficiency, allowing encoding on GPUs with approximately 18GB of memory and decoding with approximately 4GB.
Usage
Use this tool when you need to inspect or experiment with the CogVideoX VAE in isolation, debug encode/decode round-trips, measure latent space quality, or pre-compute encoded tensors for downstream tasks without running the full diffusion pipeline.
Code Reference
Source Location
- Repository: Zai_org_CogVideo
- File: inference/cli_vae_demo.py
Signature
def encode_video(model_path: str, video_path: str, dtype: torch.dtype, device: str) -> torch.Tensor
def decode_video(model_path: str, encoded_tensor_path: str, dtype: torch.dtype, device: str) -> torch.Tensor
def save_video(tensor: torch.Tensor, output_path: str) -> None
Import
from inference.cli_vae_demo import encode_video, decode_video, save_video
I/O Contract
Inputs
encode_video:
| Name | Type | Required | Description |
|---|---|---|---|
| model_path | str | Yes | Path to the pretrained CogVideoX VAE model directory |
| video_path | str | Yes | Path to the input video file (readable by imageio/ffmpeg) |
| dtype | torch.dtype | Yes | Computation data type (e.g., torch.float16 or torch.bfloat16) |
| device | str | Yes | Target device for computation (e.g., "cuda" or "cpu") |
decode_video:
| Name | Type | Required | Description |
|---|---|---|---|
| model_path | str | Yes | Path to the pretrained CogVideoX VAE model directory |
| encoded_tensor_path | str | Yes | Path to a saved encoded latent tensor file (.pt) |
| dtype | torch.dtype | Yes | Computation data type |
| device | str | Yes | Target device for computation |
save_video:
| Name | Type | Required | Description |
|---|---|---|---|
| tensor | torch.Tensor | Yes | Decoded video frames tensor of shape (1, C, T, H, W) |
| output_path | str | Yes | Directory path where output.mp4 will be saved |
Outputs
| Name | Type | Description |
|---|---|---|
| encoded_frames | torch.Tensor | Latent representation from the VAE encoder (encode_video) |
| decoded_frames | torch.Tensor | Reconstructed video frames from the VAE decoder (decode_video) |
| output.mp4 | File | Saved video file at 8 FPS in the specified output directory (save_video) |
Usage Examples
# Encode a video to latent representation
python cli_vae_demo.py \
--model_path /path/to/CogVideoX-2b/vae/ \
--video_path ../resources/videos/1.mp4 \
--mode encode
# Decode a latent tensor back to video
python cli_vae_demo.py \
--model_path /path/to/CogVideoX-2b/vae/ \
--encoded_path ./encoded.pt \
--mode decode
# Encode and decode in sequence
python cli_vae_demo.py \
--model_path /path/to/CogVideoX-2b/vae/ \
--video_path ../resources/videos/1.mp4 \
--mode both
import torch
from inference.cli_vae_demo import encode_video, decode_video, save_video
# Encode video to latent space
encoded = encode_video(
model_path="/path/to/CogVideoX-2b/vae/",
video_path="input.mp4",
dtype=torch.bfloat16,
device="cuda"
)
torch.save(encoded, "encoded.pt")
# Decode latent back to video
decoded = decode_video(
model_path="/path/to/CogVideoX-2b/vae/",
encoded_tensor_path="encoded.pt",
dtype=torch.bfloat16,
device="cuda"
)
save_video(decoded, "./output")