ciaohe / dim

Diffusion Image Mamba

Install

Install Mamba-ssm

Refer to Vision Mamba

# Install causal_conv1d and mamba
pip install -e causal_conv1d>=1.1.0
pip install -e mamba-1p1p1

Train

# Assign HF_HOME to the cache directory
export HF_HOME="/comp_robot/rentianhe/caohe/cache"

# Train with 4 GPUs
accelerate launch --multi_gpu --num_processes 4 --mixed_precision fp16 train_mamba.py --model DiM-S/2 --feature-path /shared_space/caohe/DATA/imagenet1k/train_vae --lr 5e-4

• feature-path: the path to the pre-extracted features (use SD-VAE compression rate as 8)
• model: the model name (Now supported DiM-[S|M|L|XL]/[2|4|8])
• [default] batch size : global batch size = 256 (constant for all models, suggested in DiT)
• [optional] --lr : learning rate (default 1e-4, for mamba, 5e-4 is suggested)
• [optional] --ckpt-every: default to 50_000 (= 10 epochs)

The ckpt will be saved under results/xxx-MODEL_VERSION/checkpoints/

Evaluation

Sampling scripts

# Assign HF_HOME to the cache directory
export HF_HOME="/comp_robot/rentianhe/caohe/cache"
# 
MODEL_VERSION=005-DiM-S-2 # The model version
CKPT=0100000 # The checkpoint number
torchrun --nnodes=1 --nproc_per_node=4 sample_ddp.py --model DiM-S/2 --num-fid-samples 50000 --ckpt results/$MODEL_VERSION/checkpoints/$CKPT.pt --sample-dir samples/$MODEL_VERSION-$CKPT --per-proc-batch-size 64

FID evaluation

First need to install related packages following ADM's TensorFlow evaluation suite

IN_50K_REFERENCE=/shared_space/caohe/DATA/imagenet1k/VIRTUAL_imagenet256_labeled.npz
# ! adapt the path to the sample features path
python evaluations/evaluator.py $IN_50K_REFERENCE samples/DiM-B-2-ckpt-0100000/DiM-B-2-0100000-size-256-vae-ema-cfg-1.5-seed-0.npz

dim's Issues

about Image sampling

Thanks for the Mamba based diffusion model code. Has your team successfully used this code to generate images? Thank you again for your selfless sharing and I look forward to your reply. Wish you the best at work!