FiTv2-3B-2-256/README.md

metadata

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

FiTv2-3B-2-256

Self-contained Diffusers checkpoint for FiTv2-3B/2, converted from InfImagine/FiT.

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

• model_index.json (includes ImageNet id2label)
• pipeline.py (custom FiTv2Pipeline)
• transformer/fit_transformer_2d.py and weights
• scheduler/scheduler_config.json (FlowMatchEulerDiscreteScheduler)
• vae/diffusion_pytorch_model.safetensors

Recommended inference (256×256)

Setting	Value
Resolution	256×256
Sampler	flow matching (velocity ODE)
Steps	250
CFG scale	1.5
Dtype	float32 (or bfloat16 on Ampere+)
VAE	stabilityai/sd-vae-ft-ema (bundled under vae/)

from pathlib import Path
import torch
from diffusers import DiffusionPipeline

model_dir = Path("./FiTv2-3B-2-256").resolve()
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,
)
pipe.to("cuda")

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",
    height=256,
    width=256,
    num_inference_steps=250,
    guidance_scale=1.5,
    generator=generator,
).images[0]
image.save("demo.png")