WIP

README.md

@@ -3,7 +3,7 @@
 </p>
 
 # MONAI Generative Models
-Prototyping repo for generative models to be integrated into MONAI core.
+Prototyping repository for generative models to be integrated into MONAI core.
 ## Features
 * Network architectures: Diffusion Model, Autoencoder-KL, VQ-VAE, (Multi-scale) Patch-GAN discriminator.
 * Diffusion Model Schedulers: DDPM, DDIM, and PNDM.

generative/inferers/inferer.py

@@ -347,14 +347,14 @@ def sample(
             latent = outputs
 
         with torch.no_grad():
-            image = autoencoder_model.decode_stage_2_outputs(latent) * self.scale_factor
+            image = autoencoder_model.decode_stage_2_outputs(latent / self.scale_factor)
 
         if save_intermediates:
             intermediates = []
             for latent_intermediate in latent_intermediates:
                 with torch.no_grad():
                     intermediates.append(
-                        autoencoder_model.decode_stage_2_outputs(latent_intermediate) * self.scale_factor
+                        autoencoder_model.decode_stage_2_outputs(latent_intermediate / self.scale_factor)
                     )
             return image, intermediates
 

generative/losses/perceptual.py

@@ -45,8 +45,11 @@ def __init__(
         if spatial_dims not in [2, 3]:
             raise NotImplementedError("Perceptual loss is implemented only in 2D and 3D.")
 
-        if spatial_dims == 2 and "medicalnet_" in network_type:
-            raise ValueError("MedicalNet networks are only compatible with ``spatial_dims=3``.")
+        if (spatial_dims == 2 or is_fake_3d) and "medicalnet_" in network_type:
+            raise ValueError(
+                "MedicalNet networks are only compatible with ``spatial_dims=3``."
+                "Argument is_fake_3d must be set to False."
+            )
 
         self.spatial_dims = spatial_dims
         if spatial_dims == 3 and is_fake_3d is False:

generative/metrics/ms_ssim.py

@@ -80,10 +80,7 @@ def _compute_metric(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         """
 
         if not x.shape == y.shape:
-            raise ValueError(
-                f"Input images should have the same dimensions, \
-             but got {x.shape} and {y.shape}."
-            )
+            raise ValueError(f"Input images should have the same dimensions, but got {x.shape} and {y.shape}.")
 
         for d in range(len(x.shape) - 1, 1, -1):
             x = x.squeeze(dim=d)
@@ -94,10 +91,7 @@ def _compute_metric(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         elif len(x.shape) == 5:
             avg_pool = F.avg_pool3d
         else:
-            raise ValueError(
-                f"Input images should be 4-d or 5-d tensors, but \
-                got {x.shape}"
-            )
+            raise ValueError(f"Input images should be 4-d or 5-d tensors, but got {x.shape}")
 
         if self.weights is None:
             # as per Ref 1 - Sec 3.2.
@@ -109,14 +103,14 @@ def _compute_metric(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
         for idx, shape_size in enumerate(x.shape[2:]):
             if shape_size % divisible_by != 0:
                 raise ValueError(
-                    f"Image size needs to be divisible by {divisible_by} but \
-                        dimension {idx + 2} has size {shape_size}"
+                    f"Image size needs to be divisible by {divisible_by} but "
+                    f"dimension {idx + 2} has size {shape_size}"
                 )
 
             if shape_size < bigger_than:
                 raise ValueError(
-                    f"Image size should be larger than {bigger_than} due to \
-                        the {len(self.weights) - 1} downsamplings in MS-SSIM."
+                    f"Image size should be larger than {bigger_than} due to "
+                    f"the {len(self.weights) - 1} downsamplings in MS-SSIM."
                 )
 
         levels = self.weights.shape[0]

generative/networks/nets/__init__.py

@@ -12,4 +12,5 @@
 from .autoencoderkl import AutoencoderKL
 from .diffusion_model_unet import DiffusionModelUNet
 from .patchgan_discriminator import MultiScalePatchDiscriminator, PatchDiscriminator
+from .transformer import DecoderOnlyTransformer
 from .vqvae import VQVAE