model_weights and readme

README.md

@@ -1,3 +1,42 @@
 ---
 license: mit
 ---
+
+# Med-VAE
+
+Med-VAE is a family of six large-scale, generalizable 2D and 3D variational autoencoders (VAEs) designed for medical imaging. It is trained on over one million medical images across multiple anatomical regions and modalities. Med-VAE autoencoders encode medical images as downsized latent representations and decode latent representations back to high-resolution images. Across diverse tasks obtained from 20 medical image datasets, we demonstrate that utilizing MedVAE latent representations in place of high-resolution images when training downstream models can lead to efficiency benefits (up to 70x improvement in throughput) while simultaneously preserving clinically-relevant features. 
+
+[💻 Github](https://github.com/StanfordMIMI/MedVAE)
+
+## Model Description
+| Total Compression Factor | Channels | Dimensions | Modalities | Anatomies | Config File | Model File |
+|----------|----------|----------|----------|----------|----------|----------|
+| 16 | 1 | 2D | X-ray | Chest, Breast (FFDM) | [medvae_4x1.yaml ](model_weights/medvae_4x1.yaml)| [vae_4x_1c_2D.ckpt](model_weights/vae_4x_1c_2D.ckpt)
+| 16 | 3 | 2D | X-ray | Chest, Breast (FFDM) | [medvae_4x3.yaml](model_weights/medvae_4x3.yaml) | [vae_4x_3c_2D.ckpt](model_weights/vae_4x_3c_2D.ckpt)
+| 64 | 1 | 2D | X-ray | Chest, Breast (FFDM) | [medvae_8x1.yaml](model_weights/medvae_8x1.yaml) | [vae_8x_1c_2D.ckpt](model_weights/vae_8x_1c_2D.ckpt)
+| 64 | 3 | 2D | X-ray | Chest, Breast (FFDM) | [medvae_8x4.yaml](model_weights/medvae_8x4.yaml) | [vae_8x_4c_2D.ckpt](model_weights/vae_8x_4c_2D.ckpt)
+| 64 | 1 | 3D | MRI, CT | Whole-Body | [medvae_4x1.yaml ](model_weights/medvae_4x1.yaml) | [vae_4x_1c_3D.ckpt](model_weights/vae_4x_1c_3D.ckpt)
+| 512 | 1 | 3D | MRI, CT | Whole-Body | [medvae_8x1.yaml](model_weights/medvae_8x1.yaml) | [vae_8x_1c_3D.ckpt](model_weights/vae_8x_1c_3D.ckpt)
+
+Note: Model weights and checkpoints are located in the `model_weights` folder.
+
+## Usage Instructions
+
+
+
+## Citation
+If you use Med-VAE, please cite the original paper:
+
+```bibtex
+@article{varma2025medvae,
+  title = {Med-VAE: --},
+  author = {Maya Varma, Ashwin Kumar, Rogier van der Sluijs, Sophie Ostmeier, Louis Blankemeier, Pierre Chambon, Christian Bluethgen, Jip Prince, Curtis Langlotz, Akshay Chaudhari},
+  year = {2025},
+  publisher = {Arxiv},
+  journal = {Arvix},
+  howpublished = {TODO}
+}
+```
+
+For questions, please place a Github Issues message.
+

model_weights/medvae_4x1.yaml

@@ -0,0 +1,13 @@
+embed_dim: 1
+
+ddconfig:
+  double_z: True
+  z_channels: 1
+  resolution: 512
+  in_channels: 1
+  out_ch: 1
+  ch: 128
+  ch_mult: [1,2,4]
+  num_res_blocks: 2
+  attn_resolutions: []
+  dropout: 0.0

model_weights/medvae_4x3.yaml

@@ -0,0 +1,54 @@
+model:
+  base_learning_rate: 4.5e-6
+  target: ldm.models.autoencoder.AutoencoderKL
+  params:
+    monitor: "val/rec_loss"
+    embed_dim: 3
+    lossconfig:
+      target: ldm.modules.losses.LPIPSWithDiscriminator
+      params:
+        disc_start: 50001
+        kl_weight: 0.000001
+        disc_weight: 0.5
+
+    ddconfig:
+      double_z: True
+      z_channels: 3
+      resolution: 256
+      in_channels: 3
+      out_ch: 3
+      ch: 128
+      ch_mult: [ 1,2,4 ]  # num_down = len(ch_mult)-1
+      num_res_blocks: 2
+      attn_resolutions: [ ]
+      dropout: 0.0
+
+
+data:
+  target: main.DataModuleFromConfig
+  params:
+    batch_size: 12
+    wrap: True
+    train:
+      target: ldm.data.imagenet.ImageNetSRTrain
+      params:
+        size: 256
+        degradation: pil_nearest
+    validation:
+      target: ldm.data.imagenet.ImageNetSRValidation
+      params:
+        size: 256
+        degradation: pil_nearest
+
+lightning:
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        batch_frequency: 1000
+        max_images: 8
+        increase_log_steps: True
+
+  trainer:
+    benchmark: True
+    accumulate_grad_batches: 2

model_weights/medvae_8x1.yaml

@@ -0,0 +1,13 @@
+embed_dim: 1
+
+ddconfig:
+  double_z: True
+  z_channels: 1
+  resolution: 512
+  in_channels: 1
+  out_ch: 1
+  ch: 128
+  ch_mult: [1,2,4,4]
+  num_res_blocks: 2
+  attn_resolutions: []
+  dropout: 0.0

model_weights/medvae_8x4.yaml

@@ -0,0 +1,53 @@
+model:
+  base_learning_rate: 4.5e-6
+  target: ldm.models.autoencoder.AutoencoderKL
+  params:
+    monitor: "val/rec_loss"
+    embed_dim: 4
+    lossconfig:
+      target: ldm.modules.losses.LPIPSWithDiscriminator
+      params:
+        disc_start: 50001
+        kl_weight: 0.000001
+        disc_weight: 0.5
+
+    ddconfig:
+      double_z: True
+      z_channels: 4
+      resolution: 256
+      in_channels: 3
+      out_ch: 3
+      ch: 128
+      ch_mult: [ 1,2,4,4 ]  # num_down = len(ch_mult)-1
+      num_res_blocks: 2
+      attn_resolutions: [ ]
+      dropout: 0.0
+
+data:
+  target: main.DataModuleFromConfig
+  params:
+    batch_size: 12
+    wrap: True
+    train:
+      target: ldm.data.imagenet.ImageNetSRTrain
+      params:
+        size: 256
+        degradation: pil_nearest
+    validation:
+      target: ldm.data.imagenet.ImageNetSRValidation
+      params:
+        size: 256
+        degradation: pil_nearest
+
+lightning:
+  callbacks:
+    image_logger:
+      target: main.ImageLogger
+      params:
+        batch_frequency: 1000
+        max_images: 8
+        increase_log_steps: True
+
+  trainer:
+    benchmark: True
+    accumulate_grad_batches: 2

model_weights/vae_4x_1c_2D.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e4bd8931238a6c52acb3d826025f5bcfa284aa9f98cc5455505f131d690c1ba
+size 221345538

model_weights/vae_4x_1c_3D.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c780a850d5ed20303c37a84be6022c8ebca2236509d673e3059bfcec75ce383a
+size 644085658

model_weights/vae_4x_3c_2D.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:651dfc0792fc61d004cb795b44b52dcb3c3523321776f0c84b6890561c2e5778
+size 223784534

model_weights/vae_8x_1c_2D.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4d99374ce8a6fcaf4a5b0117dd4396b520e273e5c92660772a710421337ddc52
+size 334673798

model_weights/vae_8x_1c_3D.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5416db01e6557316630a68c3367403c86100fc9726886fa2ab24595a89f0a98d
+size 983906794

model_weights/vae_8x_4c_2D.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9457e018edd267cfa526878730f5c04699f9c6a9a646a8d51d89aa9d7ccb99f8
+size 337998486