JiT Diffusers
Native diffusers implementation of JiT (Just image Transformer). Each variant folder is self-contained:
Model source
- Kind
- Other
- File size
- 501 MB–3.56 GB
- Version
- vDiffusion pytorch model
- Updated
- Oct 9, 2026
- Source
- Hugging Face
Source description
JiT-diffusers
Native diffusers implementation of JiT (Just image Transformer). Each variant folder is self-contained:
pipeline.py—JiTPipelinescheduler/scheduler_config.json—FlowMatchHeunDiscreteSchedulerconfig (defaultshift=4.0)transformer/jit_transformer_2d.py—JiTTransformer2DModel
The pipeline now supports dynamic inference resolution in __call__ with positional interpolation.
No separate jit_diffusers package; only PyPI diffusers plus local custom code in the variant directory.
Available checkpoints
| Checkpoint | Path | Resolution | Recommended CFG |
|---|---|---|---|
| JiT-B/16 | ./JiT-B-16 | 256×256 | 3.0 |
| JiT-L/16 | ./JiT-L-16 | 256×256 | 2.4 |
| JiT-H/16 | ./JiT-H-16 | 256×256 | 2.2 |
| JiT-B/32 | ./JiT-B-32 | 512×512 | 3.0 |
| JiT-L/32 | ./JiT-L-32 | 512×512 | 2.5 |
| JiT-H/32 | ./JiT-H-32 | 512×512 | 2.3 |
ImageNet class labels
Each variant keeps an English id2label map directly in its own model_index.json (DiT-style).
pipe.id2label— inspect id → English label correspondencepipe.labels— reverse map (English synonym → id), sorted for browsingpipe.get_label_ids("golden retriever")pipe(class_labels="golden retriever", ...)— string labels resolved automatically
Chinese labels are preserved in the main source repo under src/labels/id2label_cn.json for reference.
Inference
Run the bundled demo script from the repo root:
python demo_inference.py
This writes demo.png using JiT-H-32 with the settings below.
from pathlib import Path
from diffusers import DiffusionPipeline, FlowMatchHeunDiscreteScheduler
import torch
model_dir = Path("./JiT-H-32")
pipe = DiffusionPipeline.from_pretrained(
str(model_dir),
custom_pipeline=str(model_dir / "pipeline.py"),
trust_remote_code=True,
)
pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(pipe.scheduler.config, shift=4.0)
pipe.to("cuda")
# Numeric or human-readable labels
print(pipe.id2label[207])
print(pipe.get_label_ids("golden retriever"))
generator = torch.Generator(device="cuda").manual_seed(42)
image = pipe(
class_labels="golden retriever",
num_inference_steps=50,
guidance_scale=2.3,
generator=generator,
).images[0]
image.save("demo.png")
height and width default to the checkpoint's native resolution when omitted.
Load a variant subfolder (e.g. ./JiT-H-32), not the repo root.
Sources
1 source
Sources
1 sourceVerified Oct 9
Download files
Model artifacts
6 artifacts
Model artifacts
6 artifactsJiT-B-16/transformer/diffusion_pytorch_model.safetensors
safetensors · 501 MB · SHA-256 b68278f2e16a…4772 · Hugging Face
Download fileJiT-B-32/transformer/diffusion_pytorch_model.safetensors
safetensors · 509 MB · SHA-256 729654b3302f…6ba4 · Hugging Face
Download fileJiT-H-16/transformer/diffusion_pytorch_model.safetensors
safetensors · 3.55 GB · SHA-256 b6ad4cf51f5f…68e7 · Hugging Face
Download fileJiT-H-32/transformer/diffusion_pytorch_model.safetensors
safetensors · 3.56 GB · SHA-256 578fc2f9f4cc…0ab6 · Hugging Face
Download fileJiT-L-16/transformer/diffusion_pytorch_model.safetensors
safetensors · 1.71 GB · SHA-256 e9285393d92d…6b0f · Hugging Face
Download fileJiT-L-32/transformer/diffusion_pytorch_model.safetensors
safetensors · 1.72 GB · SHA-256 121d3917ab50…3266 · Hugging Face
Download fileSource excerpts
2 excerpts
Source excerpts
2 excerpts--- license: mit library_name: diffusers pipeline_tag: unconditional-image-generation tags: - diffusers - jit - image-generation - class-conditional widget: - output: url: demo.png language: - en --- # JiT-diffusers Native diffusers implementation of **JiT** (Just image Transformer). Each variant folder is self-contained: - `pipeline.py` — `JiTPipeline` - `scheduler/scheduler_config.json` — `FlowMatchHeunDiscreteScheduler` config (default `shift=4.0`) - `transformer/jit_transformer_2d.py` — `JiTTransformer2DModel` The pipeline now supports dynamic inference resolution in `__call__` with positional interpolation. No separate `jit_diffusers` package; only PyPI `diffusers` plus local custom code in the variant directory. ## Available checkpoints | Checkpoint | Path | Resolution | Recommended CFG | | --- | --- | --- | --- | | JiT-B/16 | `./JiT-B-16` | 256×256 | 3.0 | | JiT-L/16 | `./JiT-L-16` | 256×256 | 2.4 | | JiT-H/16 | `./JiT-H-16` | 256×256 | 2.2 | | JiT-B/32 | `./JiT-B-32` | 512×512 | 3.0 | | JiT-L/32 | `./JiT-L-32` | 512×512 | 2.5 | | JiT-H/32 | `./JiT-H-32` | 512×512 | 2.3 | ## ImageNet class labels Each variant keeps an English `id2label` map directly in its own `model_index.json` (DiT-style). - `pipe.id2label` — inspect id → English label correspondence - `pipe.labels` — reverse map (English synonym → id), sorted for browsing - `pipe.get_label_ids("golden retriever")` - `pipe(class_labels="golden retriever", ...)` — string labels resolved automatically Chinese labels are preserved in the main source repo under `src/labels/id2label_cn.json` for reference. ## Inference Run the bundled demo script from the repo root: ```bash python demo_inference.py ``` This writes `demo.png` using `JiT-H-32` with the settings below. ```python from pathlib import Path from diffusers import DiffusionPipeline, FlowMatchHeunDiscreteScheduler import torch model_dir = Path("./JiT-H-32") pipe = DiffusionPipeline.from_pretrained( str(model_dir), custom_pipeline=str(model_dir / "pipeline.py"), trust_remote_code=True, ) pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(pipe.scheduler.config, shift=4.0) pipe.to("cuda") # Numeric or human-readable labels print(pipe.id2label[207]) print(pipe.get_label_ids("golden retriever")) generator = torch.Generator(device="cuda").manual_seed(42) image = pipe( class_labels="golden retriever", num_inference_st...
Source context: 0 downloads · 2 likes · Pipeline unconditional-image-generation · Library diffusers · Repo BiliSakura/JiT-diffusers
