File size: 3,188 Bytes
c8bcefd d39f898 29eeed7 c8bcefd 29eeed7 c8bcefd d39f898 |
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 |
---
license: apache-2.0
library_name: diffusers
datasets:
- ILSVRC/imagenet-1k
---
# Gaussian Mixture Flow Matching Models (GMFlow)
Model used in the paper:
**Gaussian Mixture Flow Matching Models**
<br>
[Hansheng Chen](https://lakonik.github.io/)<sup>1</sup>,
[Kai Zhang](https://kai-46.github.io/website/)<sup>2</sup>,
[Hao Tan](https://research.adobe.com/person/hao-tan/)<sup>2</sup>,
[Zexiang Xu](https://zexiangxu.github.io/)<sup>3</sup>,
[Fujun Luan](https://research.adobe.com/person/fujun/)<sup>2</sup>,
[Leonidas Guibas](https://geometry.stanford.edu/?member=guibas)<sup>1</sup>,
[Gordon Wetzstein](http://web.stanford.edu/~gordonwz/)<sup>1</sup>,
[Sai Bi](https://sai-bi.github.io/)<sup>2</sup><br>
<sup>1</sup>Stanford University, <sup>2</sup>Adobe Research, <sup>3</sup>Hillbot
<br>
[[arXiv](https://arxiv.org/abs/2504.05304)] [[GitHub](https://github.com/Lakonik/GMFlow)]
<img src="gmdit.png" width="600" alt=""/>
<img src="gmdit_results.png" width="1000" alt=""/>
## Usage
Please first install the [official code repository](https://github.com/Lakonik/GMFlow).
We provide a Diffusers pipeline for easy inference. The following code demonstrates how to sample images from the pretrained GM-DiT model using the GM-ODE 2 solver and the GM-SDE 2 solver.
```python
import torch
from huggingface_hub import snapshot_download
from lib.models.diffusions.schedulers import FlowEulerODEScheduler, GMFlowSDEScheduler
from lib.pipelines.gmdit_pipeline import GMDiTPipeline
# Currently the pipeline can only load local checkpoints, so we need to download the checkpoint first
ckpt = snapshot_download(repo_id='Lakonik/gmflow_imagenet_k8_ema')
pipe = GMDiTPipeline.from_pretrained(ckpt, variant='bf16', torch_dtype=torch.bfloat16)
pipe = pipe.to('cuda')
# Pick words that exist in ImageNet
words = ['jay', 'magpie']
class_ids = pipe.get_label_ids(words)
# Sample using GM-ODE 2 solver
pipe.scheduler = FlowEulerODEScheduler.from_config(pipe.scheduler.config)
generator = torch.manual_seed(42)
output = pipe(
class_labels=class_ids,
guidance_scale=0.45,
num_inference_steps=32,
num_inference_substeps=4,
output_mode='mean',
order=2,
generator=generator)
for i, (word, image) in enumerate(zip(words, output.images)):
image.save(f'{i:03d}_{word}_gmode2_step32.png')
# Sample using GM-SDE 2 solver (the first run may be slow due to CUDA compilation)
pipe.scheduler = GMFlowSDEScheduler.from_config(pipe.scheduler.config)
generator = torch.manual_seed(42)
output = pipe(
class_labels=class_ids,
guidance_scale=0.45,
num_inference_steps=32,
num_inference_substeps=1,
output_mode='sample',
order=2,
generator=generator)
for i, (word, image) in enumerate(zip(words, output.images)):
image.save(f'{i:03d}_{word}_gmsde2_step32.png')
```
## Citation
```
@misc{gmflow,
title={Gaussian Mixture Flow Matching Models},
author={Hansheng Chen and Kai Zhang and Hao Tan and Zexiang Xu and Fujun Luan and Leonidas Guibas and Gordon Wetzstein and Sai Bi},
year={2025},
eprint={2504.05304},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2504.05304},
}
``` |