Update hf_diffusion_service.py

hf_diffusion_service.py

@@ -1,311 +1,194 @@
-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 
+import os
+import torch
+import numpy as np
+import torchvision.transforms as transforms
+from PIL import Image
+
+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):
+        # 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 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 CUDA
+        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. Using CPU (this will be slow).")
+
+        self.device = torch.device('cuda:0' if cuda_available else 'cpu')
+        self.Lambda = 25.0
+
+        # Initialize 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)
+
+        # ✅ Auto-detect model file path
+        model_candidates = [
+            "hf_model_files/pytorch_model.bin",
+            "pytorch_model.bin"
+        ]
+        self.model_path = next((path for path in model_candidates if os.path.exists(path)), None)
+
+        if not self.model_path:
+            raise FileNotFoundError("pytorch_model.bin not found in root or hf_model_files folder.")
+
+        print(f"Loading diffusion model from: {self.model_path}")
+
+        # Load model weights
+        try:
+            state_dict = torch.load(self.model_path, map_location=self.device)
+
+            # Analyze the state dict to configure model
+            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: {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: {actual_input_dim_mask}")
+
+                    # 1-channel model with 256*256 flattened mask
+                    if actual_input_channels == 1 and actual_input_dim_mask == 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:
+                        self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+                        self.input_channels = 4
+                        self.input_dim_mask = 262144
+                else:
+                    self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+                    self.input_channels = 4
+                    self.input_dim_mask = 262144
+            else:
+                # Default to original architecture
+                self.score_model = UNet(marginal_prob_std=self.marginal_prob_std_fn)
+                self.input_channels = 4
+                self.input_dim_mask = 262144
+
+            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")
+            print(f"   Input channels: {self.input_channels}, Mask dim: {self.input_dim_mask}")
+
+        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."""
+        try:
+            processed_mask = self.process_mask(mask)
+            generated_tensor = self.generate_from_mask(processed_mask)
+            return self.tensor_to_image(generated_tensor)
+        except Exception as e:
+            print(f"❌ Error generating image: {e}")
+            return None
+
+    def process_mask(self, mask):
+        """Process the input mask to the correct format for the model."""
+        try:
+            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)
+            elif isinstance(mask, np.ndarray):
+                if mask.ndim == 2:
+                    mask = mask[np.newaxis, :, :]
+                tensor = torch.from_numpy(mask).float()
+                if tensor.dim() == 3:
+                    tensor = tensor.unsqueeze(0)
+            elif isinstance(mask, torch.Tensor):
+                tensor = mask
+                if tensor.dim() == 3:
+                    tensor = tensor.unsqueeze(0)
+            else:
+                raise ValueError(f"Unsupported mask type: {type(mask)}")
+
+            # Adjust channels
+            if self.input_channels == 1:
+                if tensor.shape[1] != 1:
+                    tensor = tensor.mean(dim=1, keepdim=True)
+            else:
+                if tensor.shape[1] == 1:
+                    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 256x256 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 diffusion model."""
+        try:
+            x_shape = (1, 256, 256) if self.input_channels == 1 else (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
+                )
+            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."""
+        try:
+            if tensor.shape[1] > 1:
+                image_tensor = tensor.squeeze(0).mean(dim=0)
+            else:
+                image_tensor = tensor.squeeze(0).squeeze(0)
+
+            image_array = (image_tensor.cpu().numpy() * 255).astype(np.uint8)
+            return Image.fromarray(image_array, mode='L')
+        except Exception as e:
+            print(f"❌ Error converting tensor to image: {e}")
+            raise e