├── README.md ├── assets ├── selection.png ├── teaser1024.png ├── teaser256.png └── teaser512.png ├── benchmark_generation_mamba_simple.py ├── benchmarks └── benchmark_generation_mamba_simple.py ├── configs ├── cifar10_S_DiM.py ├── imagenet256_H_DiM.py ├── imagenet256_L_DiM.py ├── imagenet512_H_DiM_ft.py ├── imagenet512_H_DiM_upsample_3x_test.py └── imagenet512_H_DiM_upsample_test.py ├── csrc └── selective_scan │ ├── reverse_scan.cuh │ ├── selective_scan.cpp │ ├── selective_scan.h │ ├── selective_scan_bwd_bf16_complex.cu │ ├── selective_scan_bwd_bf16_real.cu │ ├── selective_scan_bwd_fp16_complex.cu │ ├── selective_scan_bwd_fp16_real.cu │ ├── selective_scan_bwd_fp32_complex.cu │ ├── selective_scan_bwd_fp32_real.cu │ ├── selective_scan_bwd_kernel.cuh │ ├── selective_scan_common.h │ ├── selective_scan_fwd_bf16.cu │ ├── selective_scan_fwd_fp16.cu │ ├── selective_scan_fwd_fp32.cu │ ├── selective_scan_fwd_kernel.cuh │ ├── static_switch.h │ └── uninitialized_copy.cuh ├── dpm_solver_pp.py ├── dpm_solver_pytorch.py ├── environment.yaml ├── eval.py ├── eval_ldm.py ├── eval_ldm_discrete.py ├── eval_t2i_discrete.py ├── evals └── lm_harness_eval.py ├── libs ├── __init__.py ├── autoencoder.py ├── clip.py ├── timm.py ├── uvit.py └── uvit_t2i.py ├── main.pdf ├── main.png ├── mamba_attn_diff ├── models │ ├── __init__.py │ ├── adapter_attn4mamba.py │ ├── attention.py │ ├── freeu.py │ ├── mamba_2d.py │ ├── normalization.py │ ├── upsample_guidance.py │ └── vim_module.py └── utils │ ├── backup_code.py │ └── init_weights.py ├── mamba_ssm ├── __init__.py ├── models │ ├── __init__.py │ ├── config_mamba.py │ └── mixer_seq_simple.py ├── modules │ ├── __init__.py │ └── mamba_simple.py ├── ops │ ├── __init__.py │ ├── selective_scan_interface.py │ └── triton │ │ ├── __init__.py │ │ ├── layernorm.py │ │ └── selective_state_update.py └── utils │ ├── __init__.py │ ├── generation.py │ └── hf.py ├── sample_t2i_discrete.py ├── scripts ├── extract_empty_feature.py ├── extract_imagenet_feature.py ├── extract_mscoco_feature.py └── extract_test_prompt_feature.py ├── sde.py ├── setup.py ├── tools ├── __init__.py ├── fid_score.py └── inception.py ├── train.py ├── train_ldm.py ├── train_ldm_discrete.py ├── train_t2i_discrete.py ├── utils.py └── uvit_datasets.py /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/README.md -------------------------------------------------------------------------------- /assets/selection.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/assets/selection.png -------------------------------------------------------------------------------- /assets/teaser1024.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/assets/teaser1024.png -------------------------------------------------------------------------------- /assets/teaser256.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/assets/teaser256.png -------------------------------------------------------------------------------- /assets/teaser512.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/assets/teaser512.png -------------------------------------------------------------------------------- /benchmark_generation_mamba_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/benchmark_generation_mamba_simple.py -------------------------------------------------------------------------------- /benchmarks/benchmark_generation_mamba_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/benchmarks/benchmark_generation_mamba_simple.py -------------------------------------------------------------------------------- /configs/cifar10_S_DiM.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/configs/cifar10_S_DiM.py -------------------------------------------------------------------------------- /configs/imagenet256_H_DiM.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/configs/imagenet256_H_DiM.py -------------------------------------------------------------------------------- /configs/imagenet256_L_DiM.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/configs/imagenet256_L_DiM.py -------------------------------------------------------------------------------- /configs/imagenet512_H_DiM_ft.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/configs/imagenet512_H_DiM_ft.py -------------------------------------------------------------------------------- /configs/imagenet512_H_DiM_upsample_3x_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/configs/imagenet512_H_DiM_upsample_3x_test.py -------------------------------------------------------------------------------- /configs/imagenet512_H_DiM_upsample_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/configs/imagenet512_H_DiM_upsample_test.py -------------------------------------------------------------------------------- /csrc/selective_scan/reverse_scan.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/reverse_scan.cuh -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan.cpp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan.cpp -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan.h -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_bf16_complex.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_bwd_bf16_complex.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_bf16_real.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_bwd_bf16_real.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp16_complex.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp16_complex.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp16_real.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp16_real.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp32_complex.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp32_complex.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_fp32_real.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_bwd_fp32_real.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_bwd_kernel.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_bwd_kernel.cuh -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_common.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_common.h -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_bf16.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_fwd_bf16.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_fp16.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_fwd_fp16.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_fp32.cu: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_fwd_fp32.cu -------------------------------------------------------------------------------- /csrc/selective_scan/selective_scan_fwd_kernel.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/selective_scan_fwd_kernel.cuh -------------------------------------------------------------------------------- /csrc/selective_scan/static_switch.h: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/static_switch.h -------------------------------------------------------------------------------- /csrc/selective_scan/uninitialized_copy.cuh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/csrc/selective_scan/uninitialized_copy.cuh -------------------------------------------------------------------------------- /dpm_solver_pp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/dpm_solver_pp.py -------------------------------------------------------------------------------- /dpm_solver_pytorch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/dpm_solver_pytorch.py -------------------------------------------------------------------------------- /environment.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/environment.yaml -------------------------------------------------------------------------------- /eval.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/eval.py -------------------------------------------------------------------------------- /eval_ldm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/eval_ldm.py -------------------------------------------------------------------------------- /eval_ldm_discrete.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/eval_ldm_discrete.py -------------------------------------------------------------------------------- /eval_t2i_discrete.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/eval_t2i_discrete.py -------------------------------------------------------------------------------- /evals/lm_harness_eval.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/evals/lm_harness_eval.py -------------------------------------------------------------------------------- /libs/__init__.py: -------------------------------------------------------------------------------- 1 | # codes from third party 2 | -------------------------------------------------------------------------------- /libs/autoencoder.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/libs/autoencoder.py -------------------------------------------------------------------------------- /libs/clip.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/libs/clip.py -------------------------------------------------------------------------------- /libs/timm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/libs/timm.py -------------------------------------------------------------------------------- /libs/uvit.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/libs/uvit.py -------------------------------------------------------------------------------- /libs/uvit_t2i.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/libs/uvit_t2i.py -------------------------------------------------------------------------------- /main.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/main.pdf -------------------------------------------------------------------------------- /main.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/main.png -------------------------------------------------------------------------------- /mamba_attn_diff/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_attn_diff/models/adapter_attn4mamba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/models/adapter_attn4mamba.py -------------------------------------------------------------------------------- /mamba_attn_diff/models/attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/models/attention.py -------------------------------------------------------------------------------- /mamba_attn_diff/models/freeu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/models/freeu.py -------------------------------------------------------------------------------- /mamba_attn_diff/models/mamba_2d.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/models/mamba_2d.py -------------------------------------------------------------------------------- /mamba_attn_diff/models/normalization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/models/normalization.py -------------------------------------------------------------------------------- /mamba_attn_diff/models/upsample_guidance.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/models/upsample_guidance.py -------------------------------------------------------------------------------- /mamba_attn_diff/models/vim_module.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/models/vim_module.py -------------------------------------------------------------------------------- /mamba_attn_diff/utils/backup_code.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/utils/backup_code.py -------------------------------------------------------------------------------- /mamba_attn_diff/utils/init_weights.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_attn_diff/utils/init_weights.py -------------------------------------------------------------------------------- /mamba_ssm/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/__init__.py -------------------------------------------------------------------------------- /mamba_ssm/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/models/config_mamba.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/models/config_mamba.py -------------------------------------------------------------------------------- /mamba_ssm/models/mixer_seq_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/models/mixer_seq_simple.py -------------------------------------------------------------------------------- /mamba_ssm/modules/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/modules/mamba_simple.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/modules/mamba_simple.py -------------------------------------------------------------------------------- /mamba_ssm/ops/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/ops/selective_scan_interface.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/ops/selective_scan_interface.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/layernorm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/ops/triton/layernorm.py -------------------------------------------------------------------------------- /mamba_ssm/ops/triton/selective_state_update.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/ops/triton/selective_state_update.py -------------------------------------------------------------------------------- /mamba_ssm/utils/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /mamba_ssm/utils/generation.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/utils/generation.py -------------------------------------------------------------------------------- /mamba_ssm/utils/hf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/mamba_ssm/utils/hf.py -------------------------------------------------------------------------------- /sample_t2i_discrete.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/sample_t2i_discrete.py -------------------------------------------------------------------------------- /scripts/extract_empty_feature.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/scripts/extract_empty_feature.py -------------------------------------------------------------------------------- /scripts/extract_imagenet_feature.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/scripts/extract_imagenet_feature.py -------------------------------------------------------------------------------- /scripts/extract_mscoco_feature.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/scripts/extract_mscoco_feature.py -------------------------------------------------------------------------------- /scripts/extract_test_prompt_feature.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/scripts/extract_test_prompt_feature.py -------------------------------------------------------------------------------- /sde.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/sde.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/setup.py -------------------------------------------------------------------------------- /tools/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tools/fid_score.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/tools/fid_score.py -------------------------------------------------------------------------------- /tools/inception.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/tools/inception.py -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/train.py -------------------------------------------------------------------------------- /train_ldm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/train_ldm.py -------------------------------------------------------------------------------- /train_ldm_discrete.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/train_ldm_discrete.py -------------------------------------------------------------------------------- /train_t2i_discrete.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/train_t2i_discrete.py -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/utils.py -------------------------------------------------------------------------------- /uvit_datasets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/tyshiwo1/DiM-DiffusionMamba/HEAD/uvit_datasets.py --------------------------------------------------------------------------------