Add JiT: Just image Transformer (JiT) for Pixel-space Diffusion

[sayakpaul]

Pixel-level diffusion seems to be catching up and I think diffusers could have a nice role to play in accelerating its adoption and research.

Code and checkpoints can be found here: https://github.com/LTH14/JiT.git.

Even though the released checkpoints were obtained with ImageNet, I believe we can support this like we supported the original DiT when it came out.

Cc: @LTH14 (first author), @kashif

[AlanPonnachan]

@sayakpaul ,
I would like to work on this. After reviewing the official official implementation, I’ve outlined a plan for the PR below:

1. Model (JiTTransformer2DModel):

   • Add src/diffusers/models/transformers/jit_transformer_2d.py.
   • Port the JiT architecture:
     • JiTBlock with RMSNorm, JiTAttention (including RoPE), and JiTSwiGLUFFN.
     • Custom embeddings: JiTPatchEmbed (bottleneck design), JiTTimestepEmbedder, JiTLabelEmbedder, and fixed sin-cos positional embeddings.
     • Support for "In-context" class token handling as per the original implementation.

2. Pipeline (JiTPipeline):

   • Add src/diffusers/pipelines/jit/pipeline_jit.py.
   • Implement a pixel-space diffusion pipeline compatible with FlowMatchEulerDiscreteScheduler.
   • The pipeline will handle the specific flow-matching parameterization used by JiT (mapping model output $x$ to velocity $v$ for the scheduler).

3. Scheduler:

   • Use the existing FlowMatchEulerDiscreteScheduler.

4. Integration:

   • Expose the model and pipeline in the respective __init__.py files.
   • Add conversion scripts for original checkpoints (if needed).

I will follow up with a PR ASAP

[sayakpaul]

Plan looks really good! Thank you!