PyTorch + GPU implementation

[Friedrich-M]

Thanks for your great work! Do you have any plans to release a PyTorch implementation with GPU support?

[HaoyiZhu]

Hi, I have re-implemented MeanFlow with PyTorch in this repo, which contains:

• A full PyTorch codebase with scripts for training, inference, and evaluation.
• A converted PyTorch checkpoint from the official JAX B/4 release, which can be used directly for evaluation and inference.

[Friedrich-M]

@HaoyiZhu Hi, thanks! Have you tested FID metrics based on your re-implemented version?

[HaoyiZhu]

Hi, I haven't had a chance to do a full retraining from scratch yet. However, I am currently in the process of evaluating the converted checkpoint directly, using the evaluation settings from the REPA repository. I will share the metrics with you here as soon as the evaluation is complete.

[HaoyiZhu]

Hi, this is my evaluation result (tested on an A800 GPU):

Inception Score: 135.00448608398438
FID: 13.386125621490976
sFID: 8.551493777249561
Precision: 0.67536
Recall: 0.4543

It is lower than the reported 11.4 FID in the readme. I am not sure whether this is due to my implementation or due to TPU/GPU inconsistency.

[zhuyu-cs]

Nice, I also reimplemented MeanFlow with PyTorch, and reproduced the results of the paper in this repo.

[qiyan98]

Nice, I also reimplemented MeanFlow with PyTorch, and reproduced the results of the paper in this repo.

Hi @zhuyu-cs , thanks a lot for your contribution! I was wondering if you were able to fully reproduce the results on CIFAR-10 (~3.0 FID)? Thanks.

[zhuyu-cs]

So far, I have not yet reproduced the CIFAR results. I have only trained for 100,000 iterations on the CIFAR branch, achieving an FID of 6.35, as I need more resources to conduct ImageNet experiments.

[antonibigata]

Hi @HaoyiZhu and @zhuyu-cs, thank you to both for your implementation. I was wondering if you found a workaround for the fact that jvp isn't implemented for flash attention (i.e., Dao-AILab/flash-attention#1672) ?
Thanks.

[HaoyiZhu]

Thanks @antonibigata ! I noticed this issue. But I have not verified its implementation yet. I may try it later when I have time. If anyone in the community finds it works, feel free to let me know!

[hu2274898]

Hi @HaoyiZhu and @zhuyu-cs, thank you to both for your implementation. I was wondering if you found a workaround for the fact that jvp isn't implemented for flash attention (i.e., Dao-AILab/flash-attention#1672) ? Thanks.

Hi,does the JAX-based implementation also make use of Flash Attention alongside JAX?
many thanks!