How lth14/jit works
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 Train and evaluate · main_jit.py Training runner · denoiser.py Denoiser and sampler · model_jit.py JiT transformer
Components 14
Training and orchestration
- Training runner
main_jit.py - Train and evaluate
engine_jit.py - Training utilities
util/misc.py - Learning-rate schedule
util/lr_sched.py
Diffusion model
- Denoiser and sampler
denoiser.py - JiT transformer
model_jit.py - Position and norm helpers
util/model_util.py
Image data
- Reference-image preparation
prepare_ref.py - Image center crop
util/crop.py
Generation and evaluation
- Generated images
- FID reference statistics
fid_stats - FID and IS metrics
Other components
- Researcher
- ImageNet
Connections 15
- 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