Update app.py

app.py

@@ -10,7 +10,7 @@ import traceback
 
 # Constants
 IMG_SIZE = 128
-TIMESTEPS = 300
+TIMESTEPS = 500
 NUM_CLASSES = 2
 
 # Global Cancellation Flag
@@ -135,47 +135,56 @@ class DiffusionModel(nn.Module):
         self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0))
 
     @torch.no_grad()
-    def sample(self, num_images, img_size, num_classes, labels, device, progress_callback=None):
-        # Initialize with noise
-        x_t = torch.randn((num_images, 3, img_size, img_size), device=device, dtype=torch.float32)
-        
-        # Convert labels to proper format
-        if labels.ndim == 1:
-            labels_one_hot = torch.zeros(num_images, num_classes, device=device)
-            labels_one_hot[torch.arange(num_images), labels] = 1
-            labels = labels_one_hot
-        else:
-            labels = labels.to(device)
+def sample(model, num_images, timesteps, img_size, num_classes, labels, device, progress_callback=None):
+    # Initialize with properly scaled noise
+    x_t = torch.randn((num_images, 3, img_size, img_size), device=device) * 0.5  # Reduced initial noise scale
+    
+    # Convert labels to proper format
+    if labels.ndim == 1:
+        labels_one_hot = torch.zeros(num_images, num_classes, device=device)
+        labels_one_hot[torch.arange(num_images), labels] = 1
+        labels = labels_one_hot
+    else:
+        labels = labels.float().to(device)
 
-        for i in reversed(range(0, self.timesteps)):
-            if cancel_event.is_set():
-                return None
-                
-            t = torch.full((num_images,), i, device=device, dtype=torch.long)
+    # Reverse diffusion process
+    for t in reversed(range(timesteps)):
+        if cancel_event.is_set():
+            return None
             
-            # Model prediction with type stability
-            pred_noise = self.model(x_t, labels, t.float())
-            
-            # Calculate diffusion parameters
-            beta_t = self.betas[t].view(-1, 1, 1, 1).to(device)
-            alpha_t = self.alphas[t].view(-1, 1, 1, 1).to(device)
-            alpha_bar_t = self.alpha_bars[t].view(-1, 1, 1, 1).to(device)
-            
-            # Improved denoising step (Fix 2)
-            if i > 0:
-                noise = torch.randn_like(x_t)
-            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 += noise * torch.sqrt(beta_t)
+        t_tensor = torch.full((num_images,), t, device=device, dtype=torch.long)
+        
+        # Model prediction with proper scaling
+        pred_noise = model.model(x_t, labels, t_tensor.float())
+        
+        # Calculate diffusion parameters
+        alpha_t = model.alphas[t].to(device)
+        alpha_bar_t = model.alpha_bars[t].to(device)
+        beta_t = model.betas[t].to(device)
+        
+        # Improved denoising step
+        if t > 0:
+            noise = torch.randn_like(x_t) * 0.5  # Reduced noise scale
+        else:
+            noise = torch.zeros_like(x_t)
             
-            if progress_callback:
-                progress_callback((self.timesteps - i) / self.timesteps)
+        # More stable prediction
+        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)
+        
+        if progress_callback:
+            progress_callback((timesteps - t) / timesteps)
+
+    # Better image normalization
+    x_t = torch.clamp(x_t, -1., 1.)
+    
+    # Alternative normalization approach
+    min_val = x_t.min()
+    max_val = x_t.max()
+    x_t = (x_t - min_val) / (max_val - min_val + 1e-8)  # Ensure we don't divide by zero
+    
+    return x_t
 
-        # Fix 3: Simplified scaling
-        x_t = torch.clamp(x_t, -1., 1.)
-        return (x_t + 1) / 2  # Scale to [0,1]
 
 def load_model(model_path, device):
     unet = UNet(num_classes=NUM_CLASSES).to(device)
@@ -231,7 +240,7 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
     if label_str not in label_map:
         raise gr.Error("Invalid condition selected")
 
-    labels = torch.zeros(num_images, NUM_CLASSES, device=device, dtype=torch.float32)
+    labels = torch.zeros(num_images, NUM_CLASSES, device=device)
     labels[:, label_map[label_str]] = 1
 
     try:
@@ -242,10 +251,9 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
 
         with torch.no_grad():
             print(f"Generating {num_images} images for {label_str}")
-            print(f"Labels shape: {labels.shape}, device: {labels.device}")
-            
             images = loaded_model.sample(
                 num_images=num_images,
+                timesteps=TIMESTEPS,
                 img_size=IMG_SIZE,
                 num_classes=NUM_CLASSES,
                 labels=labels,
@@ -256,17 +264,15 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
         if images is None:
             return None, None
             
-        # Diagnostic print
         print(f"Generated images range: {images.min().item():.3f}, {images.max().item():.3f}")
         
         processed_images = []
         for img in images:
-            # Fix 3: Improved image conversion
-            img_np = (img.cpu().numpy().transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)
-            print(f"Image range after conversion: {img_np.min()}, {img_np.max()}")
+            # Convert to numpy and ensure proper range
+            img_np = img.cpu().numpy().transpose(1, 2, 0)
+            img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
             
-            if img_np.shape[2] == 1:  # Handle grayscale if needed
-                img_np = img_np.squeeze(-1)
+            # Convert to PIL Image
             pil_img = Image.fromarray(img_np)
             processed_images.append(pil_img)
         
@@ -276,30 +282,11 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
         else:
             return None, processed_images
 
-    except torch.cuda.OutOfMemoryError:
-        torch.cuda.empty_cache()
-        raise gr.Error("Out of GPU memory - try generating fewer images")
     except Exception as e:
         traceback.print_exc()
-        if str(e) != "Generation was cancelled by user":
-            raise gr.Error(f"Generation failed: {str(e)}")
-        return None, None
+        raise gr.Error(f"Generation failed: {str(e)}")
     finally:
         torch.cuda.empty_cache()
-
-# Load model
-MODEL_NAME = "model_weights.pth"  # Updated to look in root folder
-model_path = MODEL_NAME
-print("Loading model...")
-try:
-    loaded_model = load_model(model_path, device)
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Failed to load model: {e}")
-    # Create a dummy model for demo purposes
-    print("Creating dummy model for demonstration")
-    loaded_model = DiffusionModel(UNet(num_classes=NUM_CLASSES)).to(device)
-
     
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(