NegFlow

A flow matching framework for utilizing negative samples.

A lightweight SiT flow matching codebase. This codebase is mainly built on top of SiT and REPA, but we make the following adjustments:

1. Simplicity. This codebase implements only the linear path and DINOv2-ViT-B as the external encoder .

2. Data Process. Compared with REPA, we preprocess the data on the fly, and use the random filp augmentation, which is consistent with SiT.

3. Time Embeddings. we scale it to [0,1000] while SiT doesn't.

4. Accleration. Compile and preprocess.

5. Sampling. Mainly use EM solver. And by default we apply cfg on all channels.

Preprocess:

Step 1 (Optional): Convert ImageNet to LMDB format using 'preprocess_imagenet/image2lmdb.py':

python image2lmdb.py 

Step 2 (Optional): Encode images to VAE latents:

cd ./preprocess
torchrun --nproc_per_node=8 --nnodes=1 --node_rank=0 \
    main_cache.py \
    --source_lmdb  ../data/imagenet/imagenet_train_lmdb \
    --target_lmdb ../data/imagenet/train_vae_latents_lmdb \
    --img_size 256 \
    --batch_size 128 \
    --lmdb_size_gb 400

Env setup:

conda create -n negfm python=3.10.0
conda activate negfm
pip install torch==2.1.0 torchvision==0.16.0 --index-url https://download.pytorch.org/whl/cu121
pip install -r requirements.txt

Train:

accelerate launch --mixed_precision=bf16 --num_processes=4  train.py  --model_type 'SiT-S/2' --flow_type 'standard' --neg_type 'res xt'  --neg_policy 'random' --neg_lam 0.20 --use_latent; 

Sample:

TODO:

accelerate launch --mixed_precision=bf16 --num_processes=4  train.py  --model_type 'SiT-XL/2' --flow_type 'standard' --neg_type 'res xt'  --neg_policy 'random' --neg_lam 0.20 --use_latent --use_repa; 


accelerate launch --mixed_precision=bf16 --num_processes=4  train.py  --model_type 'SiT-XL/2' --flow_type 'neg' --neg_type 'res xt'  --neg_policy 'random' --neg_lam 0.20 --use_latent --use_repa;