Update app.py

app.py

@@ -134,56 +134,46 @@ class DiffusionModel(nn.Module):
         self.alphas = 1. - self.betas
         self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0))
 
-@torch.no_grad()
-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)
-
-    # Reverse diffusion process
-    for t in reversed(range(timesteps)):
-        if cancel_event.is_set():
-            return None
-            
-        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)
+    @torch.no_grad()
+    def sample(self, num_images, timesteps, img_size, num_classes, labels, device, progress_callback=None):
+        x_t = torch.randn((num_images, 3, img_size, img_size), device=device) * 0.5
         
-        # Improved denoising step
-        if t > 0:
-            noise = torch.randn_like(x_t) * 0.5  # Reduced noise scale
+        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:
-            noise = torch.zeros_like(x_t)
+            labels = labels.float().to(device)
+
+        for t in reversed(range(timesteps)):
+            if cancel_event.is_set():
+                return None
+                
+            t_tensor = torch.full((num_images,), t, device=device, dtype=torch.long)
             
-        # 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)
+            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)
+            
+            if t > 0:
+                noise = torch.randn_like(x_t) * 0.5
+            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)
+            
+            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
+        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)
+        
+        return x_t
 
 
 def load_model(model_path, device):
@@ -210,7 +200,6 @@ def load_model(model_path, device):
             # Verify model loading
             test_input = torch.randn(1, 3, IMG_SIZE, IMG_SIZE).to(device)
             test_labels = torch.zeros(1, NUM_CLASSES).to(device)
-            test_labels[0, 0] = 1
             test_time = torch.tensor([1]).to(device)
             output = unet(test_input, test_labels, test_time)
             print(f"Model test output shape: {output.shape}")
@@ -223,6 +212,18 @@ def load_model(model_path, device):
     
     diffusion_model.eval()
     return diffusion_model
+    
+MODEL_NAME = "model_weights.pth"
+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 if loading fails
+    print("Creating dummy model for demonstration")
+    loaded_model = DiffusionModel(UNet(num_classes=NUM_CLASSES)).to(device)
 
 def cancel_generation():
     cancel_event.set()
@@ -232,7 +233,6 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
     global loaded_model
     cancel_event.clear()
     
-    # Input validation
     if num_images < 1 or num_images > 10:
         raise gr.Error("Number of images must be between 1 and 10")
     
@@ -250,7 +250,6 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
                 raise gr.Error("Generation was cancelled by user")
 
         with torch.no_grad():
-            print(f"Generating {num_images} images for {label_str}")
             images = loaded_model.sample(
                 num_images=num_images,
                 timesteps=TIMESTEPS,
@@ -264,19 +263,13 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
         if images is None:
             return None, None
             
-        print(f"Generated images range: {images.min().item():.3f}, {images.max().item():.3f}")
-        
         processed_images = []
         for img in images:
-            # 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)
-            
-            # Convert to PIL Image
             pil_img = Image.fromarray(img_np)
             processed_images.append(pil_img)
         
-        # Return appropriate outputs based on count
         if num_images == 1:
             return processed_images[0], processed_images
         else:
@@ -287,7 +280,7 @@ def generate_images(label_str, num_images, progress=gr.Progress()):
         raise gr.Error(f"Generation failed: {str(e)}")
     finally:
         torch.cuda.empty_cache()
-    
+        
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(
     primary_hue="violet",