How jy0205/pyramid-flow works
1. What this repo is
- Pyramid Flow is a video generation system, centered on an autoregressive flow-matching DiT that works in latent space.
- The repo also includes the causal video VAE used to compress/decompress videos, plus training, inference, and demo entrypoints.
2. End-to-end data flow
- Text prompt and optional input image enter the text encoder / VAE encoder side.
- The Pyramid DiT predicts latent trajectories over multiple generation stages.
- The video VAE decoder turns final latents into frames.
- Output is exported as a video file or displayed in the Gradio demo.
3. Core model boundary: video VAE
- The video_vae/ package implements the latent codec: causal 3D convolutions, encoder/decoder blocks, residual blocks, discriminator, perceptual loss, and LPIPS.
- video_vae/causal_video_vae_wrapper.py is the integration layer used by training/inference.
- This subsystem is trained separately from the DiT, and its role is to make video generation feasible at high resolution.
4. Core model boundary: Pyramid DiT
- pyramid_dit/ contains the generative transformer used for flow matching.
- There are two internal model families:
- mmdit_modules/ for the SD3-style variant
- flux_modules/ for the miniFLUX-style variant
- Both families contain embeddings, normalization, text encoding, transformer blocks, and the top-level Pyramid model wrappers.
5. Inference wrapper
- pyramid_dit/pyramid_dit_for_video_gen_pipeline.py is the orchestration layer that ties together text encoding, denoising/sampling, and VAE decoding.
- This is the main API used by app.py, notebooks, and scripts such as inference_multigpu.py.
- It exposes both text-to-video and image-to-video generation paths.
6. Schedulers / training objective
- diffusion_schedulers/ contains the sampling and training schedules:
- scheduling_flow_matching.py for the main Pyramid Flow objective
- scheduling_cosine_ddpm.py for diffusion-style scheduling support
- These modules define how latent noise is traversed during training and inference, so they sit between model prediction and generated samples.
7. Data preparation layer
- dataset/ contains dataset abstractions and loaders for video-text pairs and bucketed batching.
- The annotation/*.jsonl files provide text annotations for image/video samples.
- This layer is the boundary between raw curated data and the cached training inputs used by VAE and DiT training.
8. Offline feature/latent extraction
- tools/extract_text_features.py and tools/extract_video_vae_latents.py precompute reusable training artifacts.
- The corresponding shell scripts in scripts/ indicate these are explicit preprocessing stages, not inline steps.
- This reduces training cost by separating feature extraction from optimization.
9. Training architecture
- train/train_video_vae.py trains the causal VAE.
- train/train_pyramid_flow.py trains or fine-tunes the DiT/flow model.
- trainer_misc/ holds the distributed-training plumbing: FSDP trainer, VAE DDP trainer, communication helpers, utilities, and process-group support.
- The repo therefore has two distinct training pipelines with shared distributed infrastructure.
10. Distributed and memory-optimized inference
- app_multigpu.py, inference_multigpu.py, and the scripts/*multigpu* wrappers support multi-GPU generation.
- The README calls out sequence parallelism for inference and CPU offloading for low-memory runs.
- This is an inference-time boundary: generation can run in a single-GPU, offloaded, or multi-GPU mode without changing the model core.
11. User-facing entrypoints
- app.py is the Gradio demo entrypoint for interactive generation.
- The notebooks (image_generation_demo.ipynb, video_generation_demo.ipynb, causal_video_vae_demo.ipynb) are example clients rather than core components.
- These frontends consume the same model wrapper used in scripted inference.
12. Operational and documentation split
- docs/VAE.md and docs/DiT.md document the two training paths separately.
- scripts/ acts as the operational layer for reproducible runs, while requirements.txt and the README define the runtime environment.
- The architecture is thus split into: data prep, VAE training, DiT training, and inference/deployment.
13. Key architectural idea to capture in a diagram
- The repo is not “one model,” but a stack: curated text/video data → preprocessing caches → causal video VAE → Pyramid DiT flow model → decoder/export.
- The main variant choice is the DiT backbone (SD3/mmDiT vs miniFLUX/Flux), while the VAE and tooling provide the shared substrate underneath.
Main entry points: pyramid_dit_for_video_gen_pipeline.py Gen pipeline · train_video_vae.py VAE training · train_pyramid_flow.py Flow training · train_pyramid_flow.sh Scripts
Components 23
Data prep
Curated annotations, dataset loading, and offline feature/latent extraction.
- Annotations · JSONL annotations for image-text and video-text pairs.
annotation/image_text.jsonl - Dataset layer · Dataset abstractions, bucketed batching, and dataloader composition.
dataset/dataset_cls.py - Text features · Precomputes reusable text features for training and inference.
tools/extract_text_features.py - VAE latents · Precomputes video VAE latents for faster downstream training.
tools/extract_video_vae_latents.py
Video VAE
Causal video codec that compresses and reconstructs latent video representations.
- VAE wrapper · Adapter used by training and inference to encode/decode videos.
video_vae/causal_video_vae_wrapper.py - Causal VAE · Causal video VAE model defining the latent compression and reconstruction path.
video_vae/modeling_causal_vae.py - VAE blocks · Convolutional, residual, and encoder/decoder building blocks.
video_vae/modeling_enc_dec.py - VAE losses · Perceptual, LPIPS, and discriminator-based training losses.
video_vae/modeling_loss.py
Pyramid DiT
Flow-matching generative transformer with two backbone variants.
- mmDiT family · SD3-style Pyramid DiT variant with embeddings, blocks, normalization, and text encoder.
pyramid_dit/mmdit_modules/modeling_pyramid_mmdit.py - Flux family · miniFLUX-style Pyramid DiT variant with embeddings, blocks, normalization, and text encoder.
pyramid_dit/flux_modules/modeling_pyramid_flux.py - Gen pipeline · Orchestrates text encoding, sampling, and VAE decoding for text-to-video and image-to-video.
pyramid_dit/pyramid_dit_for_video_gen_pipeline.py - Flow schedule · Main flow-matching schedule used by Pyramid Flow training and sampling.
diffusion_schedulers/scheduling_flow_matching.py - DDPM schedule · Alternate cosine DDPM-style schedule support.
diffusion_schedulers/scheduling_cosine_ddpm.py
Training
Task-specific training entrypoints and distributed trainer plumbing.
- VAE training · Training entrypoint for the causal video VAE.
train/train_video_vae.py - Flow training · Training or fine-tuning entrypoint for Pyramid Flow.
train/train_pyramid_flow.py - Dist trainer · Distributed training helpers for FSDP, VAE DDP, communication, and process-group utilities.
trainer_misc/fsdp_trainer.py
Inference
Interactive and scripted generation entrypoints and orchestration.
- Gradio app · User-facing web demo for text-to-video and image-to-video generation.
app.py - Multi-GPU infer · Scripted inference entrypoint for multi-GPU generation.
inference_multigpu.py - Multi-GPU app · Multi-GPU demo/runtime entrypoint used by the app engine.
app_multigpu.py
Ops & docs
Shell wrappers, notebooks, docs, and environment definitions.
- Scripts · Operational shell wrappers for training, extraction, inference, and demo launches.
scripts/train_pyramid_flow.sh - Docs · Separate documentation for VAE and DiT training paths.
docs/DiT.md - Notebooks · Interactive example notebooks for image, video, and VAE demos.
video_generation_demo.ipynb - Runtime env · Repository dependency specification and top-level runtime guidance.
requirements.txt
Connections 26
- Annotations → Dataset layer: feeds · Curated annotations are consumed by dataset loaders.
- Dataset layer → Text features: prepares · Dataset inputs can be preprocessed into text features.
- Dataset layer → VAE latents: prepares · Video samples can be precomputed into VAE latents.
- Text features → Flow training: caches · Text features reduce training-time encoding cost.
- VAE latents → VAE training: caches · Latent caches support faster VAE and downstream training.
- VAE blocks → Causal VAE: builds · Encoder/decoder and residual blocks compose the causal VAE.
- VAE losses → VAE training: optimizes · Reconstruction and perceptual losses drive VAE training.
- Causal VAE → VAE wrapper: exposes · The wrapper integrates the codec into pipeline code.
- mmDiT family → Gen pipeline: backbone · mmDiT variant can serve as the generative core.
- Flux family → Gen pipeline: backbone · Flux variant can serve as the generative core.
- Flow schedule → Flow training: trains · Flow matching schedule defines the training objective.
- Flow schedule → Gen pipeline: samples · Schedule also governs inference-time latent traversal.
- DDPM schedule → Gen pipeline: optional schedule · Alternate schedule support for diffusion-style runs.
- Gen pipeline → VAE wrapper: decodes · Final latents are decoded back into video frames.
- Gen pipeline → Gradio app: serves · The pipeline is used by the interactive demo.
- Gen pipeline → Multi-GPU infer: serves · The pipeline is used by scripted inference.
- Gen pipeline → Multi-GPU app: serves · The pipeline is used by multi-GPU demo/runtime entrypoints.
- Dist trainer → VAE training: supports · Distributed plumbing supports VAE training.
- Dist trainer → Flow training: supports · Distributed plumbing supports Pyramid Flow training.
- Scripts → VAE training: launches · Shell wrappers launch the VAE training entrypoint.
- Scripts → Flow training: launches · Shell wrappers launch the flow training entrypoint.
- Scripts → Multi-GPU infer: launches · Shell wrappers launch multi-GPU inference.
- Docs → VAE training: documents · VAE documentation explains the separate codec training path.
- Docs → Flow training: documents · DiT documentation explains the flow model training path.
- Notebooks → Gradio app: examples · Notebooks act as example clients for the same generation stack.
- Runtime env → Gradio app: supports · Shared dependencies underpin runtime entrypoints.