Update app.py

app.py

@@ -125,26 +125,21 @@ class DiffusionModel(nn.Module):
         self.model = model
         self.timesteps = timesteps
         
-        # Better noise schedule for medical images
+        # More conservative noise schedule
         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.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
+        # Initialize with standard normal distribution (scale=1.0)
+        x_t = torch.randn((num_images, 3, img_size, img_size), device=device)
         
         if labels.ndim == 1:
             labels_one_hot = torch.zeros(num_images, num_classes, device=device)
@@ -167,37 +162,39 @@ class DiffusionModel(nn.Module):
             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
+            # Calculate coefficients
             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]
+            sqrt_recip_alpha_t = torch.sqrt(1.0 / alpha_t)
+            sqrt_one_minus_alpha_bar_t = torch.sqrt(1.0 - alpha_bar_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 predicted x0
+            pred_x0 = (x_t - sqrt_one_minus_alpha_bar_t * pred_noise) * sqrt_recip_alpha_t
             
             # Calculate direction pointing to x_t
-            dir_xt = torch.sqrt(1. - alpha_bar_t_prev - beta_t**2) * pred_noise
+            pred_dir = torch.sqrt(1.0 - alpha_bar_t_prev) * pred_noise
             
             # Noise for next step
             if t > 0:
-                noise = torch.randn_like(x_t) * 0.25  # Reduced noise scale
+                noise = torch.randn_like(x_t) * 0.5
             else:
                 noise = torch.zeros_like(x_t)
                 
-            # Update x_t
-            x_t = torch.sqrt(alpha_bar_t_prev) * pred_x0 + dir_xt + noise * torch.sqrt(beta_t)
+            # Update x_t with stability checks
+            x_t = torch.sqrt(alpha_bar_t_prev) * pred_x0 + pred_dir + noise * torch.sqrt(beta_t)
             
+            # Numerical stability check
+            if torch.isnan(x_t).any() or torch.isinf(x_t).any():
+                x_t = torch.randn_like(x_t) * 0.1
+                
             if progress_callback:
                 progress_callback((timesteps - t) / timesteps)
 
-        # Better normalization approach
-        x_t = torch.clamp(x_t, -1., 1.)
-        x_t = (x_t + 1) / 2  # Scale to [0, 1]
+        # Gentle normalization
+        x_t = (x_t - x_t.min()) / (x_t.max() - x_t.min() + 1e-8)  # [0, 1]
+        x_t = torch.clamp(x_t, 0, 1)  # Final safety clamp
         
         return x_t
 
-
 def load_model(model_path, device):
     unet = UNet(num_classes=NUM_CLASSES).to(device)
     diffusion_model = DiffusionModel(unet).to(device)