NEPA-DiT-XL (ImageNet-1K 256x256, 400k steps, EMA)
Class-conditional latent diffusion model on ImageNet-1K at 256x256.
- Condition encoder: NEPA-XL (
LlamaNepaModel, 28 layers, d=1152). It reads[class token, <boi>, noisy latent]and emits one embedding per latent patch (patch size 4 on the 32x32x4 SD latent). - Denoiser: Flux2-style MMDiT (
Flux2Transformer2DModel, 10 double + 18 single blocks), conditioned on the NEPA embeddings, trained with flow matching + REPA. Weights are the EMA at step 400k. - VAE:
stabilityai/sd-vae-ft-emadecoder (f8, 4 channels); latents share the SD VAE encoder space. - Sampler: flow-matching Euler ODE (
FlowMatchEulerDiscreteScheduler), classifier-free guidance with the null class.
Everything needed to run the model (encoder + pipeline code) ships in this repo; load with trust_remote_code=True.
Usage
pip install "diffusers>=0.36" "transformers>=5.3" torch safetensors
Tested with transformers 5.3.0 / diffusers 0.36.0 and transformers 5.16.1 / diffusers 0.40.0 (identical outputs).
import torch
from diffusers import DiffusionPipeline
pipe = DiffusionPipeline.from_pretrained(
"SixAILab/nepa-dit-xl-400k", trust_remote_code=True, torch_dtype=torch.bfloat16
).to("cuda")
class_ids = pipe.get_label_ids(["golden retriever", "macaw"]) # or pass ints in [0, 999] directly
images = pipe(
class_labels=class_ids,
guidance_scale=3.6,
num_inference_steps=96,
generator=torch.Generator("cuda").manual_seed(0),
).images
images[0].save("golden_retriever.png")
__call__ arguments:
| arg | default | meaning |
|---|---|---|
class_labels | required | list of ImageNet-1K ids (one image per id) |
guidance_scale | 3.6 | CFG weight; 1.0 disables guidance |
num_inference_steps | 96 | Euler ODE steps |
generator | None | torch.Generator for reproducibility |
output_type | "pil" | "pil", "np", "pt" or "latent" |
Output resolution is fixed at 256x256. bf16 is recommended (all reported numbers use it); fp16 (e.g. on a T4) and fp32 also work.
Notes
- CFG is applied on every step. For best FID with full-range CFG use
guidance_scalearound 1.4-1.5; larger values (3-4) trade diversity for fidelity, as with DiT. - Cost per step is 1 NEPA pass + 1 MMDiT pass for the conditional branch, and the same again for the
unconditional branch when
guidance_scale > 1.