Update app.py

app.py

@@ -120,64 +120,58 @@ class UNet(nn.Module):
         return output
 
 class DiffusionModel(nn.Module):
-    def __init__(self, model, timesteps=TIMESTEPS, time_dim=256):
+    def __init__(self, model, timesteps=TIMESTEPS):
         super().__init__()
         self.model = model
         self.timesteps = timesteps
         
-        # Noise schedule from working code
+        # Use the exact same noise schedule as Colab
         beta_start = 0.0001
         beta_end = 0.02
         self.betas = torch.linspace(beta_start, beta_end, timesteps, dtype=torch.float32)
-        
         self.alphas = 1. - self.betas
         self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0))
-        self.register_buffer('sqrt_one_minus_alpha_bars', torch.sqrt(1. - self.alpha_bars))
-        self.register_buffer('sqrt_recip_alphas', torch.sqrt(1. / self.alphas))
 
     @torch.no_grad()
     def sample(self, num_images, img_size, num_classes, labels, device, progress_callback=None):
-        """Improved sampling method based on working code"""
+        """Identical implementation to Colab version"""
+        # Start with random noise (same scale)
         x_t = torch.randn((num_images, 3, img_size, img_size), device=device)
         
-        # Handle labels (class indices or one-hot)
+        # Identical label handling
         if labels.ndim == 1:
             labels = torch.zeros(num_images, num_classes, device=device).scatter_(1, labels.unsqueeze(1), 1)
-        else:
-            labels = labels.float().to(device)
+        labels = labels.to(device)
 
+        # Same sampling loop
         for t in reversed(range(self.timesteps)):
             if cancel_event.is_set():
                 return None
                 
             t_tensor = torch.full((num_images,), t, device=device, dtype=torch.float)
-            
-            # Predict noise with model
-            pred_noise = self.model(x_t, labels, t_tensor)
-            
-            # Calculate coefficients from working code
+            predicted_noise = self.model(x_t, labels, t_tensor)
+
+            # Identical coefficients calculation
             beta_t = self.betas[t].to(device)
             alpha_t = self.alphas[t].to(device)
             alpha_bar_t = self.alpha_bars[t].to(device)
-            
-            # Improved reverse diffusion step
-            mean = (1 / torch.sqrt(alpha_t)) * (x_t - (beta_t / torch.sqrt(1 - alpha_bar_t)) * pred_noise)
+
+            # Same mean/variance calculation
+            mean = (1 / torch.sqrt(alpha_t)) * (x_t - (beta_t / torch.sqrt(1 - alpha_bar_t)) * predicted_noise)
             variance = beta_t
-            
+
             if t > 0:
                 noise = torch.randn_like(x_t)
             else:
                 noise = torch.zeros_like(x_t)
-                
+
             x_t = mean + torch.sqrt(variance) * noise
             
             if progress_callback:
                 progress_callback((self.timesteps - t) / self.timesteps)
 
-        # Improved normalization from working code
+        # Identical denormalization
         x_t = torch.clamp(x_t, -1., 1.)
-        
-        # Denormalize using ImageNet stats (from working code)
         mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1).to(device)
         std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1).to(device)
         x_t = std * x_t + mean