LightningDiT-diffusers

Diffusers-ready checkpoints for LightningDiT (VA-VAE–aligned latent diffusion with flow matching), converted from hustvl/lightningdit-xl-imagenet256-800ep for local/offline use.

This root folder is a model collection that contains:

• LightningDit-XL-1-256

Each subfolder is a self-contained Diffusers model repo with:

• pipeline.py (LightningDiTPipeline)
• transformer/transformer_lightningdit.py and weights
• scheduler/scheduler_config.json (FlowMatchHeunDiscreteScheduler, shift=0.3)
• vae/ (REPA-E/vavae-hf)

Each variant embeds English id2label in model_index.json, so class labels can be passed as ImageNet ids or English synonym strings.

Demo

Class-conditional sample (ImageNet class 207, golden retriever), LightningDiT-XL/1 at 256×256, 250 steps, CFG 6.7, cfg_interval_start=0.125, timestep_shift=0.3, seed 0.

Model Paths

Model	Resolution	Local path
LightningDiT-XL/1	256×256	./LightningDit-XL-1-256

Inference

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "./LightningDit-XL-1-256",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")

class_id = pipe.get_label_ids("golden retriever")[0]
image = pipe(
    class_labels=class_id,
    num_inference_steps=250,
    guidance_scale=6.7,
    cfg_interval_start=0.125,
    timestep_shift=0.3,
    generator=torch.Generator(device="cuda").manual_seed(0),
).images[0]

Citation

@inproceedings{yao2025reconstruction,
  title={Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models},
  author={Yao, Jingfeng and Yang, Bin and Wang, Xinggang},
  booktitle={CVPR},
  year={2025}
}