ebgu / bayesian_flow_networks

A replication of Bayesian Flow Networks paper with pytorch and ViT

Bayesian_Flow_Networks

A replication of Bayesian_Flow_Networks paper with PyTorch and U-ViT.

To train a new model, you can modify the yaml file and:

python multi_gpu_trainer.py example

Training data of Oxford Flowers should be split manually, and you can find the numpy version of their labels in this repo.

To run inference, please download my pretrained weight:

python sample_img.py --device "cuda:0" --load "last" --SavedDir tmp/ --ExpConfig BFN_example/BFN_example.yaml --n_sqrt 8 --steps 200

The inference process is controled by 6 parameters :

"device", usually 'cuda:0' ;

"load", best epoch or last epoch;

"SavedDir", where to save images;

"ExpConfig", the yaml file of your experiments;

"n_sqrt", you will get N² samples for each class;

"steps", n steps for sampling, the orignal paper recommands 25, but in my experiment, 200 is better.

The result should looks like the welcoming images.

Enjoy!