$ show hpcaitech/Open-Sora
Open-Sora is a video generation framework implementing text-to-video, image-to-video, and video-to-video capabilities using diffusion models and variational autoencoders. The system architecture consists of modular components for models (MMDiT, DC-AE, Hunyuan VAE), datasets with bucket-based sampling, acceleration utilities for tensor/sequence parallelism, and training/inference scripts. Key entry points include training scripts in `scripts/diffusion/train/` and `scripts/vae/train.py`, inference scripts in `scripts/diffusion/inference.py` and `scripts/vae/inference.py`, and a Gradio demo at `gradio/app.py`. The system relies on PyTorch, ColossalAI for distributed training, xformers/flash-attention for efficient attention, and safetensors for model checkpointing. Data flows involve loading video datasets, applying spatiotemporal transforms, training diffusion models with classifier-free guidance, and generating videos through denoising loops.
The LLM identified 3 files worth highlighting for reuse:
Implements a flexible registry system for dynamically instantiating models, datasets, and solvers from configuration files. Enables plug-and-play component replacement without modifying core code, demonstrating the registry pattern for extensible frameworks. Minimal dependencies (only standard library and typing) make it highly reusable in other ML projects.
Provides a batch sampler that groups video data by resolution and frame count into buckets, enabling efficient training on heterogeneous video lengths. Implements dynamic batch sizing per bucket to maximize GPU utilization while maintaining aspect ratio integrity. The sampling logic is decoupled from dataset specifics, requiring only length/resolution metadata, making it adaptable to any video dataset.
Implements tensor parallelism for 1D/2D/3D convolutional layers and GroupNorm through column/row-wise sharding and all-reduce communications. Encapsulates distributed logic behind standard PyTorch layer interfaces, allowing seamless integration into existing models. Demonstrates advanced parallelism techniques while maintaining API compatibility with nn.Module, suitable for reuse in other distributed vision projects.