| --- |
| 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_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. |
|
|
