Update hf_diffusion_service.py

hf_diffusion_service.py

@@ -1,35 +1,29 @@
-import os
 import torch
 import numpy as np
-import torchvision.transforms as transforms
 from PIL import Image
+import torchvision.transforms as transforms
+import io
+import base64
 
 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."""
-
+    """A UNet model that's compatible with saved weights (handles 1-channel input)."""
+    
     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
+        # Accept 1-channel input
         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
-            )
+            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."""
+    """Handles loading the conditional diffusion model and generating CT images."""
 
     def __init__(self):
-        # Check CUDA
         cuda_available = torch.cuda.is_available()
         print(f"CUDA available for HF diffusion: {cuda_available}")
         if not cuda_available:
@@ -37,84 +31,80 @@ class HFDiffusionService:
 
         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.")
+        # Model path (make sure pytorch_model.bin is present)
+        self.model_path = "pytorch_model.bin"
+        self.input_channels = 1
+        self.input_dim_mask = 65536
 
-        print(f"Loading diffusion model from: {self.model_path}")
+        # Load model
+        self._load_model()
 
-        # Load model weights
+    def _load_model(self):
         try:
+            print(f"Loading diffusion model from: {self.model_path}")
             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
+                self.input_channels = conv1_weight.shape[1]
+                print(f"Detected input channels: {self.input_channels}")
+            if cond_embed_weight is not None:
+                self.input_dim_mask = cond_embed_weight.shape[1]
+                print(f"Detected input_dim_mask: {self.input_dim_mask}")
+
+            # Initialize compatible UNet
+            if self.input_channels == 1 and self.input_dim_mask == 65536:
+                self.score_model = CompatibleUNet(
+                    marginal_prob_std=self.marginal_prob_std_fn,
+                    input_dim_mask=self.input_dim_mask
+                )
             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}")
+            print(f"✅ HF Diffusion model loaded successfully\n   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."""
+        """
+        Generate a CT image from a segmentation mask and return it as PIL Image.
+        """
         try:
-            processed_mask = self.process_mask(mask)
-            generated_tensor = self.generate_from_mask(processed_mask)
-            return self.tensor_to_image(generated_tensor)
+            processed_mask = self._process_mask(mask)
+            tensor_image = self._generate_from_mask(processed_mask)
+            return self._tensor_to_image(tensor_image)
         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."""
+    def generate_image_base64(self, mask):
+        """
+        Generate a CT image and return it as a base64 string (data URI).
+        """
+        image = self.generate_image(mask)
+        if image is None:
+            return None
+
+        buffer = io.BytesIO()
+        image.save(buffer, format="PNG")
+        base64_img = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        return f"data:image/png;base64,{base64_img}"
+
+    def _process_mask(self, mask):
+        """
+        Convert input mask (PIL, np.array, or tensor) into model-ready tensor.
+        """
         try:
             if isinstance(mask, Image.Image):
                 transform = transforms.Compose([
@@ -122,13 +112,15 @@ class HFDiffusionService:
                     transforms.Resize((256, 256), antialias=True),
                     transforms.ToTensor()
                 ])
-                tensor = transform(mask).unsqueeze(0)
+                tensor = transform(mask).unsqueeze(0)  # [1, 1, 256, 256]
+
             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)
+                    tensor = tensor.unsqueeze(0)  # [1, 1, 256, 256]
+
             elif isinstance(mask, torch.Tensor):
                 tensor = mask
                 if tensor.dim() == 3:
@@ -136,32 +128,23 @@ class HFDiffusionService:
             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
-                )
+            if tensor.shape[2:] != (256, 256):
+                tensor = torch.nn.functional.interpolate(tensor, size=(256, 256), mode='bilinear', align_corners=False)
+
+            if tensor.shape[1] == 1 and self.input_channels > 1:
+                tensor = tensor.repeat(1, self.input_channels, 1, 1)
 
-            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."""
+    def _generate_from_mask(self, conditioning_mask, num_steps=250, eps=1e-3):
+        """
+        Diffusion sampling given a mask, returns tensor in [0,1].
+        """
         try:
-            x_shape = (1, 256, 256) if self.input_channels == 1 else (4, 256, 256)
+            x_shape = (self.input_channels, 256, 256)
             with torch.no_grad():
                 samples = Euler_Maruyama_sampler(
                     self.score_model,
@@ -176,19 +159,22 @@ class HFDiffusionService:
                 )
             return samples.clamp(0, 1)
         except Exception as e:
-            print(f"❌ Error in generate_from_mask: {e}")
+            print(f"❌ Error in diffusion sampling: {e}")
             raise e
 
-    def tensor_to_image(self, tensor):
-        """Convert tensor to PIL Image."""
+    def _tensor_to_image(self, tensor):
+        """
+        Convert tensor -> RGB PIL image.
+        """
         try:
-            if tensor.shape[1] > 1:
-                image_tensor = tensor.squeeze(0).mean(dim=0)
+            tensor = tensor.squeeze(0)  # [C, H, W]
+            if tensor.shape[0] > 1:
+                image_array = (tensor.mean(dim=0).cpu().numpy() * 255).astype(np.uint8)
             else:
-                image_tensor = tensor.squeeze(0).squeeze(0)
+                image_array = (tensor[0].cpu().numpy() * 255).astype(np.uint8)
 
-            image_array = (image_tensor.cpu().numpy() * 255).astype(np.uint8)
-            return Image.fromarray(image_array, mode='L')
+            img_gray = Image.fromarray(image_array, mode='L')
+            return img_gray.convert("RGB")  # Always RGB for frontend
         except Exception as e:
             print(f"❌ Error converting tensor to image: {e}")
             raise e