EDM2-diffusers

Diffusers-ready checkpoints for EDM2 (Analyzing and Improving the Training Dynamics of Diffusion Models), converted from NVlabs/edm2 post-hoc reconstructions.

Official source weights: https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/

This root folder is a model collection that contains:

• edm2-img512-xs-fid
• edm2-img512-s-fid
• edm2-img512-m-fid
• edm2-img512-l-fid
• edm2-img512-l-dino
• edm2-img512-xl-fid
• edm2-img512-xxl-fid

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

• pipeline.py
• unet/unet_edm2.py
• scheduler/scheduler_config.json (EDMEulerScheduler)
• unet/diffusion_pytorch_model.safetensors
• vae/diffusion_pytorch_model.safetensors

Demo

Class-conditional sample (ImageNet class 207, golden retriever), EDM2-XXL at 512×512, 32 steps, guidance 1.0, seed 42.

Model Paths

Use paths relative to this root README:

Model	NVlabs preset	FID	Local path
EDM2-XS	edm2-img512-xs-fid	3.53	./edm2-img512-xs-fid
EDM2-S	edm2-img512-s-fid	2.56	./edm2-img512-s-fid
EDM2-M	edm2-img512-m-fid	2.25	./edm2-img512-m-fid
EDM2-L	edm2-img512-l-fid	2.06	./edm2-img512-l-fid
EDM2-L (DINO)	edm2-img512-l-dino	—	./edm2-img512-l-dino
EDM2-XL	edm2-img512-xl-fid	1.96	./edm2-img512-xl-fid
EDM2-XXL	edm2-img512-xxl-fid	1.91	./edm2-img512-xxl-fid

Inference Demo (Diffusers)

1) Load a local subfolder checkpoint

from pathlib import Path
import torch
from diffusers import DiffusionPipeline

model_dir = Path("./edm2-img512-xxl-fid")  # change to any path in the table above
pipe = DiffusionPipeline.from_pretrained(
    str(model_dir),
    local_files_only=True,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")

generator = torch.Generator(device="cuda").manual_seed(42)
image = pipe(
    class_labels=207,          # golden retriever (ImageNet id); omit for random class
    num_inference_steps=32,
    guidance_scale=1.0,        # >1.0 requires a gnet/ checkpoint
    generator=generator,
).images[0]
image.save("demo.png")

Official inference defaults (generate_images.py): num_steps=32, sigma_min=0.002, sigma_max=80, rho=7, guidance=1.0 (no gnet), S_churn=0. Heun sampling runs in float32 internally even when UNet/VAE weights are loaded in bf16/fp16.

Guided presets require a converted gnet/ folder and guidance_scale matching the NVlabs preset.

2) Convert a legacy .pkl

python scripts/convert_edm2_to_diffusers.py \
  --checkpoint models/BiliSakura/EDM2-diffusers/edm2-img512-xs-2147483-0.135.pkl \
  --output models/BiliSakura/EDM2-diffusers

Creates edm2-img512-xs-fid/ automatically from the NVlabs preset mapping.

Checkpoint preset mapping

Maps NVlabs --preset=... names from generate_images.py to source pickle filenames and local Diffusers directories.

EDM2 paper — ImageNet-512 (conditional)

NVlabs preset	Source .pkl (net)	Diffusers dir	Metric
edm2-img512-xs-fid	edm2-img512-xs-2147483-0.135.pkl	edm2-img512-xs-fid/	FID 3.53
edm2-img512-xs-dino	edm2-img512-xs-2147483-0.200.pkl	—	FDDINOv2 103.39
edm2-img512-s-fid	edm2-img512-s-2147483-0.130.pkl	edm2-img512-s-fid/	FID 2.56
edm2-img512-s-dino	edm2-img512-s-2147483-0.190.pkl	—	FDDINOv2 68.64
edm2-img512-m-fid	edm2-img512-m-2147483-0.100.pkl	edm2-img512-m-fid/	FID 2.25
edm2-img512-m-dino	edm2-img512-m-2147483-0.155.pkl	—	FDDINOv2 58.44
edm2-img512-l-fid	edm2-img512-l-1879048-0.085.pkl	edm2-img512-l-fid/	FID 2.06
edm2-img512-l-dino	edm2-img512-l-1879048-0.155.pkl	edm2-img512-l-dino/	FDDINOv2 52.25
edm2-img512-xl-fid	edm2-img512-xl-1342177-0.085.pkl	edm2-img512-xl-fid/	FID 1.96
edm2-img512-xl-dino	edm2-img512-xl-1342177-0.155.pkl	—	FDDINOv2 45.96
edm2-img512-xxl-fid	edm2-img512-xxl-0939524-0.070.pkl	edm2-img512-xxl-fid/	FID 1.91
edm2-img512-xxl-dino	edm2-img512-xxl-0939524-0.150.pkl	—	FDDINOv2 42.84

EDM2 paper — ImageNet-64 (conditional)

NVlabs preset	Source .pkl (net)	Metric
edm2-img64-s-fid	edm2-img64-s-1073741-0.075.pkl	FID 1.58
edm2-img64-m-fid	edm2-img64-m-2147483-0.060.pkl	FID 1.43
edm2-img64-l-fid	edm2-img64-l-1073741-0.040.pkl	FID 1.33
edm2-img64-xl-fid	edm2-img64-xl-0671088-0.040.pkl	FID 1.33

EDM2 paper — classifier-free guidance (ImageNet-512)

Use guidance_scale below and include the converted gnet/ checkpoint.

NVlabs preset	Source .pkl (net)	Source .pkl (gnet)	Guidance	Metric
edm2-img512-xs-guid-fid	edm2-img512-xs-2147483-0.045.pkl	edm2-img512-xs-uncond-2147483-0.045.pkl	1.40	FID 2.91
edm2-img512-xs-guid-dino	edm2-img512-xs-2147483-0.150.pkl	edm2-img512-xs-uncond-2147483-0.150.pkl	1.70	FDDINOv2 79.94
edm2-img512-s-guid-fid	edm2-img512-s-2147483-0.025.pkl	edm2-img512-xs-uncond-2147483-0.025.pkl	1.40	FID 2.23
edm2-img512-s-guid-dino	edm2-img512-s-2147483-0.085.pkl	edm2-img512-xs-uncond-2147483-0.085.pkl	1.90	FDDINOv2 52.32
edm2-img512-m-guid-fid	edm2-img512-m-2147483-0.030.pkl	edm2-img512-xs-uncond-2147483-0.030.pkl	1.20	FID 2.01
edm2-img512-m-guid-dino	edm2-img512-m-2147483-0.015.pkl	edm2-img512-xs-uncond-2147483-0.015.pkl	2.00	FDDINOv2 41.98
edm2-img512-l-guid-fid	edm2-img512-l-1879048-0.015.pkl	edm2-img512-xs-uncond-2147483-0.015.pkl	1.20	FID 1.88
edm2-img512-l-guid-dino	edm2-img512-l-1879048-0.035.pkl	edm2-img512-xs-uncond-2147483-0.035.pkl	1.70	FDDINOv2 38.20
edm2-img512-xl-guid-fid	edm2-img512-xl-1342177-0.020.pkl	edm2-img512-xs-uncond-2147483-0.020.pkl	1.20	FID 1.85
edm2-img512-xl-guid-dino	edm2-img512-xl-1342177-0.030.pkl	edm2-img512-xs-uncond-2147483-0.030.pkl	1.70	FDDINOv2 35.67
edm2-img512-xxl-guid-fid	edm2-img512-xxl-0939524-0.015.pkl	edm2-img512-xs-uncond-2147483-0.015.pkl	1.20	FID 1.81
edm2-img512-xxl-guid-dino	edm2-img512-xxl-0939524-0.015.pkl	edm2-img512-xs-uncond-2147483-0.015.pkl	1.70	FDDINOv2 33.09

Autoguidance paper

NVlabs preset	Source .pkl (net)	Source .pkl (gnet)	Guidance	Metric
edm2-img512-s-autog-fid	edm2-img512-s-2147483-0.070.pkl	edm2-img512-xs-0134217-0.125.pkl	2.10	FID 1.34
edm2-img512-s-autog-dino	edm2-img512-s-2147483-0.120.pkl	edm2-img512-xs-0134217-0.165.pkl	2.45	FDDINOv2 36.67
edm2-img512-xxl-autog-fid	edm2-img512-xxl-0939524-0.075.pkl	edm2-img512-m-0268435-0.155.pkl	2.05	FID 1.25
edm2-img512-xxl-autog-dino	edm2-img512-xxl-0939524-0.130.pkl	edm2-img512-m-0268435-0.205.pkl	2.30	FDDINOv2 24.18
edm2-img512-s-uncond-autog-fid	edm2-img512-s-uncond-2147483-0.070.pkl	edm2-img512-xs-uncond-0134217-0.110.pkl	2.85	FID 3.86
edm2-img512-s-uncond-autog-dino	edm2-img512-s-uncond-2147483-0.090.pkl	edm2-img512-xs-uncond-0134217-0.125.pkl	2.90	FDDINOv2 90.39
edm2-img64-s-autog-fid	edm2-img64-s-1073741-0.045.pkl	edm2-img64-xs-0134217-0.110.pkl	1.70	FID 1.01
edm2-img64-s-autog-dino	edm2-img64-s-1073741-0.105.pkl	edm2-img64-xs-0134217-0.175.pkl	2.20	FDDINOv2 31.85

NVlabs preset shorthand

# EDM2 paper
edm2-img512-{xs|s|m|l|xl|xxl}-{fid|dino}
edm2-img64-{s|m|l|xl}-fid
edm2-img512-{xs|s|m|l|xl|xxl}-guid-{fid|dino}

# Autoguidance paper
edm2-img512-{s|xxl}-autog-{fid|dino}
edm2-img512-s-uncond-autog-{fid|dino}
edm2-img64-s-autog-{fid|dino}

Example NVlabs command:

python generate_images.py --preset=edm2-img512-s-guid-dino --outdir=out

Equivalent expanded form:

python generate_images.py \
  --net=https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/edm2-img512-s-2147483-0.085.pkl \
  --gnet=https://nvlabs-fi-cdn.nvidia.com/edm2/posthoc-reconstructions/edm2-img512-xs-uncond-2147483-0.085.pkl \
  --guidance=1.9 \
  --outdir=out