# jy0205/pyramid-flow architecture

> Architecture diagram of the GitHub repository jy0205/pyramid-flow, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/jy0205/pyramid-flow
- Repository: https://github.com/jy0205/pyramid-flow

## Overview

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/pyramid_dit_for_video_gen_pipeline.py`](https://github.com/jy0205/pyramid-flow/blob/main/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py) (Gen pipeline), [`train/train_video_vae.py`](https://github.com/jy0205/pyramid-flow/blob/main/train/train_video_vae.py) (VAE training), [`train/train_pyramid_flow.py`](https://github.com/jy0205/pyramid-flow/blob/main/train/train_pyramid_flow.py) (Flow training), [`scripts/train_pyramid_flow.sh`](https://github.com/jy0205/pyramid-flow/blob/main/scripts/train_pyramid_flow.sh) (Scripts)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data_prep["Data prep"]
  node_annotation_data["Annotations<br/>text/video labels<br/>[image_text.jsonl]"]
  node_dataset_layer["Dataset layer<br/>loaders<br/>[dataset_cls.py]"]
  node_feature_extraction["Text features<br/>offline prep"]
  node_latent_extraction["VAE latents<br/>offline prep"]
end

subgraph group_video_vae["Video VAE"]
  node_vae_wrapper["VAE wrapper<br/>integration layer"]
  node_vae_core["Causal VAE<br/>latent codec"]
  node_vae_blocks["VAE blocks<br/>encoder/decoder blocks"]
  node_vae_loss["VAE losses<br/>reconstruction/perceptual<br/>[modeling_loss.py]"]
end

subgraph group_pyramid_dit["Pyramid DiT"]
  node_mmdit_family["mmDiT family<br/>sd3-style backbone"]
  node_flux_family["Flux family<br/>miniFLUX-style backbone"]
  node_dit_pipeline["Gen pipeline<br/>video generation pipeline"]
  node_flow_sched["Flow schedule<br/>flow matching"]
  node_ddpm_sched["DDPM schedule<br/>diffusion schedule"]
end

subgraph group_training["Training"]
  node_vae_train["VAE training<br/>train script<br/>[train_video_vae.py]"]
  node_flow_train["Flow training<br/>train script"]
  node_dist_trainer["Dist trainer<br/>fsdp/ddp plumbing<br/>[fsdp_trainer.py]"]
end

subgraph group_inference["Inference"]
  node_gradio_app["Gradio app<br/>interactive demo<br/>[app.py]"]
  node_multi_gpu_infer["Multi-GPU infer<br/>scripted inference"]
  node_multi_gpu_app["Multi-GPU app<br/>demo runtime<br/>[app_multigpu.py]"]
end

subgraph group_ops_docs["Ops & docs"]
  node_scripts_bundle["Scripts<br/>shell wrappers"]
  node_docs_bundle["Docs<br/>[DiT.md]"]
  node_notebooks_bundle["Notebooks<br/>examples"]
  node_runtime_env["Runtime env<br/>deps config<br/>[requirements.txt]"]
end

node_annotation_data -->|"feeds"| node_dataset_layer
node_dataset_layer -.->|"prepares"| node_feature_extraction
node_dataset_layer -.->|"prepares"| node_latent_extraction
node_feature_extraction -->|"caches"| node_flow_train
node_latent_extraction -->|"caches"| node_vae_train
node_vae_blocks -->|"builds"| node_vae_core
node_vae_loss -->|"optimizes"| node_vae_train
node_vae_core -->|"exposes"| node_vae_wrapper
node_mmdit_family -->|"backbone"| node_dit_pipeline
node_flux_family -->|"backbone"| node_dit_pipeline
node_flow_sched -->|"trains"| node_flow_train
node_flow_sched -->|"samples"| node_dit_pipeline
node_ddpm_sched -.->|"optional schedule"| node_dit_pipeline
node_dit_pipeline -->|"decodes"| node_vae_wrapper
node_dit_pipeline -->|"serves"| node_gradio_app
node_dit_pipeline -->|"serves"| node_multi_gpu_infer
node_dit_pipeline -->|"serves"| node_multi_gpu_app
node_dist_trainer -->|"supports"| node_vae_train
node_dist_trainer -->|"supports"| node_flow_train
node_scripts_bundle -.->|"launches"| node_vae_train
node_scripts_bundle -.->|"launches"| node_flow_train
node_scripts_bundle -.->|"launches"| node_multi_gpu_infer
node_docs_bundle -.->|"documents"| node_vae_train
node_docs_bundle -.->|"documents"| node_flow_train
node_notebooks_bundle -.->|"examples"| node_gradio_app
node_runtime_env -.->|"supports"| node_gradio_app

click node_annotation_data "https://github.com/jy0205/pyramid-flow/blob/main/annotation/image_text.jsonl"
click node_dataset_layer "https://github.com/jy0205/pyramid-flow/blob/main/dataset/dataset_cls.py"
click node_feature_extraction "https://github.com/jy0205/pyramid-flow/blob/main/tools/extract_text_features.py"
click node_latent_extraction "https://github.com/jy0205/pyramid-flow/blob/main/tools/extract_video_vae_latents.py"
click node_vae_wrapper "https://github.com/jy0205/pyramid-flow/blob/main/video_vae/causal_video_vae_wrapper.py"
click node_vae_core "https://github.com/jy0205/pyramid-flow/blob/main/video_vae/modeling_causal_vae.py"
click node_vae_blocks "https://github.com/jy0205/pyramid-flow/blob/main/video_vae/modeling_enc_dec.py"
click node_vae_loss "https://github.com/jy0205/pyramid-flow/blob/main/video_vae/modeling_loss.py"
click node_mmdit_family "https://github.com/jy0205/pyramid-flow/blob/main/pyramid_dit/mmdit_modules/modeling_pyramid_mmdit.py"
click node_flux_family "https://github.com/jy0205/pyramid-flow/blob/main/pyramid_dit/flux_modules/modeling_pyramid_flux.py"
click node_dit_pipeline "https://github.com/jy0205/pyramid-flow/blob/main/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py"
click node_flow_sched "https://github.com/jy0205/pyramid-flow/blob/main/diffusion_schedulers/scheduling_flow_matching.py"
click node_ddpm_sched "https://github.com/jy0205/pyramid-flow/blob/main/diffusion_schedulers/scheduling_cosine_ddpm.py"
click node_vae_train "https://github.com/jy0205/pyramid-flow/blob/main/train/train_video_vae.py"
click node_flow_train "https://github.com/jy0205/pyramid-flow/blob/main/train/train_pyramid_flow.py"
click node_dist_trainer "https://github.com/jy0205/pyramid-flow/blob/main/trainer_misc/fsdp_trainer.py"
click node_gradio_app "https://github.com/jy0205/pyramid-flow/blob/main/app.py"
click node_multi_gpu_infer "https://github.com/jy0205/pyramid-flow/blob/main/inference_multigpu.py"
click node_multi_gpu_app "https://github.com/jy0205/pyramid-flow/blob/main/app_multigpu.py"
click node_scripts_bundle "https://github.com/jy0205/pyramid-flow/blob/main/scripts/train_pyramid_flow.sh"
click node_docs_bundle "https://github.com/jy0205/pyramid-flow/blob/main/docs/DiT.md"
click node_notebooks_bundle "https://github.com/jy0205/pyramid-flow/blob/main/video_generation_demo.ipynb"
click node_runtime_env "https://github.com/jy0205/pyramid-flow/blob/main/requirements.txt"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_annotation_data,node_dataset_layer,node_feature_extraction,node_latent_extraction toneBlue
class node_vae_wrapper,node_vae_core,node_vae_blocks,node_vae_loss toneAmber
class node_mmdit_family,node_flux_family,node_dit_pipeline,node_flow_sched,node_ddpm_sched toneMint
class node_vae_train,node_flow_train,node_dist_trainer toneRose
class node_gradio_app,node_multi_gpu_infer,node_multi_gpu_app toneIndigo
class node_scripts_bundle,node_docs_bundle,node_notebooks_bundle,node_runtime_env toneTeal
```

## Components

### Data prep

Curated annotations, dataset loading, and offline feature/latent extraction.

- [Annotations](https://github.com/jy0205/pyramid-flow/blob/main/annotation/image_text.jsonl): text/video labels · JSONL annotations for image-text and video-text pairs. · `annotation/image_text.jsonl`
- [Dataset layer](https://github.com/jy0205/pyramid-flow/blob/main/dataset/dataset_cls.py): loaders · Dataset abstractions, bucketed batching, and dataloader composition. · `dataset/dataset_cls.py`
- [Text features](https://github.com/jy0205/pyramid-flow/blob/main/tools/extract_text_features.py): offline prep · Precomputes reusable text features for training and inference. · `tools/extract_text_features.py`
- [VAE latents](https://github.com/jy0205/pyramid-flow/blob/main/tools/extract_video_vae_latents.py): offline prep · 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](https://github.com/jy0205/pyramid-flow/blob/main/video_vae/causal_video_vae_wrapper.py): integration layer · Adapter used by training and inference to encode/decode videos. · `video_vae/causal_video_vae_wrapper.py`
- [Causal VAE](https://github.com/jy0205/pyramid-flow/blob/main/video_vae/modeling_causal_vae.py): latent codec · Causal video VAE model defining the latent compression and reconstruction path. · `video_vae/modeling_causal_vae.py`
- [VAE blocks](https://github.com/jy0205/pyramid-flow/blob/main/video_vae/modeling_enc_dec.py): encoder/decoder blocks · Convolutional, residual, and encoder/decoder building blocks. · `video_vae/modeling_enc_dec.py`
- [VAE losses](https://github.com/jy0205/pyramid-flow/blob/main/video_vae/modeling_loss.py): reconstruction/perceptual · Perceptual, LPIPS, and discriminator-based training losses. · `video_vae/modeling_loss.py`

### Pyramid DiT

Flow-matching generative transformer with two backbone variants.

- [mmDiT family](https://github.com/jy0205/pyramid-flow/blob/main/pyramid_dit/mmdit_modules/modeling_pyramid_mmdit.py): sd3-style backbone · SD3-style Pyramid DiT variant with embeddings, blocks, normalization, and text encoder. · `pyramid_dit/mmdit_modules/modeling_pyramid_mmdit.py`
- [Flux family](https://github.com/jy0205/pyramid-flow/blob/main/pyramid_dit/flux_modules/modeling_pyramid_flux.py): miniFLUX-style backbone · miniFLUX-style Pyramid DiT variant with embeddings, blocks, normalization, and text encoder. · `pyramid_dit/flux_modules/modeling_pyramid_flux.py`
- [Gen pipeline](https://github.com/jy0205/pyramid-flow/blob/main/pyramid_dit/pyramid_dit_for_video_gen_pipeline.py): video generation 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](https://github.com/jy0205/pyramid-flow/blob/main/diffusion_schedulers/scheduling_flow_matching.py): flow matching · Main flow-matching schedule used by Pyramid Flow training and sampling. · `diffusion_schedulers/scheduling_flow_matching.py`
- [DDPM schedule](https://github.com/jy0205/pyramid-flow/blob/main/diffusion_schedulers/scheduling_cosine_ddpm.py): diffusion schedule · Alternate cosine DDPM-style schedule support. · `diffusion_schedulers/scheduling_cosine_ddpm.py`

### Training

Task-specific training entrypoints and distributed trainer plumbing.

- [VAE training](https://github.com/jy0205/pyramid-flow/blob/main/train/train_video_vae.py): train script · Training entrypoint for the causal video VAE. · `train/train_video_vae.py`
- [Flow training](https://github.com/jy0205/pyramid-flow/blob/main/train/train_pyramid_flow.py): train script · Training or fine-tuning entrypoint for Pyramid Flow. · `train/train_pyramid_flow.py`
- [Dist trainer](https://github.com/jy0205/pyramid-flow/blob/main/trainer_misc/fsdp_trainer.py): fsdp/ddp plumbing · 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](https://github.com/jy0205/pyramid-flow/blob/main/app.py): interactive demo · User-facing web demo for text-to-video and image-to-video generation. · `app.py`
- [Multi-GPU infer](https://github.com/jy0205/pyramid-flow/blob/main/inference_multigpu.py): scripted inference · Scripted inference entrypoint for multi-GPU generation. · `inference_multigpu.py`
- [Multi-GPU app](https://github.com/jy0205/pyramid-flow/blob/main/app_multigpu.py): demo runtime · Multi-GPU demo/runtime entrypoint used by the app engine. · `app_multigpu.py`

### Ops & docs

Shell wrappers, notebooks, docs, and environment definitions.

- [Scripts](https://github.com/jy0205/pyramid-flow/blob/main/scripts/train_pyramid_flow.sh): shell wrappers · Operational shell wrappers for training, extraction, inference, and demo launches. · `scripts/train_pyramid_flow.sh`
- [Docs](https://github.com/jy0205/pyramid-flow/blob/main/docs/DiT.md): architecture docs · Separate documentation for VAE and DiT training paths. · `docs/DiT.md`
- [Notebooks](https://github.com/jy0205/pyramid-flow/blob/main/video_generation_demo.ipynb): examples · Interactive example notebooks for image, video, and VAE demos. · `video_generation_demo.ipynb`
- [Runtime env](https://github.com/jy0205/pyramid-flow/blob/main/requirements.txt): deps config · Repository dependency specification and top-level runtime guidance. · `requirements.txt`

## Connections

- 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.