BiliSakura / EDM2-diffusers

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:

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)

EDM2 paper — ImageNet-64 (conditional)

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

Use guidance_scale below and include the converted gnet/ checkpoint.

Autoguidance paper

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