Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+demo.png filter=lfs diff=lfs merge=lfs -text
+demo_images/jit_h32_test_inference.png filter=lfs diff=lfs merge=lfs -text

JiT-B-16/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "JiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "attention_dropout": 0.0,
+  "attn_dropout": null,
+  "dropout": 0.0,
+  "image_size": null,
+  "model_name": null,
+  "model_type": "JiT-B/16",
+  "num_class_embeds": 1000,
+  "num_classes": null,
+  "proj_dropout": null,
+  "sample_size": 256
+}

JiT-B-16/conversion_metadata.json

@@ -0,0 +1,17 @@
+{
+  "source_checkpoint": "/root/worksapce/models/raw/jit-b-16/checkpoint-last.pth",
+  "weights": "ema1",
+  "epoch": null,
+  "jit_args": {
+    "model_type": "JiT-B/16",
+    "sample_size": 256,
+    "num_class_embeds": 1000,
+    "attention_dropout": 0.0,
+    "dropout": 0.0,
+    "model": "JiT-B/16",
+    "img_size": 256,
+    "class_num": 1000,
+    "attn_dropout": 0.0,
+    "proj_dropout": 0.0
+  }
+}

JiT-B-16/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53987d6142a3dd9b0a390bd987ebecf806f95cd7b52d64beb41a4fb21199fb6b
+size 525301096

JiT-B-16/model_index.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "JiTPipeline",
+  "_diffusers_version": "0.36.0",
+  "transformer": [
+    "jit_diffusers",
+    "JiTTransformer2DModel"
+  ]
+}

JiT-B-16/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "JiTScheduler",
+  "_diffusers_version": "0.36.0",
+  "solver": "heun",
+  "timestep_start": 0.0,
+  "timestep_end": 1.0
+}

JiT-B-32/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "JiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "attention_dropout": 0.0,
+  "attn_dropout": null,
+  "dropout": 0.0,
+  "image_size": null,
+  "model_name": null,
+  "model_type": "JiT-B/32",
+  "num_class_embeds": 1000,
+  "num_classes": null,
+  "proj_dropout": null,
+  "sample_size": 512
+}

JiT-B-32/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3652a94f7815bd4b2a02575a4b004e4270a4ddbf5218c292c2cac6d16efb7399
+size 533567848

JiT-B-32/model_index.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "JiTPipeline",
+  "_diffusers_version": "0.36.0",
+  "transformer": [
+    "jit_diffusers",
+    "JiTTransformer2DModel"
+  ]
+}

JiT-B-32/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "JiTScheduler",
+  "_diffusers_version": "0.36.0",
+  "solver": "heun",
+  "timestep_start": 0.0,
+  "timestep_end": 1.0
+}

JiT-H-16/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "JiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "attention_dropout": 0.0,
+  "attn_dropout": null,
+  "dropout": 0.2,
+  "image_size": null,
+  "model_name": null,
+  "model_type": "JiT-H/16",
+  "num_class_embeds": 1000,
+  "num_classes": null,
+  "proj_dropout": null,
+  "sample_size": 256
+}

JiT-H-16/conversion_metadata.json

@@ -0,0 +1,17 @@
+{
+  "source_checkpoint": "/root/worksapce/models/raw/jit-h-16/checkpoint-last.pth",
+  "weights": "ema1",
+  "epoch": null,
+  "jit_args": {
+    "model_type": "JiT-H/16",
+    "sample_size": 256,
+    "num_class_embeds": 1000,
+    "attention_dropout": 0.0,
+    "dropout": 0.2,
+    "model": "JiT-H/16",
+    "img_size": 256,
+    "class_num": 1000,
+    "attn_dropout": 0.0,
+    "proj_dropout": 0.2
+  }
+}

JiT-H-16/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2723e68bceac95b153a01d3cc17cf10e727e8bfe7cdf3c38ee2f66adaa6e7b35
+size 3811419656

JiT-H-16/model_index.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "JiTPipeline",
+  "_diffusers_version": "0.36.0",
+  "transformer": [
+    "jit_diffusers",
+    "JiTTransformer2DModel"
+  ]
+}

JiT-H-16/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "JiTScheduler",
+  "_diffusers_version": "0.36.0",
+  "solver": "heun",
+  "timestep_start": 0.0,
+  "timestep_end": 1.0
+}

JiT-H-32/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "JiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "attention_dropout": 0.0,
+  "attn_dropout": null,
+  "dropout": 0.2,
+  "image_size": null,
+  "model_name": null,
+  "model_type": "JiT-H/32",
+  "num_class_embeds": 1000,
+  "num_classes": null,
+  "proj_dropout": null,
+  "sample_size": 512
+}

JiT-H-32/conversion_metadata.json

@@ -0,0 +1,17 @@
+{
+  "source_checkpoint": "/root/worksapce/projects/JiT-diffusers/models/raw/jit-h-32/checkpoint-last.pth",
+  "weights": "ema1",
+  "epoch": null,
+  "jit_args": {
+    "model_type": "JiT-H/32",
+    "sample_size": 512,
+    "num_class_embeds": 1000,
+    "attention_dropout": 0.0,
+    "dropout": 0.2,
+    "model": "JiT-H/32",
+    "img_size": 512,
+    "class_num": 1000,
+    "attn_dropout": 0.0,
+    "proj_dropout": 0.2
+  }
+}

JiT-H-32/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3693c9077ce262edf99fccbeff8e981ae1f7499a6d6148c0837081fe3dee15db
+size 3825584616

JiT-H-32/model_index.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "JiTPipeline",
+  "_diffusers_version": "0.36.0",
+  "transformer": [
+    "jit_diffusers",
+    "JiTTransformer2DModel"
+  ]
+}

JiT-H-32/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "JiTScheduler",
+  "_diffusers_version": "0.36.0",
+  "solver": "heun",
+  "timestep_start": 0.0,
+  "timestep_end": 1.0
+}

JiT-L-16/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "JiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "attention_dropout": 0.0,
+  "attn_dropout": null,
+  "dropout": 0.0,
+  "image_size": null,
+  "model_name": null,
+  "model_type": "JiT-L/16",
+  "num_class_embeds": 1000,
+  "num_classes": null,
+  "proj_dropout": null,
+  "sample_size": 256
+}

JiT-L-16/conversion_metadata.json

@@ -0,0 +1,17 @@
+{
+  "source_checkpoint": "/root/worksapce/models/raw/jit-l-16/checkpoint-last.pth",
+  "weights": "ema1",
+  "epoch": null,
+  "jit_args": {
+    "model_type": "JiT-L/16",
+    "sample_size": 256,
+    "num_class_embeds": 1000,
+    "attention_dropout": 0.0,
+    "dropout": 0.0,
+    "model": "JiT-L/16",
+    "img_size": 256,
+    "class_num": 1000,
+    "attn_dropout": 0.0,
+    "proj_dropout": 0.0
+  }
+}

JiT-L-16/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8d49cacb0751f4fc77d6578d71f992d65323b451fa194e35c77ebc49e21b9cbd
+size 1836598016

JiT-L-16/model_index.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "JiTPipeline",
+  "_diffusers_version": "0.36.0",
+  "transformer": [
+    "jit_diffusers",
+    "JiTTransformer2DModel"
+  ]
+}

JiT-L-16/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "JiTScheduler",
+  "_diffusers_version": "0.36.0",
+  "solver": "heun",
+  "timestep_start": 0.0,
+  "timestep_end": 1.0
+}

JiT-L-32/config.json

@@ -0,0 +1,14 @@
+{
+  "_class_name": "JiTTransformer2DModel",
+  "_diffusers_version": "0.36.0",
+  "attention_dropout": 0.0,
+  "attn_dropout": null,
+  "dropout": 0.0,
+  "image_size": null,
+  "model_name": null,
+  "model_type": "JiT-L/32",
+  "num_class_embeds": 1000,
+  "num_classes": null,
+  "proj_dropout": null,
+  "sample_size": 512
+}

JiT-L-32/conversion_metadata.json

@@ -0,0 +1,17 @@
+{
+  "source_checkpoint": "/root/worksapce/models/raw/jit-l-32/checkpoint-last.pth",
+  "weights": "ema1",
+  "epoch": null,
+  "jit_args": {
+    "model_type": "JiT-L/32",
+    "sample_size": 512,
+    "num_class_embeds": 1000,
+    "attention_dropout": 0.0,
+    "dropout": 0.0,
+    "model": "JiT-L/32",
+    "img_size": 512,
+    "class_num": 1000,
+    "attn_dropout": 0.0,
+    "proj_dropout": 0.0
+  }
+}

JiT-L-32/diffusion_pytorch_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f46ad7ed33a27906245178451c724e40c33ad13832fd82e7ab751e9620e471e0
+size 1847224072

JiT-L-32/model_index.json

@@ -0,0 +1,8 @@
+{
+  "_class_name": "JiTPipeline",
+  "_diffusers_version": "0.36.0",
+  "transformer": [
+    "jit_diffusers",
+    "JiTTransformer2DModel"
+  ]
+}

JiT-L-32/scheduler_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "JiTScheduler",
+  "_diffusers_version": "0.36.0",
+  "solver": "heun",
+  "timestep_start": 0.0,
+  "timestep_end": 1.0
+}

README.md

@@ -0,0 +1,89 @@
+---
+license: mit
+library_name: diffusers
+pipeline_tag: unconditional-image-generation
+tags:
+  - diffusers
+  - jit
+  - image-generation
+  - class-conditional
+widget:
+  - text: "ImageNet class 207 (golden retriever), JiT-H/32 test sample"
+    output:
+      url: demo.png
+---
+
+# JiT-H/32 (Diffusers)
+
+This repository is self-contained: model weights and a custom `diffusers` pipeline (`JiTPipeline`) are both included, so no external code repo is required.
+
+## Available Checkpoints (All 6)
+
+The JiT paper reports six ImageNet checkpoints across 256 and 512 resolutions. Use the following relative paths with `JiTPipeline.from_pretrained(...)`.
+
+| Checkpoint | Relative path | Resolution | Pre-trained dataset | Recommended CFG | Recommended interval | Recommended noise_scale | FID-50K |
+|---|---|---|---|---:|---|---:|---:|
+| JiT-B/16 | `./JiT-B-16` | 256x256 | ImageNet 256x256 | 3.0 | `[0.1, 1.0]` | 1.0 | 3.66 |
+| JiT-L/16 | `./JiT-L-16` | 256x256 | ImageNet 256x256 | 2.4 | `[0.1, 1.0]` | 1.0 | 2.36 |
+| JiT-H/16 | `./JiT-H-16` | 256x256 | ImageNet 256x256 | 2.2 | `[0.1, 1.0]` | 1.0 | 1.86 |
+| JiT-B/32 | `./JiT-B-32` | 512x512 | ImageNet 512x512 | 3.0 | `[0.1, 1.0]` | 2.0 | 4.02 |
+| JiT-L/32 | `./JiT-L-32` | 512x512 | ImageNet 512x512 | 2.5 | `[0.1, 1.0]` | 2.0 | 2.53 |
+| JiT-H/32 | `./JiT-H-32` | 512x512 | ImageNet 512x512 | 2.3 | `[0.1, 1.0]` | 2.0 | 1.94 |
+
+Source: [Back to Basics: Let Denoising Generative Models Denoise (arXiv:2511.13720)](https://arxiv.org/html/2511.13720).
+
+## Demo Image
+
+![JiT-H/32 test inference](demo_images/jit_h32_test_inference.png)
+
+## One-Stop Diffusers Inference
+
+```python
+from pathlib import Path
+import sys
+import torch
+
+repo_dir = Path(".").resolve()
+sys.path.insert(0, str(repo_dir))
+from jit_diffusers import JiTPipeline
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+pipe = JiTPipeline.from_pretrained("./JiT-H-32").to(device)
+pipe.transformer = pipe.transformer.to(device=device, dtype=torch.bfloat16 if device == "cuda" else torch.float32)
+pipe.transformer.eval()
+
+generator = torch.Generator(device=device).manual_seed(42)
+output = pipe(
+    class_labels=[207],
+    num_inference_steps=50,
+    guidance_scale=2.3,
+    guidance_interval_min=0.1,
+    guidance_interval_max=1.0,
+    noise_scale=2.0,
+    t_eps=5e-2,
+    sampling_method="heun",
+    generator=generator,
+    output_type="pil",
+)
+image = output.images[0]
+output_path = Path("./demo_images/jit_h32_test_inference.png")
+output_path.parent.mkdir(parents=True, exist_ok=True)
+image.save(output_path)
+print(f"Saved image to: {output_path}")
+```
+
+## Ready-to-Run Commands (All 6 Checkpoints)
+
+Run these from this repository root (`models/BiliSakura/JiT-diffusers`).
+
+```bash
+# 256x256 checkpoints
+python run_jit_diffusers_inference.py --model_path ./JiT-B-16 --output_path ./demo_images/jit_b16_test_inference.png --class_label 207 --seed 42 --steps 50 --cfg 3.0 --interval_min 0.1 --interval_max 1.0 --noise_scale 1.0 --t_eps 5e-2 --solver heun
+python run_jit_diffusers_inference.py --model_path ./JiT-L-16 --output_path ./demo_images/jit_l16_test_inference.png --class_label 207 --seed 42 --steps 50 --cfg 2.4 --interval_min 0.1 --interval_max 1.0 --noise_scale 1.0 --t_eps 5e-2 --solver heun
+python run_jit_diffusers_inference.py --model_path ./JiT-H-16 --output_path ./demo_images/jit_h16_test_inference.png --class_label 207 --seed 42 --steps 50 --cfg 2.2 --interval_min 0.1 --interval_max 1.0 --noise_scale 1.0 --t_eps 5e-2 --solver heun
+
+# 512x512 checkpoints
+python run_jit_diffusers_inference.py --model_path ./JiT-B-32 --output_path ./demo_images/jit_b32_test_inference.png --class_label 207 --seed 42 --steps 50 --cfg 3.0 --interval_min 0.1 --interval_max 1.0 --noise_scale 2.0 --t_eps 5e-2 --solver heun
+python run_jit_diffusers_inference.py --model_path ./JiT-L-32 --output_path ./demo_images/jit_l32_test_inference.png --class_label 207 --seed 42 --steps 50 --cfg 2.5 --interval_min 0.1 --interval_max 1.0 --noise_scale 2.0 --t_eps 5e-2 --solver heun
+python run_jit_diffusers_inference.py --model_path ./JiT-H-32 --output_path ./demo_images/jit_h32_test_inference.png --class_label 207 --seed 42 --steps 50 --cfg 2.3 --interval_min 0.1 --interval_max 1.0 --noise_scale 2.0 --t_eps 5e-2 --solver heun
+```

jit_diffusers/__init__.py

@@ -0,0 +1,11 @@
+from .modeling_jit_transformer_2d import JiTTransformer2DModel, JiTDiffusersModel
+from .pipeline_jit import JiTPipeline, JiTPipelineOutput
+from .scheduling_jit import JiTScheduler
+
+__all__ = [
+    "JiTTransformer2DModel",
+    "JiTDiffusersModel",
+    "JiTPipeline",
+    "JiTPipelineOutput",
+    "JiTScheduler",
+]

jit_diffusers/modeling_jit_backbone.py

@@ -0,0 +1,321 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .modeling_jit_utils import VisionRotaryEmbeddingFast, get_2d_sincos_pos_embed, RMSNorm
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class BottleneckPatchEmbed(nn.Module):
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, pca_dim=768, embed_dim=768, bias=True):
+        super().__init__()
+        img_size = (img_size, img_size)
+        patch_size = (patch_size, patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj1 = nn.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = nn.Conv2d(pca_dim, embed_dim, kernel_size=1, stride=1, bias=bias)
+
+    def forward(self, x):
+        _, _, height, width = x.shape
+        assert height == self.img_size[0] and width == self.img_size[1], (
+            f"Input image size ({height}*{width}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        )
+        x = self.proj2(self.proj1(x)).flatten(2).transpose(1, 2)
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
+            device=t.device
+        )
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_freq = t_freq.to(dtype=self.mlp[0].weight.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class LabelEmbedder(nn.Module):
+    def __init__(self, num_classes, hidden_size):
+        super().__init__()
+        self.embedding_table = nn.Embedding(num_classes + 1, hidden_size)
+        self.num_classes = num_classes
+
+    def forward(self, labels):
+        embeddings = self.embedding_table(labels)
+        return embeddings
+
+
+def scaled_dot_product_attention(query, key, value, dropout_p=0.0) -> torch.Tensor:
+    query_len, key_len = query.size(-2), key.size(-2)
+    scale_factor = 1 / math.sqrt(query.size(-1))
+    attn_bias = torch.zeros(query.size(0), 1, query_len, key_len, dtype=query.dtype, device=query.device)
+
+    with torch.amp.autocast("cuda", enabled=False):
+        attn_weight = query.float() @ key.float().transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    attn_weight = torch.dropout(attn_weight, dropout_p, train=True)
+    out = attn_weight @ value.float()
+    return out.to(query.dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=True, qk_norm=True, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+
+        self.q_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(head_dim) if qk_norm else nn.Identity()
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x, rope):
+        batch_size, num_tokens, channels = x.shape
+        qkv = self.qkv(x).reshape(batch_size, num_tokens, 3, self.num_heads, channels // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+        q = rope(q)
+        k = rope(k)
+
+        x = scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop.p if self.training else 0.0)
+        x = x.transpose(1, 2).reshape(batch_size, num_tokens, channels)
+        x = x.to(self.proj.weight.dtype)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, drop=0.0, bias=True) -> None:
+        super().__init__()
+        hidden_dim = int(hidden_dim * 2 / 3)
+        self.w12 = nn.Linear(dim, 2 * hidden_dim, bias=bias)
+        self.w3 = nn.Linear(hidden_dim, dim, bias=bias)
+        self.ffn_dropout = nn.Dropout(drop)
+
+    def forward(self, x):
+        x12 = self.w12(x)
+        x1, x2 = x12.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return self.w3(self.ffn_dropout(hidden))
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, hidden_size, patch_size, out_channels):
+        super().__init__()
+        self.norm_final = RMSNorm(hidden_size)
+        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))
+
+    def forward(self, x, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
+        x = modulate(self.norm_final(x), shift, scale)
+        x = self.linear(x)
+        return x
+
+
+class JiTBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        self.norm1 = RMSNorm(hidden_size, eps=1e-6)
+        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, qk_norm=True, attn_drop=attn_drop, proj_drop=proj_drop)
+        self.norm2 = RMSNorm(hidden_size, eps=1e-6)
+        mlp_hidden_dim = int(hidden_size * mlp_ratio)
+        self.mlp = SwiGLUFFN(hidden_size, mlp_hidden_dim, drop=proj_drop)
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))
+
+    def forward(self, x, c, feat_rope=None):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
+        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), rope=feat_rope)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class JiT(nn.Module):
+    def __init__(
+        self,
+        input_size=256,
+        patch_size=16,
+        in_channels=3,
+        hidden_size=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4.0,
+        attn_drop=0.0,
+        proj_drop=0.0,
+        num_classes=1000,
+        bottleneck_dim=128,
+        in_context_len=32,
+        in_context_start=8,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.hidden_size = hidden_size
+        self.input_size = input_size
+        self.in_context_len = in_context_len
+        self.in_context_start = in_context_start
+        self.num_classes = num_classes
+
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder = LabelEmbedder(num_classes, hidden_size)
+        self.x_embedder = BottleneckPatchEmbed(input_size, patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
+
+        num_patches = self.x_embedder.num_patches
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
+
+        if self.in_context_len > 0:
+            self.in_context_posemb = nn.Parameter(torch.zeros(1, self.in_context_len, hidden_size), requires_grad=True)
+            torch.nn.init.normal_(self.in_context_posemb, std=0.02)
+
+        half_head_dim = hidden_size // num_heads // 2
+        hw_seq_len = input_size // patch_size
+        self.feat_rope = VisionRotaryEmbeddingFast(dim=half_head_dim, pt_seq_len=hw_seq_len, num_cls_token=0)
+        self.feat_rope_incontext = VisionRotaryEmbeddingFast(dim=half_head_dim, pt_seq_len=hw_seq_len, num_cls_token=self.in_context_len)
+
+        self.blocks = nn.ModuleList(
+            [
+                JiTBlock(
+                    hidden_size,
+                    num_heads,
+                    mlp_ratio=mlp_ratio,
+                    attn_drop=attn_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
+                    proj_drop=proj_drop if (depth // 4 * 3 > i >= depth // 4) else 0.0,
+                )
+                for i in range(depth)
+            ]
+        )
+
+        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        pos_embed = get_2d_sincos_pos_embed(self.pos_embed.shape[-1], int(self.x_embedder.num_patches**0.5))
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        w1 = self.x_embedder.proj1.weight.data
+        nn.init.xavier_uniform_(w1.view([w1.shape[0], -1]))
+        w2 = self.x_embedder.proj2.weight.data
+        nn.init.xavier_uniform_(w2.view([w2.shape[0], -1]))
+        nn.init.constant_(self.x_embedder.proj2.bias, 0)
+
+        nn.init.normal_(self.y_embedder.embedding_table.weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+        for block in self.blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.linear.weight, 0)
+        nn.init.constant_(self.final_layer.linear.bias, 0)
+
+    def unpatchify(self, x, patch_size):
+        channels = self.out_channels
+        height = width = int(x.shape[1] ** 0.5)
+        assert height * width == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], height, width, patch_size, patch_size, channels))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        images = x.reshape(shape=(x.shape[0], channels, height * patch_size, height * patch_size))
+        return images
+
+    def forward(self, x, t, y):
+        t_emb = self.t_embedder(t)
+        y_emb = self.y_embedder(y)
+        c = t_emb + y_emb
+
+        x = self.x_embedder(x)
+        x += self.pos_embed
+
+        for i, block in enumerate(self.blocks):
+            if self.in_context_len > 0 and i == self.in_context_start:
+                in_context_tokens = y_emb.unsqueeze(1).repeat(1, self.in_context_len, 1)
+                in_context_tokens += self.in_context_posemb
+                x = torch.cat([in_context_tokens, x], dim=1)
+            x = block(x, c, self.feat_rope if i < self.in_context_start else self.feat_rope_incontext)
+
+        x = x[:, self.in_context_len :]
+        x = self.final_layer(x, c)
+        output = self.unpatchify(x, self.patch_size)
+        return output
+
+
+def JiT_B_16(**kwargs):
+    return JiT(depth=12, hidden_size=768, num_heads=12, bottleneck_dim=128, in_context_len=32, in_context_start=4, patch_size=16, **kwargs)
+
+
+def JiT_B_32(**kwargs):
+    return JiT(depth=12, hidden_size=768, num_heads=12, bottleneck_dim=128, in_context_len=32, in_context_start=4, patch_size=32, **kwargs)
+
+
+def JiT_L_16(**kwargs):
+    return JiT(depth=24, hidden_size=1024, num_heads=16, bottleneck_dim=128, in_context_len=32, in_context_start=8, patch_size=16, **kwargs)
+
+
+def JiT_L_32(**kwargs):
+    return JiT(depth=24, hidden_size=1024, num_heads=16, bottleneck_dim=128, in_context_len=32, in_context_start=8, patch_size=32, **kwargs)
+
+
+def JiT_H_16(**kwargs):
+    return JiT(depth=32, hidden_size=1280, num_heads=16, bottleneck_dim=256, in_context_len=32, in_context_start=10, patch_size=16, **kwargs)
+
+
+def JiT_H_32(**kwargs):
+    return JiT(depth=32, hidden_size=1280, num_heads=16, bottleneck_dim=256, in_context_len=32, in_context_start=10, patch_size=32, **kwargs)
+
+
+JiT_models = {
+    "JiT-B/16": JiT_B_16,
+    "JiT-B/32": JiT_B_32,
+    "JiT-L/16": JiT_L_16,
+    "JiT-L/32": JiT_L_32,
+    "JiT-H/16": JiT_H_16,
+    "JiT-H/32": JiT_H_32,
+}

jit_diffusers/modeling_jit_transformer_2d.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import argparse
+from collections.abc import Mapping
+from dataclasses import dataclass
+from typing import Any, Dict, Literal, Tuple
+
+import torch
+from diffusers import ConfigMixin, ModelMixin
+from diffusers.configuration_utils import register_to_config
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+from .modeling_jit_backbone import JiT_models
+
+
+def _extract_module_state_dict(
+    state_dict: Dict[str, torch.Tensor], prefixes: Tuple[str, ...] = ("transformer.", "net.")
+) -> Dict[str, torch.Tensor]:
+    """Extract module state by stripping the first fully-matching prefix.
+
+    Prefix precedence is left-to-right; `"transformer."` is preferred over legacy `"net."`.
+    """
+    for prefix in prefixes:
+        if all(key.startswith(prefix) for key in state_dict.keys()):
+            return {k[len(prefix):]: v for k, v in state_dict.items()}
+    return state_dict
+
+
+def _build_jit_kwargs(
+    image_size: int,
+    num_classes: int,
+    attn_dropout: float,
+    proj_dropout: float,
+    model_name: str | None = None,
+) -> Dict[str, object]:
+    # Keep model_name for backward-compatible internal call signatures.
+    _ = model_name
+    return {
+        "input_size": image_size,
+        "in_channels": 3,
+        "num_classes": num_classes,
+        "attn_drop": attn_dropout,
+        "proj_drop": proj_dropout,
+    }
+
+
+@dataclass
+class JiTCheckpointConfig:
+    model_name: str
+    image_size: int
+    num_classes: int
+    attn_dropout: float
+    proj_dropout: float
+
+
+def _config_from_checkpoint(ckpt_args: argparse.Namespace | Mapping[str, Any]) -> JiTCheckpointConfig:
+    if isinstance(ckpt_args, argparse.Namespace):
+        args_dict = vars(ckpt_args)
+    elif isinstance(ckpt_args, Mapping):
+        args_dict = ckpt_args
+    else:
+        raise TypeError(f"Unsupported checkpoint args type: {type(ckpt_args)}")
+
+    def _get_first_available(*keys: str, default=None):
+        for key in keys:
+            if key in args_dict and args_dict[key] is not None:
+                return args_dict[key]
+        return default
+
+    model_name = _get_first_available("model", "model_name", "model_type")
+    image_size = _get_first_available("img_size", "image_size", "sample_size")
+    num_classes = _get_first_available("class_num", "num_classes", "num_class_embeds")
+    if model_name is None or image_size is None or num_classes is None:
+        raise ValueError("Checkpoint args are missing model/image_size/num_classes information.")
+
+    return JiTCheckpointConfig(
+        model_name=str(model_name),
+        image_size=int(image_size),
+        num_classes=int(num_classes),
+        attn_dropout=float(_get_first_available("attn_dropout", "attention_dropout", default=0.0)),
+        proj_dropout=float(_get_first_available("proj_dropout", "dropout", default=0.0)),
+    )
+
+
+class JiTTransformer2DModel(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self,
+        model_type: str = "JiT-B/16",
+        sample_size: int = 256,
+        num_class_embeds: int = 1000,
+        attention_dropout: float = 0.0,
+        dropout: float = 0.0,
+        model_name: str | None = None,
+        image_size: int | None = None,
+        num_classes: int | None = None,
+        attn_dropout: float | None = None,
+        proj_dropout: float | None = None,
+    ):
+        super().__init__()
+        resolved_model_type = model_type if model_name is None else model_name
+        resolved_sample_size = sample_size if image_size is None else image_size
+        resolved_num_class_embeds = num_class_embeds if num_classes is None else num_classes
+        resolved_attention_dropout = attention_dropout if attn_dropout is None else attn_dropout
+        resolved_dropout = dropout if proj_dropout is None else proj_dropout
+
+        if resolved_model_type not in JiT_models:
+            raise ValueError(f"Unknown model '{resolved_model_type}'. Available: {list(JiT_models.keys())}")
+
+        self.transformer = JiT_models[resolved_model_type](
+            **_build_jit_kwargs(
+                image_size=resolved_sample_size,
+                num_classes=resolved_num_class_embeds,
+                attn_dropout=resolved_attention_dropout,
+                proj_dropout=resolved_dropout,
+                model_name=resolved_model_type,
+            )
+        )
+
+    def forward(
+        self,
+        sample: torch.Tensor,
+        timestep: torch.Tensor,
+        class_labels: torch.Tensor,
+        return_dict: bool = True,
+    ):
+        timestep = torch.as_tensor(timestep, device=sample.device)
+        if timestep.ndim == 0:
+            timestep = timestep.repeat(sample.shape[0])
+        else:
+            timestep = timestep.reshape(-1)
+            if timestep.shape[0] == 1 and sample.shape[0] > 1:
+                timestep = timestep.repeat(sample.shape[0])
+
+        denoised = self.transformer(sample, timestep, class_labels)
+        if not return_dict:
+            return (denoised,)
+        return Transformer2DModelOutput(sample=denoised)
+
+    @classmethod
+    def from_jit_checkpoint(
+        cls,
+        checkpoint_path: str,
+        weights: Literal["model", "ema1", "ema2"] = "ema1",
+        map_location: str = "cpu",
+        strict: bool = True,
+    ) -> Tuple["JiTTransformer2DModel", Dict[str, object]]:
+        checkpoint = torch.load(checkpoint_path, map_location=map_location)
+        if "args" not in checkpoint:
+            raise ValueError("Checkpoint is missing 'args', cannot infer JiT architecture config.")
+
+        config = _config_from_checkpoint(checkpoint["args"])
+        model = cls(
+            model_type=config.model_name,
+            sample_size=config.image_size,
+            num_class_embeds=config.num_classes,
+            attention_dropout=config.attn_dropout,
+            dropout=config.proj_dropout,
+        )
+
+        key = "model" if weights == "model" else f"model_{weights}"
+        if key not in checkpoint:
+            raise ValueError(f"Checkpoint key '{key}' not found. Available keys: {list(checkpoint.keys())}")
+
+        model_state = _extract_module_state_dict(checkpoint[key])
+        model.transformer.load_state_dict(model_state, strict=strict)
+
+        metadata = {
+            "checkpoint_path": checkpoint_path,
+            "weights": weights,
+            "epoch": checkpoint.get("epoch"),
+            "source_args": checkpoint.get("args"),
+        }
+        return model, metadata
+
+    def to_jit_checkpoint(
+        self,
+        ema_mode: Literal["none", "copy_to_both"] = "copy_to_both",
+        prefix: str = "net.",
+    ) -> Dict[str, object]:
+        base_state = {f"{prefix}{k}": v.detach().cpu() for k, v in self.transformer.state_dict().items()}
+        checkpoint = {"model": base_state}
+        if ema_mode == "copy_to_both":
+            checkpoint["model_ema1"] = {k: v.clone() for k, v in base_state.items()}
+            checkpoint["model_ema2"] = {k: v.clone() for k, v in base_state.items()}
+        elif ema_mode != "none":
+            raise ValueError(f"Unsupported ema_mode='{ema_mode}'.")
+        return checkpoint
+
+    @property
+    def net(self):
+        return self.transformer
+
+    @net.setter
+    def net(self, module):
+        self.transformer = module
+
+
+# Backward-compatible alias.
+JiTDiffusersModel = JiTTransformer2DModel

jit_diffusers/modeling_jit_utils.py

@@ -0,0 +1,129 @@
+from math import pi
+
+import numpy as np
+import torch
+from einops import rearrange, repeat
+from torch import nn
+
+
+def broadcat(tensors, dim=-1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda tensor: len(tensor.shape), tensors)))
+    assert len(shape_lens) == 1, "tensors must all have the same number of dimensions"
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda tensor: list(tensor.shape), tensors)))
+    expandable_dims = [(index, val) for index, val in enumerate(dims) if index != dim]
+    assert all([*map(lambda tensor: len(set(tensor[1])) <= 2, expandable_dims)]), "invalid dimensions for broadcastable concatenation"
+    max_dims = list(map(lambda tensor: (tensor[0], max(tensor[1])), expandable_dims))
+    expanded_dims = list(map(lambda tensor: (tensor[0], (tensor[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda tensor: tensor[1], expanded_dims)))
+    tensors = list(map(lambda tensor: tensor[0].expand(*tensor[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim=dim)
+
+
+def rotate_half(x):
+    x = rearrange(x, "... (d r) -> ... d r", r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, "... d r -> ... (d r)")
+
+
+class VisionRotaryEmbeddingFast(nn.Module):
+    def __init__(
+        self,
+        dim,
+        pt_seq_len=16,
+        ft_seq_len=None,
+        custom_freqs=None,
+        freqs_for="lang",
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+        num_cls_token=0,
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == "lang":
+            freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+        elif freqs_for == "pixel":
+            freqs = torch.linspace(1.0, max_freq / 2, dim // 2) * pi
+        elif freqs_for == "constant":
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f"unknown modality {freqs_for}")
+
+        if ft_seq_len is None:
+            ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs = torch.einsum("..., f -> ... f", t, freqs)
+        freqs = repeat(freqs, "... n -> ... (n r)", r=2)
+        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim=-1)
+
+        if num_cls_token > 0:
+            freqs_flat = freqs.view(-1, freqs.shape[-1])
+            cos_img = freqs_flat.cos()
+            sin_img = freqs_flat.sin()
+            _, dim_freq = cos_img.shape
+            cos_pad = torch.ones(num_cls_token, dim_freq, dtype=cos_img.dtype, device=cos_img.device)
+            sin_pad = torch.zeros(num_cls_token, dim_freq, dtype=sin_img.dtype, device=sin_img.device)
+            self.register_buffer("freqs_cos", torch.cat([cos_pad, cos_img], dim=0), persistent=False)
+            self.register_buffer("freqs_sin", torch.cat([sin_pad, sin_img], dim=0), persistent=False)
+        else:
+            self.register_buffer("freqs_cos", freqs.cos().view(-1, freqs.shape[-1]), persistent=False)
+            self.register_buffer("freqs_sin", freqs.sin().view(-1, freqs.shape[-1]), persistent=False)
+
+    def forward(self, tensor):
+        freqs_cos = self.freqs_cos.to(device=tensor.device, dtype=tensor.dtype)
+        freqs_sin = self.freqs_sin.to(device=tensor.device, dtype=tensor.dtype)
+        return tensor * freqs_cos + rotate_half(tensor) * freqs_sin
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0):
+    grid_h = np.arange(grid_size, dtype=np.float32)
+    grid_w = np.arange(grid_size, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0)
+    grid = grid.reshape([2, 1, grid_size, grid_size])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token and extra_tokens > 0:
+        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+    emb = np.concatenate([emb_h, emb_w], axis=1)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)
+    return emb

jit_diffusers/pipeline_jit.py

@@ -0,0 +1,156 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import List, Tuple
+
+import numpy as np
+import torch
+from diffusers import DiffusionPipeline
+from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
+from diffusers.utils import BaseOutput
+from diffusers.utils.torch_utils import randn_tensor
+
+from .modeling_jit_transformer_2d import JiTTransformer2DModel
+from .scheduling_jit import JiTScheduler
+
+
+@dataclass
+class JiTPipelineOutput(BaseOutput):
+    images: List["PIL.Image.Image"] | np.ndarray | torch.Tensor
+
+
+class JiTPipeline(DiffusionPipeline):
+    model_cpu_offload_seq = "transformer"
+
+    def __init__(self, transformer: JiTTransformer2DModel, scheduler: JiTScheduler | None = None):
+        super().__init__()
+        self.register_modules(transformer=transformer, scheduler=scheduler or JiTScheduler())
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        model_kwargs = dict(kwargs)
+        transformer_subfolder = model_kwargs.pop("transformer_subfolder", None)
+        scheduler_subfolder = model_kwargs.pop("scheduler_subfolder", None)
+        scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
+        if transformer_subfolder is not None:
+            transformer_path = str(Path(pretrained_model_name_or_path) / transformer_subfolder)
+        else:
+            transformer_path = pretrained_model_name_or_path
+        transformer = JiTTransformer2DModel.from_pretrained(transformer_path, **model_kwargs)
+        try:
+            scheduler = JiTScheduler.from_pretrained(
+                pretrained_model_name_or_path,
+                subfolder=scheduler_subfolder,
+                **scheduler_kwargs,
+            )
+        except Exception:
+            scheduler = JiTScheduler(**scheduler_kwargs)
+        return cls(transformer=transformer, scheduler=scheduler)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        class_labels: int | List[int] | torch.Tensor,
+        num_inference_steps: int = 50,
+        guidance_scale: float = 2.9,
+        guidance_interval_min: float = 0.1,
+        guidance_interval_max: float = 1.0,
+        noise_scale: float = 2.0,
+        t_eps: float = 5e-2,
+        sampling_method: str | None = None,
+        generator: torch.Generator | List[torch.Generator] | None = None,
+        output_type: str = "pil",
+        return_dict: bool = True,
+    ) -> JiTPipelineOutput | ImagePipelineOutput | Tuple:
+        if output_type not in {"pil", "np", "pt"}:
+            raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")
+        if sampling_method is not None and sampling_method not in {"heun", "euler"}:
+            raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
+        if num_inference_steps < 2:
+            raise ValueError("num_inference_steps must be >= 2.")
+        if sampling_method is not None and sampling_method != self.scheduler.config.solver:
+            self.scheduler = JiTScheduler.from_config(self.scheduler.config, solver=sampling_method)
+
+        if isinstance(class_labels, int):
+            class_labels = [class_labels]
+        if isinstance(class_labels, list):
+            class_labels = torch.tensor(class_labels, device=self._execution_device, dtype=torch.long)
+        else:
+            class_labels = class_labels.to(self._execution_device, dtype=torch.long).reshape(-1)
+
+        batch_size = class_labels.shape[0]
+        latent_size = int(self.transformer.config.sample_size)
+        latent_channels = int(getattr(self.transformer.config, "in_channels", 3))
+        num_classes = int(self.transformer.config.num_class_embeds)
+
+        class_labels = class_labels.clamp(0, num_classes - 1)
+        class_null = torch.full_like(class_labels, num_classes)
+
+        latents = randn_tensor(
+            shape=(batch_size, latent_channels, latent_size, latent_size),
+            generator=generator,
+            device=self._execution_device,
+            dtype=self.transformer.dtype,
+        ) * noise_scale
+        self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=self._execution_device)
+        timesteps = self.scheduler.timesteps.to(device=self._execution_device, dtype=latents.dtype)
+
+        def forward_cfg(z_value: torch.Tensor, t: torch.Tensor | float) -> torch.Tensor:
+            t = torch.as_tensor(t, device=self._execution_device, dtype=latents.dtype)
+            x_cond = self.transformer(sample=z_value, timestep=t.flatten(), class_labels=class_labels).sample
+            v_cond = (x_cond - z_value) / (1.0 - t).clamp_min(t_eps)
+
+            x_uncond = self.transformer(sample=z_value, timestep=t.flatten(), class_labels=class_null).sample
+            v_uncond = (x_uncond - z_value) / (1.0 - t).clamp_min(t_eps)
+
+            interval_mask = (t < guidance_interval_max) & (t > guidance_interval_min)
+            scale = torch.where(
+                interval_mask,
+                torch.tensor(guidance_scale, device=self._execution_device, dtype=latents.dtype),
+                torch.tensor(1.0, device=self._execution_device, dtype=latents.dtype),
+            )
+            return v_uncond + scale * (v_cond - v_uncond)
+
+        for i in self.progress_bar(range(num_inference_steps - 1)):
+            t, t_next = timesteps[i], timesteps[i + 1]
+            model_output = forward_cfg(latents, t)
+            if self.scheduler.config.solver == "heun":
+                latents = self.scheduler.step(
+                    model_output=model_output,
+                    timestep=t,
+                    next_timestep=t_next,
+                    sample=latents,
+                    model_fn=forward_cfg,
+                ).prev_sample
+            else:
+                latents = self.scheduler.step(
+                    model_output=model_output,
+                    timestep=t,
+                    next_timestep=t_next,
+                    sample=latents,
+                ).prev_sample
+
+        # Match the original JiT implementation: always use Euler for the final step.
+        t, t_next = timesteps[-2], timesteps[-1]
+        model_output = forward_cfg(latents, t)
+        latents = self.scheduler.euler_step(
+            model_output=model_output,
+            timestep=t,
+            next_timestep=t_next,
+            sample=latents,
+        ).prev_sample
+
+        images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
+        if output_type == "pt":
+            images = images_pt
+        else:
+            images_np = images_pt.permute(0, 2, 3, 1).numpy()
+            if output_type == "np":
+                images = images_np
+            else:
+                images = self.numpy_to_pil(images_np)
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (images,)
+        return JiTPipelineOutput(images=images)

jit_diffusers/scheduling_jit.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import SchedulerMixin, SchedulerOutput
+
+
+class JiTScheduler(SchedulerMixin, ConfigMixin):
+    order = 1
+
+    @register_to_config
+    def __init__(
+        self,
+        solver: str = "heun",
+        timestep_start: float = 0.0,
+        timestep_end: float = 1.0,
+    ):
+        if solver not in {"heun", "euler"}:
+            raise ValueError("solver must be one of: 'heun', 'euler'.")
+        if timestep_end <= timestep_start:
+            raise ValueError("timestep_end must be greater than timestep_start.")
+        self.timesteps = torch.tensor([])
+
+    def set_timesteps(self, num_inference_steps: int, device: str | torch.device | None = None):
+        if num_inference_steps < 2:
+            raise ValueError("num_inference_steps must be >= 2.")
+        self.timesteps = torch.linspace(
+            self.config.timestep_start,
+            self.config.timestep_end,
+            num_inference_steps + 1,
+            device=device,
+            dtype=torch.float32,
+        )
+
+    def euler_step(
+        self,
+        model_output: torch.Tensor,
+        timestep: torch.Tensor,
+        next_timestep: torch.Tensor,
+        sample: torch.Tensor,
+        return_dict: bool = True,
+    ) -> SchedulerOutput | tuple[torch.Tensor]:
+        prev_sample = sample + (next_timestep - timestep) * model_output
+        if not return_dict:
+            return (prev_sample,)
+        return SchedulerOutput(prev_sample=prev_sample)
+
+    def step(
+        self,
+        model_output: torch.Tensor,
+        timestep: torch.Tensor,
+        next_timestep: torch.Tensor,
+        sample: torch.Tensor,
+        model_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
+        return_dict: bool = True,
+    ) -> SchedulerOutput | tuple[torch.Tensor]:
+        if self.config.solver == "euler":
+            return self.euler_step(model_output, timestep, next_timestep, sample, return_dict=return_dict)
+
+        if model_fn is None:
+            raise ValueError("model_fn is required when solver='heun'.")
+
+        sample_euler = sample + (next_timestep - timestep) * model_output
+        model_output_next = model_fn(sample_euler, next_timestep)
+        prev_sample = sample + (next_timestep - timestep) * 0.5 * (model_output + model_output_next)
+
+        if not return_dict:
+            return (prev_sample,)
+        return SchedulerOutput(prev_sample=prev_sample)

run_jit_diffusers_inference.py

@@ -0,0 +1,96 @@
+import argparse
+from pathlib import Path
+import sys
+
+import torch
+
+SCRIPT_DIR = Path(__file__).resolve().parent
+if str(SCRIPT_DIR) not in sys.path:
+    sys.path.insert(0, str(SCRIPT_DIR))
+
+from jit_diffusers import JiTPipeline
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run single-image JiT diffusers inference.")
+    parser.add_argument("--model_path", type=str, required=True, help="Path to converted diffusers model directory.")
+    parser.add_argument("--output_path", type=str, required=True, help="Path to save output PNG image.")
+    parser.add_argument("--class_label", type=int, default=207, help="ImageNet class id for conditional generation.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
+    parser.add_argument("--steps", type=int, default=50, help="Number of ODE sampling steps.")
+    parser.add_argument("--cfg", type=float, default=2.9, help="Classifier-free guidance scale.")
+    parser.add_argument("--interval_min", type=float, default=0.1, help="CFG interval min.")
+    parser.add_argument("--interval_max", type=float, default=1.0, help="CFG interval max.")
+    parser.add_argument("--noise_scale", type=float, default=2.0, help="Initial Gaussian noise scale.")
+    parser.add_argument("--t_eps", type=float, default=5e-2, help="Small epsilon for timestep denominator.")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="auto",
+        choices=["auto", "cuda", "cpu"],
+        help="Inference device.",
+    )
+    parser.add_argument(
+        "--dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp32"],
+        help="Inference dtype. Defaults to bf16 on CUDA.",
+    )
+    parser.add_argument(
+        "--solver",
+        type=str,
+        default="scheduler",
+        choices=["scheduler", "heun", "euler"],
+        help="Sampling solver. Use scheduler to keep pipeline default.",
+    )
+    return parser.parse_args()
+
+
+def resolve_device(name: str) -> torch.device:
+    if name == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(name)
+
+
+def resolve_dtype(name: str, device: torch.device) -> torch.dtype:
+    if name == "bf16":
+        return torch.bfloat16 if device.type == "cuda" else torch.float32
+    return torch.float32
+
+
+def main() -> None:
+    args = parse_args()
+    device = resolve_device(args.device)
+    dtype = resolve_dtype(args.dtype, device)
+    if device.type == "cuda":
+        torch.set_float32_matmul_precision("high")
+
+    pipe = JiTPipeline.from_pretrained(args.model_path).to(device)
+    pipe.transformer = pipe.transformer.to(device=device, dtype=dtype)
+    pipe.transformer.eval()
+    sampling_method = None if args.solver == "scheduler" else args.solver
+
+    generator = torch.Generator(device=device).manual_seed(args.seed)
+    output = pipe(
+        class_labels=[args.class_label],
+        num_inference_steps=args.steps,
+        guidance_scale=args.cfg,
+        guidance_interval_min=args.interval_min,
+        guidance_interval_max=args.interval_max,
+        noise_scale=args.noise_scale,
+        t_eps=args.t_eps,
+        sampling_method=sampling_method,
+        generator=generator,
+        output_type="pil",
+    )
+    image = output.images[0]
+
+    output_path = Path(args.output_path)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    image.save(output_path)
+    print(f"Saved image to: {output_path}")
+
+
+if __name__ == "__main__":
+    main()