$ show NVlabs/Sana
This repository is the official implementation of the SANA family of diffusion models for high-resolution image and video generation, developed by NVIDIA Research. It is primarily a research codebase, containing training and inference pipelines for multiple model variants: SANA (text-to-image), SANA-1.5, SANA-Sprint (few-step distillation), SANA-Video, SANA-Video 2.0, SANA-Streaming (long video), LongSANA, and SANA-WM (world model for camera-controlled video). The codebase also includes post-training with reinforcement learning (Sol-RL), ControlNet support, and various schedulers (flow matching, DPM-Solver, LCM, self-forcing). While the core purpose is research, the repository also provides deployable Gradio applications, conversion scripts to Hugging Face diffusers format, and integration with ComfyUI and SGLang, giving it a hybrid application character. The code is heavily optimized with custom Triton kernels for efficient attention (LiteMLA, Gated Delta Net) and includes distributed training utilities (FSDP, context parallelism). The repository is actively maintained with frequent commits and supports a wide range of model sizes (600M to 14B parameters).
The LLM identified 5 files worth highlighting for reuse:
Provides a generic shard-based dataset implementation (ShardListDataset, IndexedTarSamples) with JSON metadata, MD5 verification, LRU caching, and distributed samplers (ChunkedSampler, DistributedChunkedSampler). This module is self-contained and can be extracted for any large-scale dataset loading pipeline that uses tar archives, supporting remote URLs and locking. It demonstrates non-obvious data engineering skill in handling sharded data with efficient caching and distributed access.
Implements Triton-optimized linear attention (LiteMLA) with fused kernels for forward and backward passes. This module is a drop-in replacement for standard attention in transformer models, offering significant speedups. It is well-isolated and can be reused in any PyTorch model that requires efficient linear attention, demonstrating high craftsmanship in custom CUDA/Triton kernel development.
Provides a JIT CUDA extension for chunkwise Gated Delta Net (GDN) attention, with automatic architecture detection and fallback to Triton/PyTorch. This module enables efficient autoregressive video generation by processing attention in chunks with state caching. It is a self-contained, reusable component for any model using GDN attention, demonstrating advanced CUDA and autograd Function implementation.
Contains utility functions for temporal chunking in causal video models, including chunk index calculation, normalization, and pattern analysis. These functions are generic and can be reused in any video generation or processing pipeline that requires splitting sequences into chunks. The module is small, well-documented, and demonstrates a clean pattern for handling variable-length video data.
Implements a factory of reward scorers (CLIP, HPSv2, PickScore, ImageReward) for reinforcement learning fine-tuning of diffusion models. Each scorer is a callable class that can be instantiated with a model name and used independently. This module is self-contained and can be extracted for any RL-based post-training pipeline, demonstrating a clean interface for multi-objective reward computation.