JiT is an open-source PyTorch implementation of a state-of-the-art image diffusion model designed around a minimalist yet powerful architecture for pixel-level generative modeling, based on the paper Back to Basics: Let Denoising Generative Models Denoise. Rather than predicting noise, JiT models directly predict clean image data, which the research suggests aligns better with the manifold structure of natural images and leads to stronger generative performance at high resolution. This implementation supports training on large datasets like ImageNet with configurable model variants, and practical scripts for setup, training, and evaluation on GPUs are included, leveraging PyTorch’s ecosystem for real-world experimentation. The repository’s layout contains modular engine, model, and training scripts enabling researchers and engineers to customize components such as training regimes, noise schedules, and evaluation routines.
Features
- PyTorch implementation of the JiT diffusion model
- Direct clean image prediction rather than noise prediction
- Training scripts with multi-GPU support
- Modular model and training engine files
- Community engagement with issues and discussions
- Suitable for high-resolution generative workflows