lambertae/drifting

How lambertae/drifting works

This is the JAX release for ImageNet one-step image generation via Drifting. The main user workflow is to train a generator or run inference from a local/Hugging Face checkpoint; inference generates samples and evaluates FID, IS, and (at 50k samples) precision/recall. A separate preparation stage encodes ImageNet into a latent cache for latent-space training. The graph preserves MAE feature-model training, checkpoint/artifact loading, and distributed runtime support. Several training and metric internals are unsampled; those links rely on the SOURCE INDEX, while optional W&B reporting is documented in the README.

Main entry points: train.py Generator training · train_mae.py MAE training · inference.py Inference entrypoint · dataset.py ImageNet loader

Components 24

Training workflows

Run dispatcher main.py
Generator training train.py
MAE training train_mae.py
Drifting loss drift_loss.py
Memory bank memory_bank.py

Models and artifacts

Generator model models/generator.py
MAE feature model models/mae_model.py
ConvNeXt encoder models/convnext.py
Artifact loading utils/init_util.py
HF artifact helpers models/hf.py
Checkpoint utilities utils/ckpt_util.py

Data preparation

Latent cache builder dataset/latent.py
VAE encode/decode dataset/vae.py
ImageNet loader dataset/dataset.py

Inference and evaluation

Inference entrypoint inference.py
FID and metric evaluation utils/fid_util.py
Inception features utils/jax_fid/inception.py
FID and PR metrics utils/jax_fid

Runtime support

Distributed sharding utils/hsdp_util.py
Run logging utils/logging.py

Other components

Engineer
ImageNet files
Hugging Face
W&B

Connections 32

Engineer → Run dispatcher: starts run README.md
Run dispatcher → Generator training: dispatches main.py
Run dispatcher → MAE training: dispatches main.py
Generator training → ImageNet loader: loads batches train.py
Generator training → Generator model: trains train.py
Generator training → Drifting loss: computes loss train.py
Generator training → Memory bank: uses train.py
Generator training → MAE feature model: uses features train.py
Generator training → Checkpoint utilities: saves state train.py
MAE training → ImageNet loader: loads batches train_mae.py
MAE training → MAE feature model: trains train_mae.py
MAE training → Checkpoint utilities: saves state train_mae.py
MAE feature model → ConvNeXt encoder: selects encoder models/mae_model.py
Generator training → Distributed sharding: shards training train.py
MAE training → Distributed sharding: shards training train_mae.py
Latent cache builder → ImageNet files: reads images README.md
Latent cache builder → VAE encode/decode: encodes batches dataset/latent.py
Latent cache builder → Distributed sharding: uses devices dataset/latent.py
Latent cache builder → ImageNet loader: writes cache for dataset/dataset.py
ImageNet loader → ImageNet files: loads split dataset/dataset.py
Inference entrypoint → Artifact loading: loads model inference.py
Artifact loading → HF artifact helpers: loads HF artifact utils/init_util.py
HF artifact helpers → Hugging Face: downloads from models/hf.py
Inference entrypoint → ImageNet loader: loads validation inference.py
Inference entrypoint → FID and metric evaluation: evaluates samples inference.py
FID and metric evaluation → Inception features: extracts features utils/fid_util.py
FID and metric evaluation → FID and PR metrics: computes metrics utils/fid_util.py
Inference entrypoint → Generator model: generates samples inference.py
Inference entrypoint → Run logging: reports results inference.py
Run logging → W&B: logs optionally README.md
Generator training → Run logging: logs progress train.py
MAE training → Run logging: logs progress train_mae.py

Drawn by GitDiagram on November 11, 2026 from the repository's file tree, README and source files. Read as Markdown ·