BiliSakura/JiT-diffusers

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

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:

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.