# lambertae/drifting architecture

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

- Interactive diagram: https://gitdiagram.com/lambertae/drifting
- Repository: https://github.com/lambertae/drifting

## Overview

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`](https://github.com/lambertae/drifting/blob/main/train.py) (Generator training), [`train_mae.py`](https://github.com/lambertae/drifting/blob/main/train_mae.py) (MAE training), [`inference.py`](https://github.com/lambertae/drifting/blob/main/inference.py) (Inference entrypoint), [`dataset/dataset.py`](https://github.com/lambertae/drifting/blob/main/dataset/dataset.py) (ImageNet loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_workflow["Training workflows"]
  node_main["Run dispatcher<br/>[main.py]"]
  node_generator_training["Generator training<br/>[train.py]"]
  node_mae_training["MAE training<br/>[train_mae.py]"]
  node_drift_loss["Drifting loss<br/>[drift_loss.py]"]
  node_memory_bank["Memory bank<br/>[memory_bank.py]"]
end

subgraph group_models["Models and artifacts"]
  node_generator["Generator model<br/>[generator.py]"]
  node_mae_model["MAE feature model<br/>[mae_model.py]"]
  node_convnext["ConvNeXt encoder<br/>[convnext.py]"]
  node_artifact_loader["Artifact loading<br/>[init_util.py]"]
  node_hf_artifacts["HF artifact helpers<br/>[hf.py]"]
  node_checkpointing["Checkpoint utilities<br/>[ckpt_util.py]"]
end

subgraph group_data["Data preparation"]
  node_latent_cache["Latent cache builder<br/>[latent.py]"]
  node_vae["VAE encode/decode<br/>[vae.py]"]
  node_imagenet_data["ImageNet loader<br/>[dataset.py]"]
end

subgraph group_evaluation["Inference and evaluation"]
  node_inference["Inference entrypoint<br/>[inference.py]"]
  node_fid_eval["FID and metric evaluation<br/>[fid_util.py]"]
  node_inception["Inception features<br/>[inception.py]"]
  node_fid_metrics["FID and PR metrics"]
end

subgraph group_runtime["Runtime support"]
  node_distribution["Distributed sharding<br/>[hsdp_util.py]"]
  node_logging["Run logging<br/>[logging.py]"]
end

node_user(("Engineer"))
node_imagenet[("ImageNet files")]
node_huggingface["Hugging Face"]
node_wandb["W&amp;B"]

node_user -->|"starts run"| node_main
node_main -->|"dispatches"| node_generator_training
node_main -->|"dispatches"| node_mae_training
node_generator_training -->|"loads batches"| node_imagenet_data
node_generator_training -->|"trains"| node_generator
node_generator_training -->|"computes loss"| node_drift_loss
node_generator_training -->|"uses"| node_memory_bank
node_generator_training -->|"uses features"| node_mae_model
node_generator_training -->|"saves state"| node_checkpointing
node_mae_training -->|"loads batches"| node_imagenet_data
node_mae_training -->|"trains"| node_mae_model
node_mae_training -->|"saves state"| node_checkpointing
node_mae_model -->|"selects encoder"| node_convnext
node_generator_training -->|"shards training"| node_distribution
node_mae_training -->|"shards training"| node_distribution
node_latent_cache -->|"reads images"| node_imagenet
node_latent_cache -->|"encodes batches"| node_vae
node_latent_cache -->|"uses devices"| node_distribution
node_latent_cache -.->|"writes cache for"| node_imagenet_data
node_imagenet_data -->|"loads split"| node_imagenet
node_inference -->|"loads model"| node_artifact_loader
node_artifact_loader -->|"loads HF artifact"| node_hf_artifacts
node_hf_artifacts -->|"downloads from"| node_huggingface
node_inference -->|"loads validation"| node_imagenet_data
node_inference -->|"evaluates samples"| node_fid_eval
node_fid_eval -->|"extracts features"| node_inception
node_fid_eval -->|"computes metrics"| node_fid_metrics
node_inference -->|"generates samples"| node_generator
node_inference -->|"reports results"| node_logging
node_logging -.->|"logs optionally"| node_wandb
node_generator_training -->|"logs progress"| node_logging
node_mae_training -->|"logs progress"| node_logging

click node_main "https://github.com/lambertae/drifting/blob/main/main.py"
click node_generator_training "https://github.com/lambertae/drifting/blob/main/train.py"
click node_mae_training "https://github.com/lambertae/drifting/blob/main/train_mae.py"
click node_drift_loss "https://github.com/lambertae/drifting/blob/main/drift_loss.py"
click node_memory_bank "https://github.com/lambertae/drifting/blob/main/memory_bank.py"
click node_generator "https://github.com/lambertae/drifting/blob/main/models/generator.py"
click node_mae_model "https://github.com/lambertae/drifting/blob/main/models/mae_model.py"
click node_convnext "https://github.com/lambertae/drifting/blob/main/models/convnext.py"
click node_artifact_loader "https://github.com/lambertae/drifting/blob/main/utils/init_util.py"
click node_hf_artifacts "https://github.com/lambertae/drifting/blob/main/models/hf.py"
click node_checkpointing "https://github.com/lambertae/drifting/blob/main/utils/ckpt_util.py"
click node_latent_cache "https://github.com/lambertae/drifting/blob/main/dataset/latent.py"
click node_vae "https://github.com/lambertae/drifting/blob/main/dataset/vae.py"
click node_imagenet_data "https://github.com/lambertae/drifting/blob/main/dataset/dataset.py"
click node_inference "https://github.com/lambertae/drifting/blob/main/inference.py"
click node_fid_eval "https://github.com/lambertae/drifting/blob/main/utils/fid_util.py"
click node_inception "https://github.com/lambertae/drifting/blob/main/utils/jax_fid/inception.py"
click node_fid_metrics "https://github.com/lambertae/drifting/tree/main/utils/jax_fid"
click node_distribution "https://github.com/lambertae/drifting/blob/main/utils/hsdp_util.py"
click node_logging "https://github.com/lambertae/drifting/blob/main/utils/logging.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_main,node_generator_training,node_mae_training,node_drift_loss,node_memory_bank toneBlue
class node_generator,node_mae_model,node_convnext,node_artifact_loader,node_hf_artifacts,node_checkpointing,node_imagenet toneAmber
class node_latent_cache,node_vae,node_imagenet_data toneMint
class node_inference,node_fid_eval,node_inception,node_fid_metrics toneRose
class node_distribution,node_logging,node_user,node_huggingface,node_wandb toneIndigo
```

## Components

### Training workflows

- [Run dispatcher](https://github.com/lambertae/drifting/blob/main/main.py): `main.py`
- [Generator training](https://github.com/lambertae/drifting/blob/main/train.py): `train.py`
- [MAE training](https://github.com/lambertae/drifting/blob/main/train_mae.py): `train_mae.py`
- [Drifting loss](https://github.com/lambertae/drifting/blob/main/drift_loss.py): `drift_loss.py`
- [Memory bank](https://github.com/lambertae/drifting/blob/main/memory_bank.py): `memory_bank.py`

### Models and artifacts

- [Generator model](https://github.com/lambertae/drifting/blob/main/models/generator.py): `models/generator.py`
- [MAE feature model](https://github.com/lambertae/drifting/blob/main/models/mae_model.py): `models/mae_model.py`
- [ConvNeXt encoder](https://github.com/lambertae/drifting/blob/main/models/convnext.py): `models/convnext.py`
- [Artifact loading](https://github.com/lambertae/drifting/blob/main/utils/init_util.py): `utils/init_util.py`
- [HF artifact helpers](https://github.com/lambertae/drifting/blob/main/models/hf.py): `models/hf.py`
- [Checkpoint utilities](https://github.com/lambertae/drifting/blob/main/utils/ckpt_util.py): `utils/ckpt_util.py`

### Data preparation

- [Latent cache builder](https://github.com/lambertae/drifting/blob/main/dataset/latent.py): `dataset/latent.py`
- [VAE encode/decode](https://github.com/lambertae/drifting/blob/main/dataset/vae.py): `dataset/vae.py`
- [ImageNet loader](https://github.com/lambertae/drifting/blob/main/dataset/dataset.py): `dataset/dataset.py`

### Inference and evaluation

- [Inference entrypoint](https://github.com/lambertae/drifting/blob/main/inference.py): `inference.py`
- [FID and metric evaluation](https://github.com/lambertae/drifting/blob/main/utils/fid_util.py): `utils/fid_util.py`
- [Inception features](https://github.com/lambertae/drifting/blob/main/utils/jax_fid/inception.py): `utils/jax_fid/inception.py`
- [FID and PR metrics](https://github.com/lambertae/drifting/tree/main/utils/jax_fid): `utils/jax_fid`

### Runtime support

- [Distributed sharding](https://github.com/lambertae/drifting/blob/main/utils/hsdp_util.py): `utils/hsdp_util.py`
- [Run logging](https://github.com/lambertae/drifting/blob/main/utils/logging.py): `utils/logging.py`

### Other components

- **Engineer**
- **ImageNet files**
- **Hugging Face**
- **W&B**

## Connections

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