W-Flow-diffusers

Native diffusers checkpoints for W-Flow (One-Step Generative Modeling via Wasserstein Gradient Flows). Code: hanjq17/W-Flow. Each variant folder is self-contained:

• pipeline.py — WFlowPipeline
• scheduler/scheduling_wflow.py — WFlowScheduler (one evaluation; the transformer predicts the clean latent)
• transformer/transformer_wflow.py — WFlowTransformer2DModel
• vae/ — stabilityai/sd-vae-ft-mse (AutoencoderKL, scaling factor 0.18215)

Load with PyPI diffusers and trust_remote_code=True. No extra W-Flow package is required. Sampling is one forward pass at 256×256. guidance_scale is the official CFG value, which the paper writes as (w+1). Weights are the released EMA. The paper reports ImageNet-256 FID 1.52 for W-Flow-B, 1.35 for W-Flow-L, and 1.29 for W-Flow-XL.

Available checkpoints

Demo

Class: golden retriever (ImageNet 207) — 256×256, 1 step, guidance_scale=1.19, torch_dtype=bfloat16, seed 1.

ImageNet class labels

Each variant keeps an English id2label map in model_index.json.

• pipe.id2label — id → English label
• pipe.labels — synonym → id
• pipe.get_label_ids("golden retriever")
• pipe(class_labels="golden retriever", ...) — string labels resolve automatically

height and width default to 256. The patch embedding is fixed to a 32×32 latent, so other sizes are rejected. On CUDA the generator autocasts to bfloat16.

Inference

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "BiliSakura/W-Flow-diffusers",
    subfolder="W-Flow-B",
    custom_pipeline="pipeline.py",
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to("cuda")

print(pipe.id2label[207])
image = pipe(
    class_labels="golden retriever",
    num_inference_steps=1,
    guidance_scale=1.19,
    generator=torch.Generator(device="cuda").manual_seed(1),
).images[0]
image.save("demo.png")

Local checkout:

from pathlib import Path
import torch
from diffusers import DiffusionPipeline

model_dir = Path("W-Flow-diffusers/W-Flow-B")
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")

Citation

@article{han2026one,
  title   = {One-Step Generative Modeling via Wasserstein Gradient Flows},
  author  = {Han, Jiaqi and Li, Puheng and Guo, Qiushan and Xu, Renyuan and Ermon, Stefano and Cand{\`e}s, Emmanuel J.},
  journal = {arXiv preprint arXiv:2605.11755},
  year    = {2026},
  url     = {https://arxiv.org/abs/2605.11755}
}