NiT-B/README.md

metadata

license: apache-2.0
library_name: diffusers
pipeline_tag: unconditional-image-generation
tags:
  - diffusers
  - nit
  - image-generation
  - class-conditional
  - imagenet
inference: true

NiT-B

Self-contained Diffusers checkpoint for NiT-B (131M), converted from GoodEnough/NiT-B-Models (model_500K.safetensors, 500K training steps).

Architecture and training settings follow the official nit_b_pack_merge_radio_65536.yaml.

Model config

Field	Value
Parameters	131M
Depth	12
Hidden size	768
Attention heads	12
Encoder depth	4
Latent channels (z_dim)	1280
Patch size	1
Input latent channels	32
Classes	1000
Class dropout	0.1
QK norm	true
VAE	mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers
Flow path type	linear

Recommended inference (256×256)

Official NiT sampling defaults for 256×256 class-conditional ImageNet generation:

Setting	Value
Resolution	256×256
Solver	SDE (Euler–Maruyama) in the official repo
Steps (NFE)	250
CFG scale	2.25
CFG interval	(0.0, 0.7)

This Diffusers port uses FlowMatchEulerDiscreteScheduler in deterministic ODE mode (stochastic_sampling=false). Keep the same step count, CFG scale, and interval as the official recipe.

Usage

from pathlib import Path
import torch
from diffusers import DiffusionPipeline

model_dir = Path(".")
pipe = DiffusionPipeline.from_pretrained(
    str(model_dir),
    local_files_only=True,
    custom_pipeline=str(model_dir / "pipeline.py"),
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda" if torch.cuda.is_available() else "cpu")

generator = torch.Generator(device=pipe.device).manual_seed(42)
image = pipe(
    class_labels="golden retriever",
    height=256,
    width=256,
    num_inference_steps=250,
    guidance_scale=2.25,
    guidance_interval=(0.0, 0.7),
    generator=generator,
).images[0]
image.save("demo_256.png")

Components

• pipeline.py — custom NiTPipeline
• model_index.json — pipeline index + ImageNet id2label
• transformer/config.json
• transformer/nit_transformer_2d.py
• transformer/diffusion_pytorch_model.safetensors
• scheduler/scheduler_config.json
• vae/config.json
• vae/diffusion_pytorch_model.safetensors