Better Source, Better Flow: Learning Condition-Dependent Source Distribution for Flow Matching (CSFM)
📄 Paper · 🌐 Project Page · 🤗 Dataset
This repository provides the official PyTorch implementation of Condition-Dependent Source Flow Matching (CSFM).
🚀 05/Feb/26 - Released training and evaluation code of CSFM!
conda create -n csfm python=3.10 -y
conda activate csfm
pip install torch==2.8.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu129
pip install -r requirements.txtImage data: ImageNet-1K dataset.
Text Caption data: CSFM-ImageNet1K-Caption
hf download junwann/CSFM-ImageNet1K-Caption \
--repo-type=dataset \
--local-dir captionsImportant: Make sure that the image paths in the caption files match your local ImageNet data paths.
We used RAE decoders as default
hf download nyu-visionx/RAE-collections \
--local-dir models Important: Make sure the model and stats paths match those defined in the config file.
src/train_t2i.py is compatible with both GPU and TPU(TorchXLA).
enable --use_xla for TPU training. We tested on torch-xla==2.8.1
Configs
- CSFM:
configs/csfm.yaml - Standard FM:
configs/fm.yaml
export ENTITY="your_wandb_entity"
export PROJECT="your_wandb_project"
export WANDB_KEY="your_wandb_key"
exp_name=""
imagenet1k_train_path=""
torchrun --standalone --nnodes=1 --nproc_per_node=4 src/train_t2i.py \
--config configs/csfm.yaml \
--data_path $imagenet1k_train_path \
--exp_name $exp_name \
--jsonl_path captions/train.jsonl \
--jsonl_path_val val_prompt.json \
--exps exps \
--last_ckpt_every 5000 \
--precision bf16 \
--wandb;Experiment Directory Structure
exps_dir/
├─ exp1/
│ ├─ checkpoints/
│ │ ├─ 010000.pt
│ │ └─ last-015000.pt
│ └─ config.yaml
├─ exp2/
│ ├─ checkpoints/
│ │ └─ ...
│ └─ config.yaml
└─ ...Download ImageNet statistics (256×256 shown here) for FID
wget https://openaipublic.blob.core.windows.net/diffusion/jul-2021/ref_batches/imagenet/256/VIRTUAL_imagenet256_labeled.npzSample/Eval Code
exps_dir=''
exp_name=''
sample_dir=''
exp_name=''
train_step=100000
sampling_step=50
batch_per_gpu=250
fid_stat_npz=''
csv_path='metrics.csv'
torchrun --standalone --nnodes=1 --nproc_per_node=4 src/test_t2i.py \
--csv_path $csv_path \
--exps_dir $exps_dir \
--sample_dir $sample_dir \
--exp_name $exp_name \
--train_step $train_step \
--num_samples 50000 \
--sampling_method euler \
--sampling_num_steps $sampling_step \
--per_proc_batch_size $batch_per_gpu \
--fid_stat_npz $fid_stat_npz \
--jsonl_path captions/val.jsonl \
--precision fp32 \
--tf32 \
--global_seed 0 \
To follow the ADM evaluation use --make_npz to create npz file and follow this setup
git clone https://github.com/openai/guided-diffusion.git
cd guided-diffusion/evaluation
conda create -n adm-fid python=3.10
conda activate adm-fid
pip install 'tensorflow[and-cuda]'==2.19 scipy requests tqdm
python evaluator.py VIRTUAL_imagenet256_labeled.npz /path/to/samples.npzThis code is built upon the following repositories:
- RAE: for RAE and diffusion implementation
- Lumina-Image 2.0: for Unified Next-DiT implementation
The toy experiments in the paper are conducted upon following repository:
- C2OT: for toy experiments (not provided in this repo)
@misc{kim2026bettersourcebetterflow,
title={Better Source, Better Flow: Learning Condition-Dependent Source Distribution for Flow Matching},
author={Junwan Kim and Jiho Park and Seonghu Jeon and Seungryong Kim},
year={2026},
eprint={2602.05951},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2602.05951},
}