Update app.py

app.py

@@ -124,18 +124,26 @@ class DiffusionModel(nn.Module):
         super().__init__()
         self.model = model
         self.timesteps = timesteps
-        self.time_dim = time_dim
-
-        # Fix 1: Ensure consistent float32 types
+        
+        # Better noise schedule for medical images
         scale = 1000 / timesteps
         beta_start = scale * 0.0001
         beta_end = scale * 0.02
         self.betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float32)
+        
+        # More stable alpha calculations
         self.alphas = 1. - self.betas
-        self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0))
+        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
+        self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(self.alphas_cumprod))
+        self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1. - self.alphas_cumprod))
+        
+        # Parameters for posterior variance
+        self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / self.alphas_cumprod))
+        self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / self.alphas_cumprod - 1))
 
     @torch.no_grad()
     def sample(self, num_images, timesteps, img_size, num_classes, labels, device, progress_callback=None):
+        # Initialize with proper scale
         x_t = torch.randn((num_images, 3, img_size, img_size), device=device) * 0.5
         
         if labels.ndim == 1:
@@ -151,27 +159,41 @@ class DiffusionModel(nn.Module):
                 
             t_tensor = torch.full((num_images,), t, device=device, dtype=torch.long)
             
+            # Predict noise with model
             pred_noise = self.model(x_t, labels, t_tensor.float())
             
-            alpha_t = self.alphas[t].to(device)
-            alpha_bar_t = self.alpha_bars[t].to(device)
-            beta_t = self.betas[t].to(device)
+            # Get current alpha values
+            alpha_t = self.alphas[t]
+            alpha_bar_t = self.alphas_cumprod[t]
+            alpha_bar_t_prev = self.alphas_cumprod[t-1] if t > 0 else torch.tensor(1.0)
+            
+            # Calculate coefficients for cleaner sampling
+            beta_t = self.betas[t]
+            sqrt_recip_alphas_t = self.sqrt_recip_alphas_cumprod[t]
+            sqrt_one_minus_alphas_bar_t = self.sqrt_one_minus_alphas_cumprod[t]
+            
+            # Main denoising equation
+            pred_x0 = (x_t - sqrt_one_minus_alphas_bar_t * pred_noise) / sqrt_recip_alphas_t
+            pred_x0 = torch.clamp(pred_x0, -1., 1.)
+            
+            # Calculate direction pointing to x_t
+            dir_xt = torch.sqrt(1. - alpha_bar_t_prev - beta_t**2) * pred_noise
             
+            # Noise for next step
             if t > 0:
-                noise = torch.randn_like(x_t) * 0.5
+                noise = torch.randn_like(x_t) * 0.25  # Reduced noise scale
             else:
                 noise = torch.zeros_like(x_t)
                 
-            x_t = (x_t - (1 - alpha_t)/torch.sqrt(1 - alpha_bar_t) * pred_noise) / torch.sqrt(alpha_t)
-            x_t = x_t + noise * torch.sqrt(beta_t)
+            # Update x_t
+            x_t = torch.sqrt(alpha_bar_t_prev) * pred_x0 + dir_xt + noise * torch.sqrt(beta_t)
             
             if progress_callback:
                 progress_callback((timesteps - t) / timesteps)
 
+        # Better normalization approach
         x_t = torch.clamp(x_t, -1., 1.)
-        min_val = x_t.min()
-        max_val = x_t.max()
-        x_t = (x_t - min_val) / (max_val - min_val + 1e-8)
+        x_t = (x_t + 1) / 2  # Scale to [0, 1]
         
         return x_t