| --- |
| 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} |
| } |
| ``` |
|
|
