d-jepa/djepa-imagenet

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.