# lth14/jit architecture

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

- Interactive diagram: https://gitdiagram.com/lth14/jit
- Repository: https://github.com/lth14/jit

## Overview

JiT is a PyTorch implementation of class-conditioned, pixel-space diffusion for high-resolution image generation. The main workflow is ImageNet training through a denoiser and JiT transformer, followed by sampling and FID/IS evaluation; the repository also provides reference-image preparation and precomputed evaluation statistics. The graph reflects documented capabilities where source excerpts are incomplete; in particular, training data-loader wiring and some model internals are not fully shown in the sampled source.

Main entry points: [`engine_jit.py`](https://github.com/lth14/jit/blob/main/engine_jit.py) (Train and evaluate), [`main_jit.py`](https://github.com/lth14/jit/blob/main/main_jit.py) (Training runner), [`denoiser.py`](https://github.com/lth14/jit/blob/main/denoiser.py) (Denoiser and sampler), [`model_jit.py`](https://github.com/lth14/jit/blob/main/model_jit.py) (JiT transformer)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_runtime["Training and orchestration"]
  node_main["Training runner<br/>[main_jit.py]"]
  node_engine["Train and evaluate<br/>[engine_jit.py]"]
  node_training_utils["Training utilities<br/>[misc.py]"]
  node_schedule["Learning-rate schedule<br/>[lr_sched.py]"]
end

subgraph group_model["Diffusion model"]
  node_denoiser["Denoiser and sampler<br/>[denoiser.py]"]
  node_jit["JiT transformer<br/>[model_jit.py]"]
  node_model_helpers["Position and norm helpers<br/>[model_util.py]"]
end

subgraph group_data["Image data"]
  node_reference_prep["Reference-image preparation<br/>[prepare_ref.py]"]
  node_crop["Image center crop<br/>[crop.py]"]
end

subgraph group_evaluation["Generation and evaluation"]
  node_generated_images["Generated images"]
  node_fid_stats[("FID reference statistics")]
  node_metrics["FID and IS metrics"]
end

node_user(("Researcher"))
node_imagenet[("ImageNet")]

node_user -->|"starts training"| node_main
node_imagenet -->|"training data"| node_main
node_main -->|"creates model"| node_denoiser
node_main -->|"runs epochs"| node_engine
node_main -->|"saves checkpoints"| node_training_utils
node_engine -->|"trains and samples"| node_denoiser
node_engine -->|"adjusts learning rate"| node_schedule
node_engine -->|"logs and synchronizes"| node_training_utils
node_denoiser -->|"predicts denoised image"| node_jit
node_jit -.->|"uses model helpers"| node_model_helpers
node_imagenet -->|"reference images"| node_reference_prep
node_reference_prep -->|"crops images"| node_crop
node_engine -->|"writes samples"| node_generated_images
node_engine -->|"evaluates against"| node_fid_stats
node_engine -->|"computes FID and IS"| node_metrics

click node_main "https://github.com/lth14/jit/blob/main/main_jit.py"
click node_engine "https://github.com/lth14/jit/blob/main/engine_jit.py"
click node_training_utils "https://github.com/lth14/jit/blob/main/util/misc.py"
click node_schedule "https://github.com/lth14/jit/blob/main/util/lr_sched.py"
click node_denoiser "https://github.com/lth14/jit/blob/main/denoiser.py"
click node_jit "https://github.com/lth14/jit/blob/main/model_jit.py"
click node_model_helpers "https://github.com/lth14/jit/blob/main/util/model_util.py"
click node_reference_prep "https://github.com/lth14/jit/blob/main/prepare_ref.py"
click node_crop "https://github.com/lth14/jit/blob/main/util/crop.py"
click node_fid_stats "https://github.com/lth14/jit/tree/main/fid_stats"

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_main,node_engine,node_training_utils,node_schedule toneBlue
class node_denoiser,node_jit,node_model_helpers,node_imagenet toneAmber
class node_reference_prep,node_crop toneMint
class node_generated_images,node_fid_stats,node_metrics toneRose
class node_user toneIndigo
```

## Components

### Training and orchestration

- [Training runner](https://github.com/lth14/jit/blob/main/main_jit.py): `main_jit.py`
- [Train and evaluate](https://github.com/lth14/jit/blob/main/engine_jit.py): `engine_jit.py`
- [Training utilities](https://github.com/lth14/jit/blob/main/util/misc.py): `util/misc.py`
- [Learning-rate schedule](https://github.com/lth14/jit/blob/main/util/lr_sched.py): `util/lr_sched.py`

### Diffusion model

- [Denoiser and sampler](https://github.com/lth14/jit/blob/main/denoiser.py): `denoiser.py`
- [JiT transformer](https://github.com/lth14/jit/blob/main/model_jit.py): `model_jit.py`
- [Position and norm helpers](https://github.com/lth14/jit/blob/main/util/model_util.py): `util/model_util.py`

### Image data

- [Reference-image preparation](https://github.com/lth14/jit/blob/main/prepare_ref.py): `prepare_ref.py`
- [Image center crop](https://github.com/lth14/jit/blob/main/util/crop.py): `util/crop.py`

### Generation and evaluation

- **Generated images**
- [FID reference statistics](https://github.com/lth14/jit/tree/main/fid_stats): `fid_stats`
- **FID and IS metrics**

### Other components

- **Researcher**
- **ImageNet**

## Connections

- Researcher → Training runner: starts training
- ImageNet → Training runner: training data
- Training runner → Denoiser and sampler: creates model
- Training runner → Train and evaluate: runs epochs
- Training runner → Training utilities: saves checkpoints
- Train and evaluate → Denoiser and sampler: trains and samples
- Train and evaluate → Learning-rate schedule: adjusts learning rate
- Train and evaluate → Training utilities: logs and synchronizes
- Denoiser and sampler → JiT transformer: predicts denoised image
- JiT transformer → Position and norm helpers: uses model helpers
- ImageNet → Reference-image preparation: reference images
- Reference-image preparation → Image center crop: crops images
- Train and evaluate → Generated images: writes samples
- Train and evaluate → FID reference statistics: evaluates against
- Train and evaluate → FID and IS metrics: computes FID and IS