Upload 4 files

hf_diffusion_service.py

@@ -0,0 +1,311 @@
+import os
+import sys
+import torch
+import numpy as np
+import torchvision.transforms as transforms
+from PIL import Image
+
+# Add the hf_model_files directory to the path
+sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'hf_model_files'))
+
+from model import UNet, marginal_prob_std, diffusion_coeff, Euler_Maruyama_sampler
+
+class CompatibleUNet(UNet):
+    """A UNet model that's compatible with the saved weights."""
+    
+    def __init__(self, marginal_prob_std, channels=[32, 64, 128, 256, 512], embed_dim=256,
+                 embed_dim_mask=256, input_dim_mask=1*256*256):  # Changed to 1*256*256
+        # Override the parent's __init__ to set the correct input channels
+        super().__init__(marginal_prob_std, channels, embed_dim, embed_dim_mask, input_dim_mask)
+        
+        # Replace the first conv layer to accept 1 input channel instead of 4
+        self.conv1 = torch.nn.Conv2d(1, channels[0], 3, stride=2, bias=False, padding=1)
+        
+        # Also need to fix the output layer if it exists
+        if hasattr(self, 'tconv0'):
+            self.tconv0 = torch.nn.ConvTranspose2d(channels[0], 1, 3, stride=1, padding=1, output_padding=0)
+
+class HFDiffusionService:
+    """Service class for the Hugging Face conditional diffusion model."""
+    
+    def __init__(self):
+        # Check if CUDA is available and print status
+        cuda_available = torch.cuda.is_available()
+        print(f"CUDA available for HF diffusion: {cuda_available}")
+        if not cuda_available:
+            print("Warning: CUDA is not available for HF diffusion. Using CPU instead. This might be slower.")
+            
+        self.device = torch.device('cuda:0' if cuda_available else 'cpu')
+        self.Lambda = 25.0
+        
+        # Initialize the model functions
+        self.marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=self.Lambda, device=self.device)
+        self.diffusion_coeff_fn = lambda t: diffusion_coeff(t, Lambda=self.Lambda, device=self.device)
+        
+        # Model path for the downloaded Hugging Face model
+        self.model_path = os.path.join("hf_model_files", "pytorch_model.bin")
+        
+        try:
+            # Load the state dict first to understand the architecture
+            state_dict = torch.load(self.model_path, map_location=self.device)
+            
+            # Analyze the state dict to determine the correct architecture
+            conv1_weight = state_dict.get('conv1.weight', None)
+            cond_embed_weight = state_dict.get('cond_embed.1.weight', None)
+            
+            if conv1_weight is not None:
+                actual_input_channels = conv1_weight.shape[1]
+                print(f"Detected input channels from state dict: {actual_input_channels}")
+                
+                if cond_embed_weight is not None:
+                    actual_input_dim_mask = cond_embed_weight.shape[1]
+                    print(f"Detected input_dim_mask from state dict: {actual_input_dim_mask}")
+                    
+                    # Create a compatible model
+                    if actual_input_channels == 1 and actual_input_dim_mask == 65536:
+                        # The saved model expects 1 input channel and 65536 flattened input
+                        # This suggests it was trained with 1*256*256 = 65536
+                        self.score_model = CompatibleUNet(
+                            marginal_prob_std=self.marginal_prob_std_fn,
+                            input_dim_mask=65536
+                        )
+                        self.input_channels = 1
+                        self.input_dim_mask = 65536
+                    else:
+                        # Use the original architecture
+                        self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+                        self.input_channels = 4
+                        self.input_dim_mask = 262144
+                else:
+                    # Fallback to original
+                    self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+                    self.input_channels = 4
+                    self.input_dim_mask = 262144
+            else:
+                # Fallback to original
+                self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+                self.input_channels = 4
+                self.input_dim_mask = 262144
+            
+            # Load the weights
+            self.score_model.load_state_dict(state_dict)
+            self.score_model.to(self.device)
+            self.score_model.eval()
+            
+            print(f"HF Diffusion model loaded successfully from {self.model_path}")
+            print(f"Model configured for {self.input_channels} input channels and {self.input_dim_mask} mask dimensions")
+            
+        except Exception as e:
+            print(f"Error loading HF diffusion model: {e}")
+            raise e
+
+    def generate_image(self, mask):
+        """
+        Generate a medical image based on a conditioning mask.
+        
+        Args:
+            mask: Conditioning mask tensor of shape (1, 4, 256, 256) or PIL Image
+            
+        Returns:
+            Generated image as PIL Image
+        """
+        try:
+            # Process the mask input
+            processed_mask = self.process_mask(mask)
+            
+            # Generate the image
+            generated_tensor = self.generate_from_mask(processed_mask)
+            
+            # Convert tensor to PIL Image
+            return self.tensor_to_image(generated_tensor)
+            
+        except Exception as e:
+            print(f"Error generating HF diffusion image: {e}")
+            return None
+
+    def process_mask(self, mask):
+        """
+        Process the input mask to the correct format for the model.
+        
+        Args:
+            mask: Input mask (PIL Image, numpy array, or tensor)
+            
+        Returns:
+            Processed mask tensor of shape (1, 1, 256, 256) for 1-channel model
+        """
+        try:
+            # If mask is a PIL Image, convert to tensor
+            if isinstance(mask, Image.Image):
+                transform = transforms.Compose([
+                    transforms.Grayscale(num_output_channels=1),
+                    transforms.Resize((256, 256), antialias=True),
+                    transforms.ToTensor()
+                ])
+                tensor = transform(mask).unsqueeze(0)  # Add batch dimension
+            elif isinstance(mask, np.ndarray):
+                # Convert numpy array to tensor
+                if mask.ndim == 2:
+                    mask = mask[np.newaxis, :, :]  # Add channel dimension
+                tensor = torch.from_numpy(mask).float()
+                if tensor.dim() == 3:
+                    tensor = tensor.unsqueeze(0)  # Add batch dimension
+            elif isinstance(mask, torch.Tensor):
+                tensor = mask
+                if tensor.dim() == 3:
+                    tensor = tensor.unsqueeze(0)  # Add batch dimension
+            else:
+                raise ValueError(f"Unsupported mask type: {type(mask)}")
+            
+            # Ensure the tensor has the correct shape based on model input
+            if self.input_channels == 1:
+                # Model expects 1 channel
+                if tensor.shape[1] != 1:
+                    # Take the first channel or average if multiple channels
+                    if tensor.shape[1] > 1:
+                        tensor = tensor.mean(dim=1, keepdim=True)
+                    else:
+                        tensor = tensor[:, :1, :, :]
+            else:
+                # Model expects 4 channels
+                if tensor.shape[1] == 1:
+                    # If single channel, repeat to 4 channels
+                    tensor = tensor.repeat(1, 4, 1, 1)
+                elif tensor.shape[1] != 4:
+                    raise ValueError(f"Expected 1 or 4 channels, got {tensor.shape[1]}")
+            
+            # Ensure correct size
+            if tensor.shape[2] != 256 or tensor.shape[3] != 256:
+                tensor = torch.nn.functional.interpolate(tensor, size=(256, 256), mode='bilinear', align_corners=False)
+            
+            print(f"Processed mask shape: {tensor.shape}")
+            return tensor.to(self.device)
+            
+        except Exception as e:
+            print(f"Error processing mask: {e}")
+            raise e
+
+    def generate_from_mask(self, conditioning_mask, num_steps=250, eps=1e-3):
+        """
+        Generate image from conditioning mask using the diffusion model.
+        
+        Args:
+            conditioning_mask: Conditioning mask tensor
+            num_steps: Number of sampling steps
+            eps: Smallest time step for numerical stability
+            
+        Returns:
+            Generated image tensor
+        """
+        try:
+            # Determine the output shape based on the model
+            if self.input_channels == 1:
+                x_shape = (1, 256, 256)
+            else:
+                x_shape = (4, 256, 256)
+            
+            with torch.no_grad():
+                samples = Euler_Maruyama_sampler(
+                    self.score_model,
+                    self.marginal_prob_std_fn,
+                    self.diffusion_coeff_fn,
+                    batch_size=1,
+                    x_shape=x_shape,
+                    num_steps=num_steps,
+                    device=self.device,
+                    eps=eps,
+                    y=conditioning_mask
+                )
+            
+            # Clamp values to [0, 1] range
+            return samples.clamp(0, 1)
+            
+        except Exception as e:
+            print(f"Error in generate_from_mask: {e}")
+            raise e
+
+    def tensor_to_image(self, tensor):
+        """
+        Convert tensor to PIL Image.
+        
+        Args:
+            tensor: Generated tensor
+            
+        Returns:
+            PIL Image
+        """
+        try:
+            # Take the first channel for visualization (or average all channels)
+            if tensor.shape[1] > 1:
+                # Average the channels
+                image_tensor = tensor.squeeze(0).mean(dim=0)
+            else:
+                image_tensor = tensor.squeeze(0).squeeze(0)
+            
+            # Convert to numpy and scale to 0-255
+            image_array = (image_tensor.cpu().numpy() * 255).astype(np.uint8)
+            
+            # Create PIL Image
+            image = Image.fromarray(image_array, mode='L')
+            
+            return image
+            
+        except Exception as e:
+            print(f"Error converting tensor to image: {e}")
+            raise e
+
+    def generate_batch(self, masks, num_steps=250, eps=1e-3):
+        """
+        Generate multiple images from a batch of masks.
+        
+        Args:
+            masks: List of masks or batch tensor
+            num_steps: Number of sampling steps
+            eps: Smallest time step for numerical stability
+            
+        Returns:
+            List of generated PIL Images
+        """
+        try:
+            if isinstance(masks, list):
+                # Process each mask individually
+                results = []
+                for mask in masks:
+                    result = self.generate_image(mask)
+                    results.append(result)
+                return results
+            else:
+                # Process as batch
+                processed_masks = self.process_mask(masks)
+                batch_size = processed_masks.shape[0]
+                
+                # Determine the output shape based on the model
+                if self.input_channels == 1:
+                    x_shape = (1, 256, 256)
+                else:
+                    x_shape = (4, 256, 256)
+                
+                with torch.no_grad():
+                    samples = Euler_Maruyama_sampler(
+                        self.score_model,
+                        self.marginal_prob_std_fn,
+                        self.diffusion_coeff_fn,
+                        batch_size=batch_size,
+                        x_shape=x_shape,
+                        num_steps=num_steps,
+                        device=self.device,
+                        eps=eps,
+                        y=processed_masks
+                    )
+                
+                # Convert each sample to image
+                results = []
+                for i in range(batch_size):
+                    sample = samples[i:i+1]
+                    image = self.tensor_to_image(sample)
+                    results.append(image)
+                
+                return results
+                
+        except Exception as e:
+            print(f"Error in generate_batch: {e}")
+            raise e 

inference.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""
+Inference script for the conditional diffusion model.
+This script provides easy-to-use functions for generating medical images.
+"""
+
+import torch
+import numpy as np
+import matplotlib.pyplot as plt
+from model import UNet, marginal_prob_std, diffusion_coeff, Euler_Maruyama_sampler
+
+
+class ConditionalDiffusionInference:
+    """Wrapper class for easy inference with the conditional diffusion model."""
+    
+    def __init__(self, model_path, device='cuda'):
+        """
+        Initialize the inference model.
+        
+        Args:
+            model_path: Path to the trained model checkpoint
+            device: Device to run inference on ('cuda' or 'cpu')
+        """
+        self.device = device
+        self.Lambda = 25.0
+        
+        # Initialize the model
+        self.marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=self.Lambda, device=self.device)
+        self.diffusion_coeff_fn = lambda t: diffusion_coeff(t, Lambda=self.Lambda, device=self.device)
+        
+        self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+        self.score_model.load_state_dict(torch.load(model_path, map_location=self.device))
+        self.score_model.to(self.device)
+        self.score_model.eval()
+        
+        print(f"Model loaded successfully on {self.device}")
+    
+    def generate_image(self, conditioning_mask, num_steps=250, eps=1e-3):
+        """
+        Generate a medical image based on a conditioning mask.
+        
+        Args:
+            conditioning_mask: Conditioning mask tensor of shape (1, 4, 256, 256)
+            num_steps: Number of sampling steps
+            eps: Smallest time step for numerical stability
+            
+        Returns:
+            Generated image tensor of shape (1, 4, 256, 256)
+        """
+        if not isinstance(conditioning_mask, torch.Tensor):
+            conditioning_mask = torch.tensor(conditioning_mask, dtype=torch.float32)
+        
+        if conditioning_mask.dim() == 3:
+            conditioning_mask = conditioning_mask.unsqueeze(0)
+        
+        conditioning_mask = conditioning_mask.to(self.device)
+        
+        with torch.no_grad():
+            samples = Euler_Maruyama_sampler(
+                self.score_model,
+                self.marginal_prob_std_fn,
+                self.diffusion_coeff_fn,
+                batch_size=1,
+                x_shape=(4, 256, 256),
+                num_steps=num_steps,
+                device=self.device,
+                eps=eps,
+                y=conditioning_mask
+            )
+        
+        return samples.clamp(0, 1)
+    
+    def generate_batch(self, conditioning_masks, num_steps=250, eps=1e-3):
+        """
+        Generate multiple images based on conditioning masks.
+        
+        Args:
+            conditioning_masks: Conditioning mask tensor of shape (B, 4, 256, 256)
+            num_steps: Number of sampling steps
+            eps: Smallest time step for numerical stability
+            
+        Returns:
+            Generated images tensor of shape (B, 4, 256, 256)
+        """
+        if not isinstance(conditioning_masks, torch.Tensor):
+            conditioning_masks = torch.tensor(conditioning_masks, dtype=torch.float32)
+        
+        if conditioning_masks.dim() == 3:
+            conditioning_masks = conditioning_masks.unsqueeze(0)
+        
+        conditioning_masks = conditioning_masks.to(self.device)
+        batch_size = conditioning_masks.shape[0]
+        
+        with torch.no_grad():
+            samples = Euler_Maruyama_sampler(
+                self.score_model,
+                self.marginal_prob_std_fn,
+                self.diffusion_coeff_fn,
+                batch_size=batch_size,
+                x_shape=(4, 256, 256),
+                num_steps=num_steps,
+                device=self.device,
+                eps=eps,
+                y=conditioning_masks
+            )
+        
+        return samples.clamp(0, 1)
+    
+    def visualize_generation(self, conditioning_mask, generated_image, save_path=None):
+        """
+        Visualize the conditioning mask and generated image.
+        
+        Args:
+            conditioning_mask: Conditioning mask tensor
+            generated_image: Generated image tensor
+            save_path: Optional path to save the visualization
+        """
+        fig, axes = plt.subplots(2, 4, figsize=(16, 8))
+        
+        # Plot conditioning mask
+        for i in range(4):
+            axes[0, i].imshow(conditioning_mask[0, i].cpu().numpy(), cmap='gray')
+            axes[0, i].set_title(f'Conditioning Mask {i+1}')
+            axes[0, i].axis('off')
+        
+        # Plot generated image
+        for i in range(4):
+            axes[1, i].imshow(generated_image[0, i].cpu().numpy(), cmap='gray')
+            axes[1, i].set_title(f'Generated Image {i+1}')
+            axes[1, i].axis('off')
+        
+        plt.tight_layout()
+        
+        if save_path:
+            plt.savefig(save_path, dpi=150, bbox_inches='tight')
+            print(f"Visualization saved to {save_path}")
+        
+        plt.show()
+
+
+def main():
+    """Example usage of the inference model."""
+    
+    # Initialize the model
+    model_path = "ckpt_3D_v2.pth"  # Update with your model path
+    inference_model = ConditionalDiffusionInference(model_path, device='cuda')
+    
+    # Create a random conditioning mask (replace with your actual mask)
+    conditioning_mask = torch.randn(1, 4, 256, 256)
+    
+    # Generate image
+    print("Generating image...")
+    generated_image = inference_model.generate_image(conditioning_mask)
+    
+    # Visualize results
+    inference_model.visualize_generation(conditioning_mask, generated_image, "generation_result.png")
+    
+    print("Generation complete!")
+
+
+if __name__ == "__main__":
+    main() 

model.py

@@ -0,0 +1,334 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from tqdm import tqdm
+
+class GaussianFourierProjection(nn.Module):
+    """Gaussian random features for encoding time steps."""
+    def __init__(self, embed_dim, scale=30.):
+        super().__init__()
+        # Randomly sample weights (frequencies) during initialization.
+        # These weights (frequencies) are fixed during optimization and are not trainable.
+        self.W = nn.Parameter(torch.randn(embed_dim // 2) * scale, requires_grad=False)
+    
+    def forward(self, x):
+        # Cosine(2 pi freq x), Sine(2 pi freq x)
+        x_proj = x[:, None] * self.W[None, :] * 2 * np.pi
+        return torch.cat([torch.sin(x_proj), torch.cos(x_proj)], dim=-1)
+
+
+class Dense(nn.Module):
+    """
+    Maps an embedding vector to a bias/scale tensor that can be broadcast over a
+    2-D feature map (B, C, H, W) – output shape is (B, C, 1, 1).
+    """
+    def __init__(self, input_dim: int, output_dim: int):
+        super().__init__()
+        self.dense = nn.Linear(input_dim, output_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B = x.size(0)
+        x = x.view(B, -1)                 # (B, input_dim)
+        return self.dense(x).view(B, -1, 1, 1)   # (B, C, 1, 1)
+
+
+class UNet(nn.Module):
+    """A time-dependent score-based model built upon U-Net architecture."""
+
+    def __init__(self, marginal_prob_std, channels=[32, 64, 128, 256, 512], embed_dim=256,
+                 embed_dim_mask=256, input_dim_mask=4*256*256):
+        """Initialize a time-dependent score-based network.
+
+        Args:
+            marginal_prob_std: A function that takes time t and gives the standard
+                deviation of the perturbation kernel p_{0t}(x(t) | x(0)).
+            channels: The number of channels for feature maps of each resolution.
+            embed_dim: The dimensionality of Gaussian random feature embeddings.
+        """
+        super().__init__()
+        # Gaussian random feature embedding layer for time
+        self.time_embed = nn.Sequential(
+            GaussianFourierProjection(embed_dim=embed_dim),
+            nn.Linear(embed_dim, embed_dim)
+        )
+        
+        # flatten the mask and apply a linear layer
+        self.cond_embed = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(input_dim_mask, embed_dim_mask)
+        )
+
+        # Encoding layers where the resolution decreases
+        self.conv1 = nn.Conv2d(4, channels[0], 3, stride=2, bias=False, padding=1)
+        self.t_mod1 = Dense(embed_dim, channels[0])
+        self.gnorm1 = nn.GroupNorm(4, num_channels=channels[0])
+
+        self.conv1a = nn.Conv2d(channels[0], channels[0], 3, stride=1, bias=False, padding=1)
+        self.t_mod1a = Dense(embed_dim, channels[0])
+        self.gnorm1a = nn.GroupNorm(4, num_channels=channels[0])
+
+        self.conv2 = nn.Conv2d(channels[0], channels[1], 3, stride=2, bias=False, padding=1)
+        self.t_mod2 = Dense(embed_dim, channels[1])
+        self.y_mod2 = Dense(embed_dim, channels[1])
+        self.gnorm2 = nn.GroupNorm(32, num_channels=channels[1])
+        
+        self.conv2a = nn.Conv2d(channels[1], channels[1], 3, stride=1, bias=False, padding=1)
+        self.t_mod2a = Dense(embed_dim, channels[1])
+        self.y_mod2a = Dense(embed_dim, channels[1])
+        self.gnorm2a = nn.GroupNorm(32, num_channels=channels[1])
+
+        self.conv3 = nn.Conv2d(channels[1], channels[2], 3, stride=2, bias=False, padding=1)
+        self.t_mod3 = Dense(embed_dim, channels[2])
+        self.y_mod3 = Dense(embed_dim, channels[2])
+        self.gnorm3 = nn.GroupNorm(32, num_channels=channels[2])
+        
+        self.conv3a = nn.Conv2d(channels[2], channels[2], 3, stride=1, bias=False, padding=1)
+        self.t_mod3a = Dense(embed_dim, channels[2])
+        self.y_mod3a = Dense(embed_dim, channels[2])
+        self.gnorm3a = nn.GroupNorm(32, num_channels=channels[2])
+
+        self.conv4 = nn.Conv2d(channels[2], channels[3], 3, stride=2, bias=False, padding=1)
+        self.t_mod4 = Dense(embed_dim, channels[3])
+        self.y_mod4 = Dense(embed_dim, channels[3])
+        self.gnorm4 = nn.GroupNorm(32, num_channels=channels[3])
+        
+        self.conv4a = nn.Conv2d(channels[3], channels[3], 3, stride=1, bias=False, padding=1)
+        self.t_mod4a = Dense(embed_dim, channels[3])
+        self.y_mod4a = Dense(embed_dim, channels[3])
+        self.gnorm4a = nn.GroupNorm(32, num_channels=channels[3])
+
+        self.conv5 = nn.Conv2d(channels[3], channels[4], 3, stride=2, bias=False, padding=1)
+        self.t_mod5 = Dense(embed_dim, channels[4])
+        self.y_mod5 = Dense(embed_dim, channels[4])
+        self.gnorm5 = nn.GroupNorm(32, num_channels=channels[4])
+
+        self.conv5a = nn.Conv2d(channels[4], channels[4], 3, stride=1, bias=False, padding=1)
+        self.t_mod5a = Dense(embed_dim, channels[4])
+        self.y_mod5a = Dense(embed_dim, channels[4])
+        self.gnorm5a = nn.GroupNorm(32, num_channels=channels[4])
+        
+        # Decoding layers where the resolution increases
+        self.tconv5b = nn.Conv2d(channels[4], channels[4], 3, stride=1, bias=False, padding=1)
+        self.t_mod6b = Dense(embed_dim, channels[4])
+        self.y_mod6b = Dense(embed_dim, channels[4])
+        self.tgnorm5b = nn.GroupNorm(32, num_channels=channels[4])
+        
+        self.tconv5 = nn.ConvTranspose2d(2*channels[4], channels[3], 3, stride=2, bias=False, padding=1, output_padding=1)
+        self.t_mod6 = Dense(embed_dim, channels[3])
+        self.y_mod6 = Dense(embed_dim, channels[3])
+        self.tgnorm5 = nn.GroupNorm(32, num_channels=channels[3])
+        
+        self.tconv4b = nn.Conv2d(2*channels[3], channels[3], 3, stride=1, bias=False, padding=1)
+        self.t_mod7b = Dense(embed_dim, channels[3])
+        self.y_mod7b = Dense(embed_dim, channels[3])
+        self.tgnorm4b = nn.GroupNorm(32, num_channels=channels[3])
+
+        self.tconv4 = nn.ConvTranspose2d(2*channels[3], channels[2], 3, stride=2, bias=False, padding=1, output_padding=1)
+        self.t_mod7 = Dense(embed_dim, channels[2])
+        self.y_mod7 = Dense(embed_dim, channels[2])
+        self.tgnorm4 = nn.GroupNorm(32, num_channels=channels[2])
+        
+        self.tconv3b = nn.Conv2d(2*channels[2], channels[2], 3, stride=1, bias=False, padding=1)
+        self.t_mod8b = Dense(embed_dim, channels[2])
+        self.y_mod8b = Dense(embed_dim, channels[2])
+        self.tgnorm3b = nn.GroupNorm(32, num_channels=channels[2])
+        
+        self.tconv3 = nn.ConvTranspose2d(2*channels[2], channels[1], 3, stride=2, bias=False, padding=1, output_padding=1)
+        self.t_mod8 = Dense(embed_dim, channels[1])
+        self.y_mod8 = Dense(embed_dim, channels[1])
+        self.tgnorm3 = nn.GroupNorm(32, num_channels=channels[1])
+        
+        self.tconv2b = nn.Conv2d(2*channels[1], channels[1], 3, stride=1, bias=False, padding=1)
+        self.t_mod9b = Dense(embed_dim, channels[1])
+        self.y_mod9b = Dense(embed_dim, channels[1])
+        self.tgnorm2b = nn.GroupNorm(32, num_channels=channels[1])
+
+        self.tconv2 = nn.ConvTranspose2d(2*channels[1], channels[0], 3, stride=2, bias=False, padding=1, output_padding=1)
+        self.t_mod9 = Dense(embed_dim, channels[0])
+        self.y_mod9 = Dense(embed_dim, channels[0])
+        self.tgnorm2 = nn.GroupNorm(32, num_channels=channels[0])
+
+        self.tconv1b = nn.Conv2d(2*channels[0], channels[0], 3, stride=1, bias=False, padding=1)
+        self.t_mod10b = Dense(embed_dim, channels[0])
+        self.y_mod10b = Dense(embed_dim, channels[0])
+        self.tgnorm1b = nn.GroupNorm(32, num_channels=channels[0])
+
+        self.tconv1 = nn.ConvTranspose2d(2*channels[0], channels[0], 3, stride=2, bias=False, padding=1, output_padding=1)
+        self.t_mod10 = Dense(embed_dim, channels[0])
+        self.y_mod10 = Dense(embed_dim, channels[0])
+        self.tgnorm1 = nn.GroupNorm(32, num_channels=channels[0])
+        
+        self.tconv0 = nn.ConvTranspose2d(channels[0], 4, 3, stride=1, padding=1, output_padding=0)
+
+        # The swish activation function
+        self.act = nn.SiLU()
+        # A restricted version of the `marginal_prob_std` function, after specifying a Lambda.
+        self.marginal_prob_std = marginal_prob_std
+
+    def forward(self, x, t, y=None):
+        # Obtain the Gaussian random feature embedding for t
+        embed = self.act(self.time_embed(t))
+        y_embed = self.cond_embed(y)
+
+        # Encoding path, downsampling
+        h1 = self.conv1(x) + self.t_mod1(embed)
+        h1 = self.act(self.gnorm1(h1))
+        
+        h1a = self.conv1a(h1) + self.t_mod1a(embed)
+        h1a = self.act(self.gnorm1a(h1a))
+
+        # 2nd conv
+        h2 = self.conv2(h1a) + self.t_mod2(embed)
+        h2 = h2 * self.y_mod2(y_embed)
+        h2 = self.act(self.gnorm2(h2))
+
+        h2a = self.conv2a(h2) + self.t_mod2a(embed)
+        h2a = h2a * self.y_mod2a(y_embed)
+        h2a = self.act(self.gnorm2a(h2a))
+
+        # 3rd conv
+        h3 = self.conv3(h2a) + self.t_mod3(embed)
+        h3 = h3 * self.y_mod3(y_embed)
+        h3 = self.act(self.gnorm3(h3))
+
+        h3a = self.conv3a(h3) + self.t_mod3a(embed)
+        h3a = h3a * self.y_mod3a(y_embed)
+        h3a = self.act(self.gnorm3a(h3a))
+
+        # 4th conv
+        h4 = self.conv4(h3a) + self.t_mod4(embed)
+        h4 = h4 * self.y_mod4(y_embed)
+        h4 = self.act(self.gnorm4(h4))
+
+        h4a = self.conv4a(h4) + self.t_mod4a(embed)
+        h4a = h4a * self.y_mod4a(y_embed)
+        h4a = self.act(self.gnorm4a(h4a))
+
+        # 5th conv
+        h5 = self.conv5(h4a) + self.t_mod5(embed)
+        h5 = h5 * self.y_mod5(y_embed)
+        h5 = self.act(self.gnorm5(h5))
+
+        h5a = self.conv5a(h5) + self.t_mod5a(embed)
+        h5a = h5a * self.y_mod5a(y_embed)
+        h5a = self.act(self.gnorm5a(h5a))
+
+        # Decoding path up sampling
+        h = self.tconv5b(h5a) + self.t_mod6b(embed)
+        h = h * self.y_mod5(y_embed)
+        h = self.act(self.tgnorm5b(h))
+        
+        # Skip connection from the encoding path
+        h = self.tconv5(torch.cat([h, h5], dim=1)) + self.t_mod6(embed)
+        h = h * self.y_mod6(y_embed)
+        h = self.act(self.tgnorm5(h))
+
+        h = self.tconv4b(torch.cat([h, h4a], dim=1)) + self.t_mod7b(embed)
+        h = h * self.y_mod7b(y_embed)
+        h = self.act(self.tgnorm4b(h))
+
+        h = self.tconv4(torch.cat([h, h4], dim=1)) + self.t_mod7(embed)
+        h = h * self.y_mod7(y_embed)
+        h = self.act(self.tgnorm4(h))
+
+        h = self.tconv3b(torch.cat([h, h3a], dim=1)) + self.t_mod8b(embed)
+        h = h * self.y_mod8b(y_embed)
+        h = self.act(self.tgnorm3b(h))
+
+        h = self.tconv3(torch.cat([h, h3], dim=1)) + self.t_mod8(embed)
+        h = h * self.y_mod8(y_embed)
+        h = self.act(self.tgnorm3(h))
+
+        h = self.tconv2b(torch.cat([h, h2a], dim=1)) + self.t_mod9b(embed)
+        h = h * self.y_mod9b(y_embed)
+        h = self.act(self.tgnorm2b(h))
+
+        h = self.tconv2(torch.cat([h, h2], dim=1)) + self.t_mod9(embed)
+        h = h * self.y_mod9(y_embed)
+        h = self.act(self.tgnorm2(h))
+
+        h = self.tconv1b(torch.cat([h, h1a], dim=1)) + self.t_mod10b(embed)
+        h = h * self.y_mod10b(y_embed)
+        h = self.act(self.tgnorm1b(h))
+
+        h = self.tconv1(torch.cat([h, h1], dim=1)) + self.t_mod10(embed)
+        h = h * self.y_mod10(y_embed)
+        h = self.act(self.tgnorm1(h))
+
+        h = self.tconv0(h)
+        
+        # Normalize output
+        h = h / self.marginal_prob_std(t)[:, None, None, None]
+
+        return h
+
+
+def marginal_prob_std(t, Lambda, device='cpu'):
+    """Compute the standard deviation of $p_{0t}(x(t) | x(0))$.
+
+    Args:
+        t: A vector of time steps.
+        Lambda: The $\lambda$ in our SDE.
+
+    Returns:
+        std : The standard deviation.
+    """
+    t = t.to(device)
+    std = torch.sqrt((Lambda**(2 * t) - 1.) / 2. / np.log(Lambda))
+    return std
+
+
+def diffusion_coeff(t, Lambda, device='cpu'):
+    """Compute the diffusion coefficient of our SDE.
+
+    Args:
+        t: A vector of time steps.
+        Lambda: The $\lambda$ in our SDE.
+
+    Returns:
+        diff_coeff : The vector of diffusion coefficients.
+    """
+    diff_coeff = Lambda**t
+    return diff_coeff.to(device)
+
+
+def Euler_Maruyama_sampler(score_model,
+                          marginal_prob_std,
+                          diffusion_coeff,
+                          batch_size=1,
+                          x_shape=(4, 256, 256),
+                          num_steps=250,
+                          device='cuda',
+                          eps=1e-3, 
+                          y=None):
+    """Generate samples from score-based models with the Euler-Maruyama solver.
+
+    Args:
+        score_model: A PyTorch model that represents the time-dependent score-based model.
+        marginal_prob_std: A function that gives the standard deviation of
+            the perturbation kernel.
+        diffusion_coeff: A function that gives the diffusion coefficient of the SDE.
+        batch_size: The number of samplers to generate by calling this function once.
+        num_steps: The number of sampling steps.
+            Equivalent to the number of discretized time steps.
+        device: 'cuda' for running on GPUs, and 'cpu' for running on CPUs.
+        eps: The smallest time step for numerical stability.
+
+    Returns:
+        Samples.
+    """
+    t = torch.ones(batch_size).to(device)
+    r = torch.randn(batch_size, *x_shape).to(device)
+    init_x = r * marginal_prob_std(t)[:, None, None, None]
+    init_x = init_x.to(device)
+    time_steps = torch.linspace(1., eps, num_steps).to(device)
+    step_size = time_steps[0] - time_steps[1]
+    x = init_x
+    with torch.no_grad():
+        for time_step in tqdm(time_steps):
+            batch_time_step = torch.ones(batch_size, device=device) * time_step
+            g = diffusion_coeff(batch_time_step)
+            mean_x = x + (g**2)[:, None, None, None] * score_model(x, batch_time_step, y=y) * step_size
+            x = mean_x + torch.sqrt(step_size) * g[:, None, None, None] * torch.randn_like(x)
+    # Do not include any noise in the last sampling step.
+    return mean_x 

requirements.txt

@@ -0,0 +1,10 @@
+torch>=1.9.0
+torchvision>=0.10.0
+numpy>=1.21.0
+tqdm>=4.62.0
+matplotlib>=3.5.0
+seaborn>=0.11.0
+nibabel>=3.2.0
+monai>=0.9.0
+transformers>=4.20.0
+huggingface_hub>=0.10.0