# d-jepa/djepa-imagenet architecture

> Architecture diagram of the GitHub repository d-jepa/djepa-imagenet, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-03). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/d-jepa/djepa-imagenet
- Repository: https://github.com/d-jepa/djepa-imagenet

## Overview

#### Purpose and entry points
D-JEPA is the official PyTorch implementation of “Denoising with a Joint-Embedding Predictive Architecture” for ImageNet-scale latent image generation. The workflow first caches pretrained-VAE latents, then trains or evaluates a D-JEPA model with configurable base/large/huge presets.

Primary entry points:
- `main_cache.py` — dataset-to-latent caching job.
- `main_djepa.py` — D-JEPA training/evaluation job.
- `scripts/cache_vae.sh`, `scripts/djepa_base.sh`, `scripts/djepa_large.sh`, `scripts/djepa_huge.sh` — launch presets.

#### Latent-data preparation boundary
Images are encoded once by a pretrained KL16 VAE and persisted as cached latents, separating expensive image encoding from model training. The VAE checkpoint is an external artifact expected under `pretrained_models/vae/kl16.ckpt`; it is not versioned in this tree.

Relevant implementation:
- `main_cache.py`
- `models/pretrained_vae.py`
- `scripts/cache_vae.sh`

This establishes the central data boundary: cache generation consumes source images and the VAE; D-JEPA consumes latent representations rather than raw images.

#### D-JEPA model architecture
The principal generative/predictive model is defined in `models/djepa.py`, with Transformer-style building blocks in `models/block.py`. It likely combines joint-embedding prediction with denoising objectives and exposes model construction used by the training engine.

Core model/loss units:
- `models/djepa.py`
- `models/block.py`
- `models/jepaloss.py`
- `models/denoising_mlp.py`

The architecture boundary is explicit: model modules own network and objective definitions, while engine modules own execution of training and evaluation steps.

#### Denoising and sampling backends
The repository supports multiple diffusion/flow-style loss and sampling implementations beneath a shared D-JEPA training flow:
- Gaussian diffusion: `models/diffloss/gaussian_diffusion/gaussian_diffusion.py`, `models/diffloss/gaussian_diffusion/diffloss.py`
- Flow matching: `models/diffloss/flow_match/flow_match.py`, `models/diffloss/flow_match/flowloss.py`

These backends are architecturally important because they determine the denoising target, time/noise trajectory, and generation procedure applied to latent tokens.

#### Training, evaluation, and distributed execution
`main_djepa.py` delegates training/evaluation orchestration to the D-JEPA engine. The engine coordinates batches, forward/loss computation, optimization, checkpoint/evaluation behavior, and likely distributed-runtime concerns.

Relevant paths:
- `main_djepa.py`
- `engine/djepa.py`
- `engine/__init__.py`

Model-scale launch configurations are represented by:
- `scripts/djepa_base.sh`
- `scripts/djepa_large.sh`
- `scripts/djepa_huge.sh`

#### Data loading and augmentation flow
Dataset ingestion and latent/cache access are centralized in utility modules, with crop logic separated from loader construction:
- `util/loader.py`
- `util/crop.py`
- `main_cache.py`

Control flow is therefore: source dataset → loader/cropping → VAE cache job, then cached latent loader → D-JEPA engine → denoising loss/sampler.

#### Optimization, evaluation, and external assets
Training schedules and distributed/general runtime utilities are modularized:
- `util/lr_sched.py`
- `util/wd_sched.py`
- `util/ema_sched.py`
- `util/misc.py`

Evaluation quality is tied to bundled ImageNet ADM FID statistics:
- `fid_stats/adm_in256_stats.npz`

Material external dependencies/services are Conda-managed Python/PyTorch dependencies, the downloadable KL16-VAE checkpoint, and Hugging Face-hosted pretrained D-JEPA weights referenced by the README.

Main entry points: [`engine/djepa.py`](https://github.com/d-jepa/djepa-imagenet/blob/main/engine/djepa.py) (D-JEPA engine), [`util/loader.py`](https://github.com/d-jepa/djepa-imagenet/blob/main/util/loader.py) (Dataset loader), [`models/djepa.py`](https://github.com/d-jepa/djepa-imagenet/blob/main/models/djepa.py) (D-JEPA model), [`main_cache.py`](https://github.com/d-jepa/djepa-imagenet/blob/main/main_cache.py) (Latent cache job)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_prepare["Latent preparation"]
  node_cache_launch["VAE cache launcher<br/>shell preset<br/>[cache_vae.sh]"]
  node_cache_main["Latent cache job<br/>Python entry point<br/>[main_cache.py]"]
  node_source_images[("Source image dataset")]
  node_loader["Dataset loader<br/>data loading<br/>[loader.py]"]
  node_crop["Crop transforms<br/>augmentation<br/>[crop.py]"]
  node_vae["Pretrained KL16 VAE<br/>latent encoder<br/>[pretrained_vae.py]"]
  node_latent_cache[("Cached latents<br/>persistent dataset")]
end

subgraph group_train["Training runtime"]
  node_djepa_launchers["D-JEPA scale presets<br/>shell presets<br/>[djepa_base.sh]"]
  node_djepa_main["D-JEPA job<br/>Python entry point<br/>[main_djepa.py]"]
  node_engine["D-JEPA engine<br/>training engine<br/>[djepa.py]"]
  node_runtime["Distributed runtime<br/>runtime utilities<br/>[misc.py]"]
  node_schedules["Optimizer schedules<br/>training schedules<br/>[lr_sched.py]"]
end

subgraph group_model["Model and objectives"]
  node_djepa_model["D-JEPA model<br/>joint-embedding model<br/>[djepa.py]"]
  node_transformer_blocks["Transformer blocks<br/>model building blocks<br/>[block.py]"]
  node_jepa_loss["JEPA objective<br/>loss<br/>[jepaloss.py]"]
  node_denoising_mlp["Denoising MLP<br/>denoising head<br/>[denoising_mlp.py]"]
  node_gaussian_backend["Gaussian diffusion backend<br/>denoising backend"]
  node_flow_backend["Flow matching backend<br/>denoising backend<br/>[flow_match.py]"]
end

subgraph group_assets["External assets"]
  node_vae_checkpoint["KL16 VAE checkpoint<br/>external checkpoint"]
  node_fid_stats["ImageNet ADM FID stats<br/>evaluation asset"]
  node_published_weights["Published D-JEPA weights<br/>external weights"]
end

node_cache_launch -->|"launches"| node_cache_main
node_source_images -->|"ingested by"| node_loader
node_crop -->|"applies crops"| node_loader
node_loader -->|"image batches"| node_cache_main
node_vae_checkpoint -->|"loads"| node_vae
node_vae -->|"encodes images"| node_cache_main
node_cache_main -->|"persists"| node_latent_cache
node_djepa_launchers -->|"launches"| node_djepa_main
node_latent_cache -->|"read as training data"| node_loader
node_loader -->|"latent batches"| node_engine
node_djepa_main -->|"delegates to"| node_engine
node_runtime -->|"supports execution"| node_engine
node_schedules -->|"controls optimization"| node_engine
node_engine -->|"trains and evaluates"| node_djepa_model
node_transformer_blocks -->|"compose"| node_djepa_model
node_djepa_model -->|"optimizes"| node_jepa_loss
node_djepa_model -->|"uses"| node_denoising_mlp
node_denoising_mlp -.->|"selects backend"| node_gaussian_backend
node_denoising_mlp -.->|"selects backend"| node_flow_backend
node_engine -->|"evaluates against"| node_fid_stats
node_published_weights -.->|"pretrained artifact"| node_djepa_model

click node_cache_launch "https://github.com/d-jepa/djepa-imagenet/blob/main/scripts/cache_vae.sh"
click node_cache_main "https://github.com/d-jepa/djepa-imagenet/blob/main/main_cache.py"
click node_loader "https://github.com/d-jepa/djepa-imagenet/blob/main/util/loader.py"
click node_crop "https://github.com/d-jepa/djepa-imagenet/blob/main/util/crop.py"
click node_vae "https://github.com/d-jepa/djepa-imagenet/blob/main/models/pretrained_vae.py"
click node_djepa_launchers "https://github.com/d-jepa/djepa-imagenet/blob/main/scripts/djepa_base.sh"
click node_djepa_main "https://github.com/d-jepa/djepa-imagenet/blob/main/main_djepa.py"
click node_engine "https://github.com/d-jepa/djepa-imagenet/blob/main/engine/djepa.py"
click node_runtime "https://github.com/d-jepa/djepa-imagenet/blob/main/util/misc.py"
click node_schedules "https://github.com/d-jepa/djepa-imagenet/blob/main/util/lr_sched.py"
click node_djepa_model "https://github.com/d-jepa/djepa-imagenet/blob/main/models/djepa.py"
click node_transformer_blocks "https://github.com/d-jepa/djepa-imagenet/blob/main/models/block.py"
click node_jepa_loss "https://github.com/d-jepa/djepa-imagenet/blob/main/models/jepaloss.py"
click node_denoising_mlp "https://github.com/d-jepa/djepa-imagenet/blob/main/models/denoising_mlp.py"
click node_gaussian_backend "https://github.com/d-jepa/djepa-imagenet/blob/main/models/diffloss/gaussian_diffusion/gaussian_diffusion.py"
click node_flow_backend "https://github.com/d-jepa/djepa-imagenet/blob/main/models/diffloss/flow_match/flow_match.py"
click node_fid_stats "https://github.com/d-jepa/djepa-imagenet/blob/main/fid_stats/adm_in256_stats.npz"

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_cache_launch,node_cache_main,node_source_images,node_loader,node_crop,node_vae,node_latent_cache toneBlue
class node_djepa_launchers,node_djepa_main,node_engine,node_runtime,node_schedules toneAmber
class node_djepa_model,node_transformer_blocks,node_jepa_loss,node_denoising_mlp,node_gaussian_backend,node_flow_backend toneMint
class node_vae_checkpoint,node_fid_stats,node_published_weights toneRose
```

## Components

### Latent preparation

Encode source images once before D-JEPA training.

- [VAE cache launcher](https://github.com/d-jepa/djepa-imagenet/blob/main/scripts/cache_vae.sh): shell preset · `scripts/cache_vae.sh`
- [Latent cache job](https://github.com/d-jepa/djepa-imagenet/blob/main/main_cache.py): Python entry point · Creates persistent VAE latents from image batches. · `main_cache.py`
- **Source image dataset**: dataset
- [Dataset loader](https://github.com/d-jepa/djepa-imagenet/blob/main/util/loader.py): data loading · Builds image and cached-latent data access. · `util/loader.py`
- [Crop transforms](https://github.com/d-jepa/djepa-imagenet/blob/main/util/crop.py): augmentation · `util/crop.py`
- [Pretrained KL16 VAE](https://github.com/d-jepa/djepa-imagenet/blob/main/models/pretrained_vae.py): latent encoder · Encodes images into training latents. · `models/pretrained_vae.py`
- **Cached latents**: persistent dataset · Boundary between image encoding and D-JEPA execution.

### Training runtime

Train or evaluate D-JEPA on cached latents.

- [D-JEPA scale presets](https://github.com/d-jepa/djepa-imagenet/blob/main/scripts/djepa_base.sh): shell presets · Base, large, and huge launch configurations. · `scripts/djepa_base.sh`
- [D-JEPA job](https://github.com/d-jepa/djepa-imagenet/blob/main/main_djepa.py): Python entry point · Starts training or evaluation. · `main_djepa.py`
- [D-JEPA engine](https://github.com/d-jepa/djepa-imagenet/blob/main/engine/djepa.py): training engine · Coordinates batches, optimization, checkpoints, and evaluation. · `engine/djepa.py`
- [Distributed runtime](https://github.com/d-jepa/djepa-imagenet/blob/main/util/misc.py): runtime utilities · Shared distributed and general execution support. · `util/misc.py`
- [Optimizer schedules](https://github.com/d-jepa/djepa-imagenet/blob/main/util/lr_sched.py): training schedules · Learning-rate, weight-decay, and EMA scheduling. · `util/lr_sched.py`

### Model and objectives

Network definitions and interchangeable denoising backends.

- [D-JEPA model](https://github.com/d-jepa/djepa-imagenet/blob/main/models/djepa.py): joint-embedding model · Main predictive and generative network. · `models/djepa.py`
- [Transformer blocks](https://github.com/d-jepa/djepa-imagenet/blob/main/models/block.py): model building blocks · `models/block.py`
- [JEPA objective](https://github.com/d-jepa/djepa-imagenet/blob/main/models/jepaloss.py): loss · `models/jepaloss.py`
- [Denoising MLP](https://github.com/d-jepa/djepa-imagenet/blob/main/models/denoising_mlp.py): denoising head · `models/denoising_mlp.py`
- [Gaussian diffusion backend](https://github.com/d-jepa/djepa-imagenet/blob/main/models/diffloss/gaussian_diffusion/gaussian_diffusion.py): denoising backend · Defines Gaussian diffusion loss and sampling trajectory. · `models/diffloss/gaussian_diffusion/gaussian_diffusion.py`
- [Flow matching backend](https://github.com/d-jepa/djepa-imagenet/blob/main/models/diffloss/flow_match/flow_match.py): denoising backend · Defines flow-matching loss and generation trajectory. · `models/diffloss/flow_match/flow_match.py`

### External assets

Checkpoints, statistics, and published weights.

- **KL16 VAE checkpoint**: external checkpoint · Downloaded artifact; not versioned in the repository.
- [ImageNet ADM FID stats](https://github.com/d-jepa/djepa-imagenet/blob/main/fid_stats/adm_in256_stats.npz): evaluation asset · `fid_stats/adm_in256_stats.npz`
- **Published D-JEPA weights**: external weights · Hugging Face-hosted pretrained model artifacts.

## Connections

- VAE cache launcher → Latent cache job: launches
- Source image dataset → Dataset loader: ingested by
- Crop transforms → Dataset loader: applies crops
- Dataset loader → Latent cache job: image batches
- KL16 VAE checkpoint → Pretrained KL16 VAE: loads
- Pretrained KL16 VAE → Latent cache job: encodes images
- Latent cache job → Cached latents: persists
- D-JEPA scale presets → D-JEPA job: launches
- Cached latents → Dataset loader: read as training data
- Dataset loader → D-JEPA engine: latent batches
- D-JEPA job → D-JEPA engine: delegates to
- Distributed runtime → D-JEPA engine: supports execution
- Optimizer schedules → D-JEPA engine: controls optimization
- D-JEPA engine → D-JEPA model: trains and evaluates
- Transformer blocks → D-JEPA model: compose
- D-JEPA model → JEPA objective: optimizes
- D-JEPA model → Denoising MLP: uses
- Denoising MLP → Gaussian diffusion backend: selects backend
- Denoising MLP → Flow matching backend: selects backend
- D-JEPA engine → ImageNet ADM FID stats: evaluates against
- Published D-JEPA weights → D-JEPA model: pretrained artifact