File size: 3,389 Bytes
f43f486 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | ---
license: apache-2.0
library_name: diffusers
pipeline_tag: unconditional-image-generation
tags:
- diffusers
- w-flow
- image-generation
- class-conditional
- one-step
- imagenet
widget:
- text: golden retriever
output:
url: demo.png
language:
- en
---
# W-Flow-diffusers
Native [diffusers](https://github.com/huggingface/diffusers) checkpoints for **W-Flow** ([One-Step Generative Modeling via Wasserstein Gradient Flows](https://arxiv.org/abs/2605.11755)). Code: [hanjq17/W-Flow](https://github.com/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
| Checkpoint | Path | Resolution | Steps | Guidance |
| --- | --- | --- | ---: | ---: |
| W-Flow-B | [`W-Flow-B/`](W-Flow-B/) | 256×256 | 1 | 1.19 |
| W-Flow-L | [`W-Flow-L/`](W-Flow-L/) | 256×256 | 1 | 1.14 |
| W-Flow-XL | [`W-Flow-XL/`](W-Flow-XL/) | 256×256 | 1 | 1.09 |
## 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
```python
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:
```python
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
```bibtex
@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}
}
```
|