# leaplabthu/lineardiff architecture

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

- Interactive diagram: https://gitdiagram.com/leaplabthu/lineardiff
- Repository: https://github.com/leaplabthu/lineardiff

## Overview

LinearDiff is a PyTorch research implementation with two user-facing workflows: ImageNet classification using VCA-DeiT, and class-conditioned image generation using VCA-DiT or VCA-SiT, followed by sample evaluation. The graph preserves those separate training and inference paths and the SiT transport layer. Several links to generation internals rely on the documented workflow and repository layout; unsampled wiring is not inferred beyond those documented roles.

Main entry points: [`image_classification/main.py`](https://github.com/leaplabthu/lineardiff/blob/master/image_classification/main.py) (Train / evaluate), [`image_generation/models/dit/sample_ddp_nonpz.py`](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/sample_ddp_nonpz.py) (Sample DiT), [`image_generation/models/sit/train.py`](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/train.py) (Train SiT), [`image_generation/models/sit/sample_ddp_nonpz.py`](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/sample_ddp_nonpz.py) (Sample SiT)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_classification["Classification"]
  node_class_entry["Train / evaluate<br/>[main.py]"]
  node_class_data["Data pipeline<br/>[build.py]"]
  node_class_storage["Image dataset access"]
  node_class_model["VCA-DeiT<br/>[vca_deit.py]"]
  node_class_optim["Optimization schedule<br/>[optimizer.py]"]
  node_class_scheduler["Learning-rate schedule<br/>[lr_scheduler.py]"]
  node_class_result["Classifier checkpoint"]
end

subgraph group_dit["DiT Generation"]
  node_dit_train["Train DiT<br/>[train.py]"]
  node_dit_model["VCA-DiT model<br/>[models.py]"]
  node_dit_diffusion["Gaussian diffusion"]
  node_dit_sample["Sample DiT"]
  node_dit_checkpoint["DiT checkpoint"]
end

subgraph group_sit["SiT Generation"]
  node_sit_train["Train SiT<br/>[train.py]"]
  node_sit_model["VCA-SiT model<br/>[models.py]"]
  node_sit_transport["Transport and sampler<br/>[transport.py]"]
  node_sit_sample["Sample SiT"]
  node_sit_checkpoint["SiT checkpoint"]
end

subgraph group_evaluation["Evaluation"]
  node_reference_tool["Build reference stats<br/>[create_ref_npz.py]"]
  node_evaluator["FID evaluation<br/>[evaluator_read.py]"]
end

node_operator(("Researcher"))
node_imagenet[("ImageNet data")]
node_samples["Generated samples"]
node_reference["Reference statistics"]
node_metrics["Evaluation metrics"]

node_operator -->|"starts"| node_class_entry
node_imagenet -->|"provides images"| node_class_data
node_class_entry -->|"builds loaders"| node_class_data
node_class_data -->|"loads images"| node_class_storage
node_class_entry -->|"builds model"| node_class_model
node_class_entry -->|"builds optimizer"| node_class_optim
node_class_entry -->|"builds schedule"| node_class_scheduler
node_class_entry -->|"saves checkpoint"| node_class_result
node_operator -->|"starts"| node_dit_train
node_imagenet -->|"provides images"| node_dit_train
node_dit_train -->|"trains"| node_dit_model
node_dit_train -->|"saves checkpoint"| node_dit_checkpoint
node_operator -->|"starts"| node_dit_sample
node_dit_checkpoint -->|"loads checkpoint"| node_dit_sample
node_dit_sample -->|"samples with"| node_dit_model
node_dit_sample -->|"uses diffusion"| node_dit_diffusion
node_dit_sample -->|"writes samples"| node_samples
node_operator -->|"starts"| node_sit_train
node_imagenet -->|"provides images"| node_sit_train
node_sit_train -->|"trains"| node_sit_model
node_sit_train -->|"computes losses"| node_sit_transport
node_sit_train -->|"saves checkpoint"| node_sit_checkpoint
node_operator -->|"starts"| node_sit_sample
node_sit_checkpoint -->|"loads checkpoint"| node_sit_sample
node_sit_sample -->|"samples with"| node_sit_model
node_sit_sample -->|"uses sampler"| node_sit_transport
node_sit_sample -->|"writes samples"| node_samples
node_imagenet -->|"provides images"| node_reference_tool
node_operator -->|"creates reference"| node_reference_tool
node_reference_tool -->|"writes statistics"| node_reference
node_operator -->|"runs evaluation"| node_evaluator
node_samples -->|"provides samples"| node_evaluator
node_reference -->|"provides baseline"| node_evaluator
node_evaluator -->|"computes metrics"| node_metrics

click node_class_entry "https://github.com/leaplabthu/lineardiff/blob/master/image_classification/main.py"
click node_class_data "https://github.com/leaplabthu/lineardiff/blob/master/image_classification/data/build.py"
click node_class_storage "https://github.com/leaplabthu/lineardiff/blob/master/image_classification/data/cached_image_folder.py"
click node_class_model "https://github.com/leaplabthu/lineardiff/blob/master/image_classification/models/vca_deit.py"
click node_class_optim "https://github.com/leaplabthu/lineardiff/blob/master/image_classification/optimizer.py"
click node_class_scheduler "https://github.com/leaplabthu/lineardiff/blob/master/image_classification/lr_scheduler.py"
click node_dit_train "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/train.py"
click node_dit_model "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/models.py"
click node_dit_diffusion "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/diffusion/gaussian_diffusion.py"
click node_dit_sample "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/sample_ddp_nonpz.py"
click node_sit_train "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/train.py"
click node_sit_model "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/models.py"
click node_sit_transport "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/transport/transport.py"
click node_sit_sample "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/sample_ddp_nonpz.py"
click node_reference_tool "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/tools/create_ref_npz.py"
click node_evaluator "https://github.com/leaplabthu/lineardiff/blob/master/image_generation/tools/evaluator_read.py"

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_class_entry,node_class_data,node_class_storage,node_class_model,node_class_optim,node_class_scheduler,node_class_result toneBlue
class node_dit_train,node_dit_model,node_dit_diffusion,node_dit_sample,node_dit_checkpoint,node_imagenet toneAmber
class node_sit_train,node_sit_model,node_sit_transport,node_sit_sample,node_sit_checkpoint toneMint
class node_reference_tool,node_evaluator toneRose
class node_operator,node_samples,node_reference,node_metrics toneIndigo
```

## Components

### Classification

- [Train / evaluate](https://github.com/leaplabthu/lineardiff/blob/master/image_classification/main.py): `image_classification/main.py`
- [Data pipeline](https://github.com/leaplabthu/lineardiff/blob/master/image_classification/data/build.py): `image_classification/data/build.py`
- [Image dataset access](https://github.com/leaplabthu/lineardiff/blob/master/image_classification/data/cached_image_folder.py): `image_classification/data/cached_image_folder.py`
- [VCA-DeiT](https://github.com/leaplabthu/lineardiff/blob/master/image_classification/models/vca_deit.py): `image_classification/models/vca_deit.py`
- [Optimization schedule](https://github.com/leaplabthu/lineardiff/blob/master/image_classification/optimizer.py): `image_classification/optimizer.py`
- [Learning-rate schedule](https://github.com/leaplabthu/lineardiff/blob/master/image_classification/lr_scheduler.py): `image_classification/lr_scheduler.py`
- **Classifier checkpoint**

### DiT Generation

- [Train DiT](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/train.py): `image_generation/models/dit/train.py`
- [VCA-DiT model](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/models.py): `image_generation/models/dit/models.py`
- [Gaussian diffusion](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/diffusion/gaussian_diffusion.py): `image_generation/models/dit/diffusion/gaussian_diffusion.py`
- [Sample DiT](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/dit/sample_ddp_nonpz.py): `image_generation/models/dit/sample_ddp_nonpz.py`
- **DiT checkpoint**

### SiT Generation

- [Train SiT](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/train.py): `image_generation/models/sit/train.py`
- [VCA-SiT model](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/models.py): `image_generation/models/sit/models.py`
- [Transport and sampler](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/transport/transport.py): `image_generation/models/sit/transport/transport.py`
- [Sample SiT](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/models/sit/sample_ddp_nonpz.py): `image_generation/models/sit/sample_ddp_nonpz.py`
- **SiT checkpoint**

### Evaluation

- [Build reference stats](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/tools/create_ref_npz.py): `image_generation/tools/create_ref_npz.py`
- [FID evaluation](https://github.com/leaplabthu/lineardiff/blob/master/image_generation/tools/evaluator_read.py): `image_generation/tools/evaluator_read.py`

### Other components

- **Researcher**
- **ImageNet data**
- **Generated samples**
- **Reference statistics**
- **Evaluation metrics**

## Connections

- Researcher → Train / evaluate: starts
- ImageNet data → Data pipeline: provides images
- Train / evaluate → Data pipeline: builds loaders
- Data pipeline → Image dataset access: loads images
- Train / evaluate → VCA-DeiT: builds model
- Train / evaluate → Optimization schedule: builds optimizer
- Train / evaluate → Learning-rate schedule: builds schedule
- Train / evaluate → Classifier checkpoint: saves checkpoint
- Researcher → Train DiT: starts
- ImageNet data → Train DiT: provides images
- Train DiT → VCA-DiT model: trains
- Train DiT → DiT checkpoint: saves checkpoint
- Researcher → Sample DiT: starts
- DiT checkpoint → Sample DiT: loads checkpoint
- Sample DiT → VCA-DiT model: samples with
- Sample DiT → Gaussian diffusion: uses diffusion
- Sample DiT → Generated samples: writes samples
- Researcher → Train SiT: starts
- ImageNet data → Train SiT: provides images
- Train SiT → VCA-SiT model: trains
- Train SiT → Transport and sampler: computes losses
- Train SiT → SiT checkpoint: saves checkpoint
- Researcher → Sample SiT: starts
- SiT checkpoint → Sample SiT: loads checkpoint
- Sample SiT → VCA-SiT model: samples with
- Sample SiT → Transport and sampler: uses sampler
- Sample SiT → Generated samples: writes samples
- ImageNet data → Build reference stats: provides images
- Researcher → Build reference stats: creates reference
- Build reference stats → Reference statistics: writes statistics
- Researcher → FID evaluation: runs evaluation
- Generated samples → FID evaluation: provides samples
- Reference statistics → FID evaluation: provides baseline
- FID evaluation → Evaluation metrics: computes metrics