feat(images): add matched JiT latent-denoise training

external/JiT/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Tianhong Li
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

external/JiT/README.md

@@ -0,0 +1,150 @@
+## Just image Transformer (JiT) for Pixel-space Diffusion
+
+[![arXiv](https://img.shields.io/badge/arXiv%20paper-2511.13720-b31b1b.svg)](https://arxiv.org/abs/2511.13720) 
+
+<p align="center">
+  <img src="demo/visual.jpg" width="100%">
+</p>
+
+
+This is a PyTorch/GPU re-implementation of the paper [Back to Basics: Let Denoising Generative Models Denoise](https://arxiv.org/abs/2511.13720):
+
+```
+@article{li2025jit,
+  title={Back to Basics: Let Denoising Generative Models Denoise},
+  author={Li, Tianhong and He, Kaiming},
+  journal={arXiv preprint arXiv:2511.13720},
+  year={2025}
+}
+```
+
+JiT adopts a minimalist and self-contained design for pixel-level high-resolution image diffusion. 
+The original implementation was in JAX+TPU. This re-implementation is in PyTorch+GPU.
+
+<p align="center">
+  <img src="demo/jit.jpg" width="40%">
+</p>
+
+### Dataset
+Download [ImageNet](http://image-net.org/download) dataset, and place it in your `IMAGENET_PATH`.
+
+### Installation
+
+Download the code:
+```
+git clone https://github.com/LTH14/JiT.git
+cd JiT
+```
+
+A suitable [conda](https://conda.io/) environment named `jit` can be created and activated with:
+
+```
+conda env create -f environment.yaml
+conda activate jit
+```
+
+If you get ```undefined symbol: iJIT_NotifyEvent``` when importing ```torch```, simply
+```
+pip uninstall torch
+pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu124
+```
+Check this [issue](https://github.com/conda/conda/issues/13812#issuecomment-2071445372) for more details.
+
+### Training
+The below training scripts have been tested on 8 H200 GPUs.
+
+Example script for training JiT-B/16 on ImageNet 256x256 for 600 epochs:
+```
+torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
+main_jit.py \
+--model JiT-B/16 \
+--proj_dropout 0.0 \
+--P_mean -0.8 --P_std 0.8 \
+--img_size 256 --noise_scale 1.0 \
+--batch_size 128 --blr 5e-5 \
+--epochs 600 --warmup_epochs 5 \
+--gen_bsz 128 --num_images 50000 --cfg 2.9 --interval_min 0.1 --interval_max 1.0 \
+--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
+--data_path ${IMAGENET_PATH} --online_eval
+```
+
+Example script for training JiT-B/32 on ImageNet 512x512 for 600 epochs:
+```
+torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
+main_jit.py \
+--model JiT-B/32 \
+--proj_dropout 0.0 \
+--P_mean -0.8 --P_std 0.8 \
+--img_size 512 --noise_scale 2.0 \
+--batch_size 128 --blr 5e-5 \
+--epochs 600 --warmup_epochs 5 \
+--gen_bsz 128 --num_images 50000 --cfg 2.9 --interval_min 0.1 --interval_max 1.0 \
+--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
+--data_path ${IMAGENET_PATH} --online_eval
+```
+
+Example script for training JiT-H/16 on ImageNet 256x256 for 600 epochs:
+```
+torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
+main_jit.py \
+--model JiT-H/16 \
+--proj_dropout 0.2 \
+--P_mean -0.8 --P_std 0.8 \
+--img_size 256 --noise_scale 1.0 \
+--batch_size 128 --blr 5e-5 \
+--epochs 600 --warmup_epochs 5 \
+--gen_bsz 128 --num_images 50000 --cfg 2.2 --interval_min 0.1 --interval_max 1.0 \
+--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
+--data_path ${IMAGENET_PATH} --online_eval
+```
+
+### Evaluation
+
+PyTorch pre-trained models are available [here](https://www.dropbox.com/scl/fo/3ken1avtsd81ip67b9qpi/AK218ZNvXKSv74igVvht4PQ?rlkey=14gjrblmljewpl6ygxzlr3njm&st=ffkl77al&dl=0).
+
+Evaluate pre-trained JiT-B:
+```
+torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
+main_jit.py \
+--model JiT-B/16 (or JiT-B/32) \
+--img_size 256 (or 512) --noise_scale 1.0 (or 2.0) \
+--gen_bsz 256 --num_images 50000 --cfg 3.0 --interval_min 0.1 --interval_max 1.0 \
+--output_dir ${CKPT_DIR} --resume ${CKPT_DIR} \
+--data_path ${IMAGENET_PATH} --evaluate_gen
+```
+
+Evaluate pre-trained JiT-L:
+```
+torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
+main_jit.py \
+--model JiT-L/16 (or JiT-L/32) \
+--img_size 256 (or 512) --noise_scale 1.0 (or 2.0) \
+--gen_bsz 256 --num_images 50000 --cfg 2.4 (or 2.5) --interval_min 0.1 --interval_max 1.0 \
+--output_dir ${CKPT_DIR} --resume ${CKPT_DIR} \
+--data_path ${IMAGENET_PATH} --evaluate_gen
+```
+
+Evaluate pre-trained JiT-H:
+```
+torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
+main_jit.py \
+--model JiT-H/16 (or JiT-H/32) \
+--img_size 256 (or 512) --noise_scale 1.0 (or 2.0) \
+--gen_bsz 256 --num_images 50000 --cfg 2.2 (or 2.3) --interval_min 0.1 --interval_max 1.0 \
+--output_dir ${CKPT_DIR} --resume ${CKPT_DIR} \
+--data_path ${IMAGENET_PATH} --evaluate_gen
+```
+
+We use a customized [```torch-fidelity```](https://github.com/LTH14/torch-fidelity)
+to evaluate FID and IS against a reference image folder or statistics. You can use ```prepare_ref.py```
+to prepare the reference image folder, or directly use our pre-computed reference stats
+under ```fid_stats```.
+
+### Acknowledgements
+
+We thank Google TPU Research Cloud (TRC) for granting us access to TPUs, and the MIT
+ORCD Seed Fund Grants for supporting GPU resources.
+
+### Contact
+
+If you have any questions, feel free to contact me through email (tianhong@mit.edu). Enjoy!

external/JiT/UPSTREAM.md

@@ -0,0 +1,11 @@
+# Upstream provenance
+
+This directory vendors the training-relevant files from the MIT-licensed
+PyTorch JiT implementation from
+https://github.com/LTH14/JiT at commit
+`cbc743a2ada5e9762697da2c83f8c4f8379e8c17`.
+
+The 512px FID statistics and demo images are omitted. The only upstream
+compatibility change is device-agnostic rotary buffers in
+`util/model_util.py`. The matched experiment lives in `matched_models.py` and
+`train_matched.py`; the original JiT model and objective remain intact.

external/JiT/denoiser.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn as nn
+from model_jit import JiT_models
+
+
+class Denoiser(nn.Module):
+    def __init__(
+        self,
+        args
+    ):
+        super().__init__()
+        self.net = JiT_models[args.model](
+            input_size=args.img_size,
+            in_channels=3,
+            num_classes=args.class_num,
+            attn_drop=args.attn_dropout,
+            proj_drop=args.proj_dropout,
+        )
+        self.img_size = args.img_size
+        self.num_classes = args.class_num
+
+        self.label_drop_prob = args.label_drop_prob
+        self.P_mean = args.P_mean
+        self.P_std = args.P_std
+        self.t_eps = args.t_eps
+        self.noise_scale = args.noise_scale
+
+        # ema
+        self.ema_decay1 = args.ema_decay1
+        self.ema_decay2 = args.ema_decay2
+        self.ema_params1 = None
+        self.ema_params2 = None
+
+        # generation hyper params
+        self.method = args.sampling_method
+        self.steps = args.num_sampling_steps
+        self.cfg_scale = args.cfg
+        self.cfg_interval = (args.interval_min, args.interval_max)
+
+    def drop_labels(self, labels):
+        drop = torch.rand(labels.shape[0], device=labels.device) < self.label_drop_prob
+        out = torch.where(drop, torch.full_like(labels, self.num_classes), labels)
+        return out
+
+    def sample_t(self, n: int, device=None):
+        z = torch.randn(n, device=device) * self.P_std + self.P_mean
+        return torch.sigmoid(z)
+
+    def forward(self, x, labels):
+        labels_dropped = self.drop_labels(labels) if self.training else labels
+
+        t = self.sample_t(x.size(0), device=x.device).view(-1, *([1] * (x.ndim - 1)))
+        e = torch.randn_like(x) * self.noise_scale
+
+        z = t * x + (1 - t) * e
+        v = (x - z) / (1 - t).clamp_min(self.t_eps)
+
+        x_pred = self.net(z, t.flatten(), labels_dropped)
+        v_pred = (x_pred - z) / (1 - t).clamp_min(self.t_eps)
+
+        # l2 loss
+        loss = (v - v_pred) ** 2
+        loss = loss.mean(dim=(1, 2, 3)).mean()
+
+        return loss
+
+    @torch.no_grad()
+    def generate(self, labels):
+        device = labels.device
+        bsz = labels.size(0)
+        z = self.noise_scale * torch.randn(bsz, 3, self.img_size, self.img_size, device=device)
+        timesteps = torch.linspace(0.0, 1.0, self.steps+1, device=device).view(-1, *([1] * z.ndim)).expand(-1, bsz, -1, -1, -1)
+
+        if self.method == "euler":
+            stepper = self._euler_step
+        elif self.method == "heun":
+            stepper = self._heun_step
+        else:
+            raise NotImplementedError
+
+        # ode
+        for i in range(self.steps - 1):
+            t = timesteps[i]
+            t_next = timesteps[i + 1]
+            z = stepper(z, t, t_next, labels)
+        # last step euler
+        z = self._euler_step(z, timesteps[-2], timesteps[-1], labels)
+        return z
+
+    @torch.no_grad()
+    def _forward_sample(self, z, t, labels):
+        # conditional
+        x_cond = self.net(z, t.flatten(), labels)
+        v_cond = (x_cond - z) / (1.0 - t).clamp_min(self.t_eps)
+
+        # unconditional
+        x_uncond = self.net(z, t.flatten(), torch.full_like(labels, self.num_classes))
+        v_uncond = (x_uncond - z) / (1.0 - t).clamp_min(self.t_eps)
+
+        # cfg interval
+        low, high = self.cfg_interval
+        interval_mask = (t < high) & ((low == 0) | (t > low))
+        cfg_scale_interval = torch.where(interval_mask, self.cfg_scale, 1.0)
+
+        return v_uncond + cfg_scale_interval * (v_cond - v_uncond)
+
+    @torch.no_grad()
+    def _euler_step(self, z, t, t_next, labels):
+        v_pred = self._forward_sample(z, t, labels)
+        z_next = z + (t_next - t) * v_pred
+        return z_next
+
+    @torch.no_grad()
+    def _heun_step(self, z, t, t_next, labels):
+        v_pred_t = self._forward_sample(z, t, labels)
+
+        z_next_euler = z + (t_next - t) * v_pred_t
+        v_pred_t_next = self._forward_sample(z_next_euler, t_next, labels)
+
+        v_pred = 0.5 * (v_pred_t + v_pred_t_next)
+        z_next = z + (t_next - t) * v_pred
+        return z_next
+
+    @torch.no_grad()
+    def update_ema(self):
+        source_params = list(self.parameters())
+        for targ, src in zip(self.ema_params1, source_params):
+            targ.detach().mul_(self.ema_decay1).add_(src, alpha=1 - self.ema_decay1)
+        for targ, src in zip(self.ema_params2, source_params):
+            targ.detach().mul_(self.ema_decay2).add_(src, alpha=1 - self.ema_decay2)