End of training

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+val_imgs_grid.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,52 @@
+
+---
+license: creativeml-openrail-m
+base_model: yurman/uncond_sd2-base
+tags:
+- stable-diffusion
+- stable-diffusion-diffusers
+- diffusers
+inference: true
+---
+
+# Unconditioned stable diffusion finetuning - yurman/uncond-sd2-base-complex-4
+
+This pipeline was finetuned from **yurman/uncond_sd2-base**
+for brain image generation.
+Below are some example images generated with the finetuned pipeline: 
+
+![val_imgs_grid](./val_imgs_grid.png)
+
+
+## Pipeline usage
+
+You can use the pipeline like so:
+
+```python
+from diffusers import StableDiffusionUnconditionalPipeline
+import torch
+
+pipeline = StableDiffusionUnconditionalPipeline.from_pretrained("yurman/uncond-sd2-base-complex-4", torch_dtype=torch.float32)
+image = pipeline(1).images[0]
+image.save("brain_image.png")
+```
+
+## Training info
+These are the key hyperparameters used during training:
+
+* Epochs: 400
+* Max Train Steps: 100000
+* Learning rate: 5e-05
+* Batch size: 18
+* VAE scaling: 0.12
+* VAE type: MEDVAE
+* Input perturbation: 0.0
+* Noise offset: 0.0
+* Gradient accumulation steps: 3
+* Image resolution: 256
+* Mixed-precision: no
+* Max rotation degree: 10
+* Prediction Type: v_prediction
+* SNR Gamma: 5.0
+
+More information on all the CLI arguments and the environment are available on your [`wandb` run page](https://wandb.ai/mri-diffusion/uncond-sd2-base-complex/runs/fpsgaddb).

checkpoint-100000/optimizer.bin

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

checkpoint-100000/random_states_0.pkl

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

checkpoint-100000/scheduler.bin

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

checkpoint-100000/unet/config.json

@@ -0,0 +1,68 @@
+{
+  "_class_name": "UNet2DModel",
+  "_diffusers_version": "0.26.0.dev0",
+  "_name_or_path": "yurman/uncond_sd2-base",
+  "act_fn": "silu",
+  "add_mid_block_attention": true,
+  "attention_head_dim": [
+    5,
+    10,
+    20,
+    20
+  ],
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "down_block_types": [
+    "SelfAttnDownBlock2D",
+    "SelfAttnDownBlock2D",
+    "SelfAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "downsample_type": "conv",
+  "dropout": 0.0,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DSelfAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "num_train_timesteps": null,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "SelfAttnUpBlock2D",
+    "SelfAttnUpBlock2D",
+    "SelfAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "upsample_type": "conv",
+  "use_linear_projection": true,
+  "use_transformer_attentions": true
+}

checkpoint-100000/unet/diffusion_pytorch_model.safetensors

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

checkpoint-100000/unet_ema/config.json

@@ -0,0 +1,75 @@
+{
+  "_class_name": "UNet2DModel",
+  "_diffusers_version": "0.26.0.dev0",
+  "_name_or_path": "yurman/uncond_sd2-base",
+  "act_fn": "silu",
+  "add_mid_block_attention": true,
+  "attention_head_dim": [
+    5,
+    10,
+    20,
+    20
+  ],
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "decay": 0.9999,
+  "down_block_types": [
+    "SelfAttnDownBlock2D",
+    "SelfAttnDownBlock2D",
+    "SelfAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "downsample_type": "conv",
+  "dropout": 0.0,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "inv_gamma": 1.0,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DSelfAttn",
+  "min_decay": 0.0,
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "num_train_timesteps": null,
+  "optimization_step": 100000,
+  "out_channels": 4,
+  "power": 0.6666666666666666,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "SelfAttnUpBlock2D",
+    "SelfAttnUpBlock2D",
+    "SelfAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "update_after_step": 0,
+  "upsample_type": "conv",
+  "use_ema_warmup": false,
+  "use_linear_projection": true,
+  "use_transformer_attentions": true
+}

checkpoint-100000/unet_ema/diffusion_pytorch_model.safetensors

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

model_index.json

@@ -0,0 +1,18 @@
+{
+  "_class_name": "StableDiffusionUnconditionalPipeline",
+  "_diffusers_version": "0.26.0.dev0",
+  "_name_or_path": "yurman/uncond_sd2-base",
+  "requires_safety_checker": false,
+  "scheduler": [
+    "diffusers",
+    "DDPMScheduler"
+  ],
+  "unet": [
+    "diffusers",
+    "UNet2DModel"
+  ],
+  "vae": [
+    "pips.model.med_vae_wrapper",
+    "MedVAEWrapper"
+  ]
+}

scheduler/scheduler_config.json

@@ -0,0 +1,21 @@
+{
+  "_class_name": "DDPMScheduler",
+  "_diffusers_version": "0.26.0.dev0",
+  "beta_end": 0.012,
+  "beta_schedule": "scaled_linear",
+  "beta_start": 0.00085,
+  "clip_sample": false,
+  "clip_sample_range": 1.0,
+  "dynamic_thresholding_ratio": 0.995,
+  "num_train_timesteps": 1000,
+  "prediction_type": "v_prediction",
+  "rescale_betas_zero_snr": false,
+  "sample_max_value": 1.0,
+  "set_alpha_to_one": false,
+  "skip_prk_steps": true,
+  "steps_offset": 1,
+  "thresholding": false,
+  "timestep_spacing": "leading",
+  "trained_betas": null,
+  "variance_type": "fixed_small"
+}

unet/config.json

@@ -0,0 +1,68 @@
+{
+  "_class_name": "UNet2DModel",
+  "_diffusers_version": "0.26.0.dev0",
+  "_name_or_path": "yurman/uncond_sd2-base",
+  "act_fn": "silu",
+  "add_mid_block_attention": true,
+  "attention_head_dim": [
+    5,
+    10,
+    20,
+    20
+  ],
+  "attention_type": "default",
+  "block_out_channels": [
+    320,
+    640,
+    1280,
+    1280
+  ],
+  "center_input_sample": false,
+  "class_embed_type": null,
+  "class_embeddings_concat": false,
+  "conv_in_kernel": 3,
+  "conv_out_kernel": 3,
+  "down_block_types": [
+    "SelfAttnDownBlock2D",
+    "SelfAttnDownBlock2D",
+    "SelfAttnDownBlock2D",
+    "DownBlock2D"
+  ],
+  "downsample_padding": 1,
+  "downsample_type": "conv",
+  "dropout": 0.0,
+  "flip_sin_to_cos": true,
+  "freq_shift": 0,
+  "in_channels": 4,
+  "layers_per_block": 2,
+  "mid_block_scale_factor": 1,
+  "mid_block_type": "UNetMidBlock2DSelfAttn",
+  "norm_eps": 1e-05,
+  "norm_num_groups": 32,
+  "num_attention_heads": null,
+  "num_class_embeds": null,
+  "num_train_timesteps": null,
+  "out_channels": 4,
+  "projection_class_embeddings_input_dim": null,
+  "resnet_out_scale_factor": 1.0,
+  "resnet_skip_time_act": false,
+  "resnet_time_scale_shift": "default",
+  "reverse_transformer_layers_per_block": null,
+  "sample_size": 64,
+  "time_cond_proj_dim": null,
+  "time_embedding_act_fn": null,
+  "time_embedding_dim": null,
+  "time_embedding_type": "positional",
+  "timestep_post_act": null,
+  "transformer_layers_per_block": 1,
+  "up_block_types": [
+    "UpBlock2D",
+    "SelfAttnUpBlock2D",
+    "SelfAttnUpBlock2D",
+    "SelfAttnUpBlock2D"
+  ],
+  "upcast_attention": false,
+  "upsample_type": "conv",
+  "use_linear_projection": true,
+  "use_transformer_attentions": true
+}

unet/diffusion_pytorch_model.safetensors

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

vae/config.json

@@ -0,0 +1,12 @@
+{
+  "_class_name": "MedVAEWrapper",
+  "_diffusers_version": "0.26.0.dev0",
+  "block_out_channels": [
+    1,
+    1,
+    1
+  ],
+  "downsampling_factor": 4,
+  "in_channels": 2,
+  "scaling_factor": 0.12
+}

vae/diffusion_pytorch_model.safetensors

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

vae/pips.model.med_vae_wrapper.py

@@ -0,0 +1,86 @@
+"""
+Wrapper to match VAE interface to that of SD VAE.
+"""
+
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from medvae.models import AutoencoderKL_2D
+from medvae.utils.factory import (
+    FILE_DICT_ASSOCIATIONS,
+    create_model,
+    download_model_weights,
+)
+from omegaconf import OmegaConf
+
+
+class LatentDist:
+    def __init__(self, dist):
+        self.latent_dist = dist
+
+    def sample(self):
+        return self.latent_dist.mode()
+
+    # alias
+    def mode(self):
+        return self.latent_dist.mode()
+
+
+class MedVAEWrapper(ModelMixin, ConfigMixin):
+    config_name = "config.json"
+    ignore_for_config = ["vae"]
+
+    @register_to_config
+    def __init__(self, vae=None, scaling_factor=1.0, downsampling_factor=4):
+        super().__init__()
+        assert downsampling_factor in [
+            4,
+            8,
+        ], "Only 4x and 8x downsampling are currently supported"
+        if vae is None:
+            model_name = (
+                "medvae_4_4_2d_c" if downsampling_factor == 4 else "medvae_8_4_2d_c"
+            )
+            config_fpath = download_model_weights(
+                FILE_DICT_ASSOCIATIONS[model_name]["config"]
+            )
+            if model_name == "medvae_8_4_2d_c":
+                config_fpath = "/data/yurman/repos/fast-mri-ldm/submodules/medvae/configs/ours-8x1-new.yaml"
+
+            conf = OmegaConf.load(config_fpath)
+            conf.embed_dim = 4
+            conf.ddconfig.z_channels = 4
+            conf["ddconfig"]["in_channels"] = 2
+            conf["ddconfig"]["out_ch"] = 2
+
+            vae = AutoencoderKL_2D(
+                ddconfig=conf.ddconfig,
+                embed_dim=conf.embed_dim,
+            )
+
+        self.vae = vae
+        # When using SD pipeline it uses `block_out_channels` to determine the size of the image based on
+        # 2 ** (len(block_out_channels) - 1)
+        n_blocks = int(np.log2(downsampling_factor)) + 1
+        self.register_to_config(
+            block_out_channels=[
+                1,
+            ]
+            * n_blocks,
+            in_channels=2,
+            scaling_factor=scaling_factor,
+            downsampling_factor=downsampling_factor,
+        )
+
+    def encode(self, x):
+        dist = self.vae.encode(x)
+
+        return SimpleNamespace(latent_dist=LatentDist(dist))
+
+    def decode(self, x, return_dict=False, generator=None):
+        with torch.amp.autocast(device_type="cuda", enabled=False):
+            x = self.vae.decode(x)
+        return (x,)