iMF-diffusers

Native diffusers implementation of Improved Mean Flows (iMF). Each variant folder is self-contained:

• pipeline.py — IMFPipeline
• scheduler/scheduler_config.json — FlowMatchEulerDiscreteScheduler config
• transformer/transformer_imf.py — IMFTransformer2DModel
• vae/ — bundled stabilityai/sd-vae-ft-mse (AutoencoderKL)

Demo

Class-conditional sample (ImageNet class 207, golden retriever), iMF-XL/2 at 256×256, 1 step, CFG 1.8, interval [0.0, 1.0], seed 42.

Available checkpoints

Checkpoint	Path	Latent size	FID eval CFG (ω)	FID eval interval
iMF-B/2	./iMF-B-2	32×32	8.0	[0.40, 0.65]
iMF-L/2	./iMF-L-2	32×32	10.5	[0.40, 0.60]
iMF-XL/2	./iMF-XL-2	32×32	8.0	[0.42, 0.62]

FID eval settings follow upstream imeanflow eval config.

Inference

from pathlib import Path
from diffusers import DiffusionPipeline
import torch

model_dir = Path("./iMF-XL-2")
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")

generator = torch.Generator(device="cuda").manual_seed(42)
image = pipe(
    class_labels="golden retriever",
    num_inference_steps=1,
    guidance_scale=1.8,
    guidance_interval_start=0.0,
    guidance_interval_end=1.0,
    generator=generator,
).images[0]
image.save("demo.png")

Load a variant subfolder (e.g. ./iMF-XL-2), not the repo root.