{"id":17714302,"url":"https://github.com/yuanzhi-zhu/SlimFlow","last_synced_at":"2025-03-13T22:32:21.421Z","repository":{"id":248914334,"uuid":"829066974","full_name":"yuanzhi-zhu/SlimFlow","owner":"yuanzhi-zhu","description":"\"SlimFlow: Training Smaller One-Step Diffusion Models with Rectified Flow\", Yuanzhi Zhu, Xingchao Liu, Qiang Liu","archived":false,"fork":false,"pushed_at":"2024-11-26T23:44:33.000Z","size":77,"stargazers_count":40,"open_issues_count":1,"forks_count":0,"subscribers_count":4,"default_branch":"main","last_synced_at":"2024-11-27T00:27:28.674Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"https://arxiv.org/abs/2407.12718","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/yuanzhi-zhu.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-07-15T17:29:50.000Z","updated_at":"2024-11-26T23:44:36.000Z","dependencies_parsed_at":"2024-07-17T22:28:37.325Z","dependency_job_id":"109f1c97-c4ab-4e81-af38-23f1ea3c729d","html_url":"https://github.com/yuanzhi-zhu/SlimFlow","commit_stats":null,"previous_names":["yuanzhi-zhu/slimflow"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhi-zhu%2FSlimFlow","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhi-zhu%2FSlimFlow/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhi-zhu%2FSlimFlow/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/yuanzhi-zhu%2FSlimFlow/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/yuanzhi-zhu","download_url":"https://codeload.github.com/yuanzhi-zhu/SlimFlow/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":243494530,"owners_count":20299829,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-10-25T11:02:20.869Z","updated_at":"2025-03-13T22:32:21.416Z","avatar_url":"https://github.com/yuanzhi-zhu.png","language":"Python","funding_links":[],"categories":["Accelerate"],"sub_categories":[],"readme":"# SlimFlow\n\nThis is the official implementation of ECCV2024 paper\n## [SlimFlow: Training Smaller One-Step Diffusion Models with Rectified Flow](https://arxiv.org/abs/2407.12718) \nby [Yuanzhi Zhu](https://yuanzhi-zhu.github.io/about/), [Xingcaho Liu](https://gnobitab.github.io/), [Qiang Liu](https://www.cs.utexas.edu/~lqiang/)\n\nThis code is based on [RectifiedFlow](https://github.com/gnobitab/RectifiedFlow).\n\n## usage\n\n### Train 1-Rectified Flow\n\n```bash\npython ./train.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py  \\\n    --config.expr 1_rectified_flow \\\n```\n\n### Evaluation \n\nevaluate FID of ckpts from `config.eval.begin_ckpt` in `ckpt_dir`\n#### one step\n```bash\npython ./evaluation_fid.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py  \\\n    --ckpt_dir logs/1_rectified_flow \\\n    --config.eval.batch_size 512 --config.eval.num_samples 50000 \\\n    --config.eval.begin_ckpt 1 --config.eval.end_ckpt 0 \\\n    --config.sampling.sample_N 1 --config.sampling.use_ode_sampler euler \\\n```\n\n#### rk45 by default\n```bash\npython ./evaluation_fid.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py  \\\n    --ckpt_dir logs/1_rectified_flow \\\n    --config.eval.batch_size 512 --config.eval.num_samples 50000 \\\n    --config.eval.begin_ckpt 1 --config.eval.end_ckpt 0 \\\n```\n\n### Image Sampling\n\nsampling all ckpts in `sampling_dir`\n```bash\npython ./image_sampling.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py \\\n    --sampling_dir \"logs/1_rectified_flow\" \\\n    --config.eval.batch_size 64\n```\n\n##### Image Sampling Configurations\n- Sample from 1flows: `--config.sampling.use_ode_sampler rk45`\n- Sample from 2flows: [`--config.sampling.use_ode_sampler rk45`, `--config.sampling.use_ode_sampler heun` + `--config.sampling.sample_N 3`, `--config.sampling.use_ode_sampler euler` + `--config.sampling.sample_N 1`]\n- Sample from distilled one-step models: `--config.sampling.use_ode_sampler euler` + `--config.sampling.sample_N 1`\n\n##### Model Configurations\n- ImageNet64 80.7M: `--config.model.name DhariwalUNet --config.model.nf 128 --config.model.num_res_blocks 2 --config.model.ch_mult '(1, 2, 2, 4)' --config.data.num_classes 1000 --config.data.image_size 64 --config.model.attn_resolutions '32, 16'`\n- ImageNet 44.7MM: `--config.model.name DhariwalUNet --config.model.nf 128 --config.model.num_res_blocks 2 --config.model.ch_mult '(1, 2, 2, 2)' --config.data.num_classes 1000 --config.data.image_size 64 --config.model.attn_resolutions '32, 16'`\n- FFHQ64 27.9M: `--config.model.nf 128 --config.model.num_res_blocks 2 --config.data.image_size 64 --config.model.ch_mult '(1, 2, 2)'`\n- FFHQ64 15.7M: `--config.model.nf 96 --config.model.num_res_blocks 2 --config.data.image_size 64 --config.model.ch_mult '(1, 2, 2)'`\n- FFHQ64 7.0M: `--config.model.nf 64 --config.model.num_res_blocks 2 --config.data.image_size 64 --config.model.ch_mult '(1, 2, 2)'`\n- FFHQ64 3.4M: `--config.model.nf 64 --config.model.num_res_blocks 1 --config.data.image_size 64 --config.model.ch_mult '(1, 1, 2)'`\n- CIFAR32 27.9M: `--config.model.nf 128 --config.model.num_res_blocks 2 --config.data.image_size 32 --config.model.ch_mult '(1, 2, 2)'`\n- CIFAR32 15.7M: `--config.model.nf 96 --config.model.num_res_blocks 2 --config.data.image_size 32 --config.model.ch_mult '(1, 2, 2)'`\n- CIFAR32 7.0M: `--config.model.nf 64 --config.model.num_res_blocks 2 --config.data.image_size 32 --config.model.ch_mult '(1, 2, 2)'`\n- CIFAR32 3.4M: `--config.model.nf 64 --config.model.num_res_blocks 1 --config.data.image_size 32 --config.model.ch_mult '(1, 1, 2)'`\n\n### Generate Data Pair\n\n#### z0--\u003ez1 by default\n```bash\npython ./generate_data.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py  \\\n    --ckpt_path \"logs/1_rectified_flow/checkpoints/checkpoint_14.pth\" \\\n    --data_root \"reflow_data/1_rectified_flow_50000/\" \\\n    --config.sampling.total_number_of_samples 50000 --config.seed 0 \\\n    --config.training.batch_size 512 \\\n    --config.sampling.direction from_z0 \\\n```\n\n`config.sampling.direction` has 3 options: 'from_z0', 'from_z1', 'random_paired'\n\n\n### Reflow to get 2-Rectified Flow with the Generated Data Pair\n\n```bash\npython ./train.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py  \\\n    --config.data.reflow_data_root \"reflow_data/1_rectified_flow_50000/\" \\\n    --config.flow.flow_t_schedule uniform \\\n    --config.expr 2_rectified_flow \\\n    --config.flow.h_flip=true \\\n    --config.flow.pre_train_model /logs/1_rectified_flow/checkpoints/checkpoint_14.pth \\\n```\n\n### Annealing Reflow\n\n```bash\npython ./train.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py  \\\n    --config.expr 2_rectified_flow_500001flow_flip_warmup_300000_28m \\\n    --config.flow.h_flip=true \\\n    --config.training.x0_randomness warmup_300000 \\\n    --config.training.snapshot_freq 50000 \\\n    --config.training.snapshot_sampling 10000 \\\n    --config.data.reflow_data_root \"reflow_data/1_rectified_flow_50000/\" \\\n    --config.model.nf 128 --config.model.num_res_blocks 2 \\\n    --config.model.ch_mult '(1, 2, 2)' \\\n```\n\nmust specify `config.data.data_root` for reflow training\n\nif `config.flow.pre_train_model` is not specified, the model will be trained from scratch.\n\n### Distill to get one-step model\n\u003c!-- distillation as special case of reflow with different `flow_t_schedule` and `flow_alpha_t` --\u003e\n\n```bash\npython ./train.py \\\n    --config ./configs/rectified_flow/cifar10_rf_gaussian.py  \\\n    --config.data.reflow_data_root \"reflow_data/1_rectified_flow_50000/\" \\\n    --config.flow.flow_t_schedule t0 \\\n    --config.training.loss_type lpips \\\n    --config.flow.use_teacher true \\\n    --config.expr 2_rectified_flow_500000bigflow_28m_distill_lpips_use_teacher \\\n    --config.flow.pre_train_model \"./logs/2_rectified_flow_500001flow_flip_warmup_300000_28m/checkpoints/checkpoint_16.pth\" \\\n    --config.model.nf 128 --config.model.num_res_blocks 2 \\\n    --config.model.ch_mult '(1, 2, 2)' \\\n```\n\n## Checkpoints\ncheckpoints can be found here on HuggingFace: https://huggingface.co/Yuanzhi/SlimFlow\nTo sample from these checkpoints, please follow the instructions in the README.md of the HuggingFace model.\n\n## Citation\nIf you find this repo helpful, please cite:\n\n```bibtex\n@inproceedings{zhu2025slimflow,\n  title={SlimFlow: Training Smaller One-Step Diffusion Models with Rectified Flow},\n  author={Zhu, Yuanzhi and Liu, Xingchao and Liu, Qiang},\n  booktitle={European Conference on Computer Vision},\n  pages={342--359},\n  year={2025},\n  organization={Springer}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyuanzhi-zhu%2FSlimFlow","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fyuanzhi-zhu%2FSlimFlow","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fyuanzhi-zhu%2FSlimFlow/lists"}