BiliSakura/JiT-diffusers

1---2license: mit3library_name: diffusers4pipeline_tag: unconditional-image-generation5tags:6- diffusers7- jit8- image-generation9- class-conditional10widget:11- output:12    url: demo.png13language:14- en15---16 17# JiT-diffusers18 19Native diffusers implementation of **JiT** (Just image Transformer). Each variant folder is self-contained:20 21- `pipeline.py` — `JiTPipeline`22- `scheduler/scheduler_config.json` — `FlowMatchHeunDiscreteScheduler` config (default `shift=4.0`)23- `transformer/jit_transformer_2d.py` — `JiTTransformer2DModel`24 25The pipeline now supports dynamic inference resolution in `__call__` with positional interpolation.26 27No separate `jit_diffusers` package; only PyPI `diffusers` plus local custom code in the variant directory.28 29## Available checkpoints30 31| Checkpoint | Path | Resolution | Recommended CFG |32| --- | --- | --- | --- |33| JiT-B/16 | `./JiT-B-16` | 256×256 | 3.0 |34| JiT-L/16 | `./JiT-L-16` | 256×256 | 2.4 |35| JiT-H/16 | `./JiT-H-16` | 256×256 | 2.2 |36| JiT-B/32 | `./JiT-B-32` | 512×512 | 3.0 |37| JiT-L/32 | `./JiT-L-32` | 512×512 | 2.5 |38| JiT-H/32 | `./JiT-H-32` | 512×512 | 2.3 |39 40## ImageNet class labels41 42Each variant keeps an English `id2label` map directly in its own `model_index.json` (DiT-style).43 44- `pipe.id2label` — inspect id → English label correspondence45- `pipe.labels` — reverse map (English synonym → id), sorted for browsing46- `pipe.get_label_ids("golden retriever")`47- `pipe(class_labels="golden retriever", ...)` — string labels resolved automatically48 49Chinese labels are preserved in the main source repo under `src/labels/id2label_cn.json` for reference.50 51## Inference52 53Run the bundled demo script from the repo root:54 55```bash56python demo_inference.py57```58 59This writes `demo.png` using `JiT-H-32` with the settings below.60 61```python62from pathlib import Path63from diffusers import DiffusionPipeline, FlowMatchHeunDiscreteScheduler64import torch65 66model_dir = Path("./JiT-H-32")67pipe = DiffusionPipeline.from_pretrained(68    str(model_dir),69    custom_pipeline=str(model_dir / "pipeline.py"),70    trust_remote_code=True,71)72pipe.scheduler = FlowMatchHeunDiscreteScheduler.from_config(pipe.scheduler.config, shift=4.0)73pipe.to("cuda")74 75# Numeric or human-readable labels76print(pipe.id2label[207])77print(pipe.get_label_ids("golden retriever"))78 79generator = torch.Generator(device="cuda").manual_seed(42)80image = pipe(81    class_labels="golden retriever",82    num_inference_steps=50,83    guidance_scale=2.3,84    generator=generator,85).images[0]86image.save("demo.png")87```88 89`height` and `width` default to the checkpoint's native resolution when omitted.90 91Load a **variant subfolder** (e.g. `./JiT-H-32`), not the repo root.92