Update app.py

app.py

@@ -5,16 +5,21 @@ from PIL import Image
 import numpy as np
 import math
 import os
+from threading import Event
+import traceback
 
-# Constants (update these to match your training config)
+# Constants
 IMG_SIZE = 128
 TIMESTEPS = 300
 NUM_CLASSES = 2
 
-# Define the device
+# Global Cancellation Flag
+cancel_event = Event()
+
+# Device Configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Define the SinusoidalPositionEmbeddings class
+# --- Model Definitions (from second file) ---
 class SinusoidalPositionEmbeddings(nn.Module):
     def __init__(self, dim):
         super().__init__()
@@ -34,11 +39,9 @@ class SinusoidalPositionEmbeddings(nn.Module):
         output = torch.cat([embeddings.sin(), embeddings.cos()], dim=-1)
         return output
 
-# Define the UNet class
 class UNet(nn.Module):
     def __init__(self, in_channels=3, out_channels=3, num_classes=2, time_dim=256):
         super().__init__()
-
         self.num_classes = num_classes
         self.label_embedding = nn.Embedding(num_classes, time_dim)
 
@@ -121,7 +124,6 @@ class UNet(nn.Module):
         output = self.outc(x)
         return output
 
-# Define the DiffusionModel class
 class DiffusionModel(nn.Module):
     def __init__(self, model, timesteps=500, time_dim=256):
         super().__init__()
@@ -154,7 +156,7 @@ class DiffusionModel(nn.Module):
         return predicted_noise, noise, t
 
 @torch.no_grad()
-def sample(model, num_images, timesteps, img_size, num_classes, labels, device):
+def sample(model, num_images, timesteps, img_size, num_classes, labels, device, progress_callback=None):
     x_t = torch.randn(num_images, 3, img_size, img_size).to(device)
 
     if labels.ndim == 1:
@@ -165,6 +167,9 @@ def sample(model, num_images, timesteps, img_size, num_classes, labels, device):
         labels = labels.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.float)
 
         predicted_noise = model.model(x_t, labels, t_tensor)
@@ -182,6 +187,9 @@ def sample(model, num_images, timesteps, img_size, num_classes, labels, device):
             noise = torch.zeros_like(x_t)
 
         x_t = mean + torch.sqrt(variance) * noise
+        
+        if progress_callback:
+            progress_callback((timesteps - t) / timesteps)
 
     x_0 = torch.clamp(x_t, -1., 1.)
 
@@ -192,14 +200,12 @@ def sample(model, num_images, timesteps, img_size, num_classes, labels, device):
 
     return x_0
 
-# Load the trained model with improved error handling
 def load_model(model_path, device):
     unet_model = UNet(num_classes=NUM_CLASSES).to(device)
     diffusion_model = DiffusionModel(unet_model, timesteps=TIMESTEPS).to(device)
     
     try:
         checkpoint = torch.load(model_path, map_location=device)
-        # Handle both full model and state_dict loading
         if 'model_state_dict' in checkpoint:
             diffusion_model.model.load_state_dict(checkpoint['model_state_dict'])
         else:
@@ -212,19 +218,20 @@ def load_model(model_path, device):
     diffusion_model.eval()
     return diffusion_model
 
-# Improved image generation function
-def generate_image(label_str):
+def cancel_generation():
+    cancel_event.set()
+    return "Generation cancelled"
+
+def generate_single_image(label_str):
     label_map = {'Pneumonia': 0, 'Pneumothorax': 1}
     try:
         label_index = label_map[label_str]
     except KeyError:
         raise gr.Error(f"Invalid label '{label_str}'. Please select either 'Pneumonia' or 'Pneumothorax'.")
 
-    # Create one-hot encoded label
     labels = torch.zeros(1, NUM_CLASSES, device=device)
     labels[0, label_index] = 1
 
-    # Generate image
     with torch.no_grad():
         generated_image = sample(
             model=loaded_model,
@@ -236,39 +243,155 @@ def generate_image(label_str):
             device=device
         )
 
-    # Convert to PIL Image
     img_np = generated_image.squeeze(0).cpu().permute(1, 2, 0).numpy()
-    img_np = np.clip(img_np, 0, 1)  # Ensure proper range
+    img_np = np.clip(img_np, 0, 1)
     img_pil = Image.fromarray((img_np * 255).astype(np.uint8))
     
     return img_pil
 
-# Model paths (update these for your deployment)
-MODEL_DIR = "models"
-MODEL_NAME = "diffusion_unet_xray.pth"  # Update with your actual filename
-model_path = os.path.join(MODEL_DIR, MODEL_NAME)
+def generate_batch_images(label_str, num_images, progress=gr.Progress()):
+    global loaded_model
+    cancel_event.clear()
+    
+    if num_images < 1 or num_images > 10:
+        raise gr.Error("Number of images must be between 1 and 10")
+    
+    label_map = {'Pneumonia': 0, 'Pneumothorax': 1}
+    if label_str not in label_map:
+        raise gr.Error("Invalid condition selected")
+
+    labels = torch.zeros(num_images, NUM_CLASSES, device=device)
+    labels[:, label_map[label_str]] = 1
+
+    try:
+        def progress_callback(progress_val):
+            progress(progress_val, desc="Generating...")
+            if cancel_event.is_set():
+                raise gr.Error("Generation was cancelled by user")
+
+        with torch.no_grad():
+            images = sample(
+                model=loaded_model,
+                num_images=num_images,
+                timesteps=TIMESTEPS,
+                img_size=IMG_SIZE,
+                num_classes=NUM_CLASSES,
+                labels=labels,
+                device=device,
+                progress_callback=progress_callback
+            )
+        
+        if images is None:
+            return None
+            
+        processed_images = []
+        for img in images:
+            img_np = img.cpu().permute(1, 2, 0).numpy()
+            img_np = np.clip(img_np, 0, 1)
+            pil_img = Image.fromarray((img_np * 255).astype(np.uint8))
+            processed_images.append(pil_img)
+        
+        return 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
+    finally:
+        torch.cuda.empty_cache()
 
 # Load model
+MODEL_DIR = "models"
+MODEL_NAME = "diffusion_unet_xray.pth"
+model_path = os.path.join(MODEL_DIR, MODEL_NAME)
 print("Loading model...")
 loaded_model = load_model(model_path, device)
 print("Model loaded successfully!")
 
-# Gradio interface
-iface = gr.Interface(
-    fn=generate_image,
-    inputs=gr.Dropdown(
-        choices=["Pneumonia", "Pneumothorax"],
-        label="Select Condition",
-        value="Pneumonia"  # Default value
-    ),
-    outputs=gr.Image(
-        type="pil",
-        label="Generated X-ray Image"
-    ),
-    title="Medical X-ray Image Generator",
-    description="Generate synthetic chest X-ray images using a diffusion model. Select a condition to generate.",
-    examples=[["Pneumonia"], ["Pneumothorax"]]
-)
+# --- Gradio UI (from first file with modifications) ---
+with gr.Blocks(theme=gr.themes.Soft(
+    primary_hue="violet",
+    neutral_hue="slate",
+    font=[gr.themes.GoogleFont("Poppins")],
+    text_size="md"
+)) as demo:
+    gr.Markdown("""
+    <center>
+    <h1>Synthetic X-ray Generator</h1>
+    <p><em>Generate synthetic chest X-rays conditioned on pathology</em></p>
+    </center>
+    """)
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            condition = gr.Dropdown(
+                ["Pneumonia", "Pneumothorax"],
+                label="Select Condition",
+                value="Pneumonia",
+                interactive=True
+            )
+            num_images = gr.Slider(
+                1, 10, value=1, step=1,
+                label="Number of Images",
+                interactive=True
+            )
+            
+            with gr.Row():
+                submit_btn = gr.Button("Generate", variant="primary")
+                cancel_btn = gr.Button("Cancel", variant="stop")
+            
+            gr.Markdown("""
+            <div style="text-align: center; margin-top: 10px;">
+                <small>Note: Generation may take several seconds per image</small>
+            </div>
+            """)
+        
+        with gr.Column(scale=2):
+            with gr.Tab("Single Image"):
+                single_image = gr.Image(
+                    type="pil",
+                    label="Generated X-ray",
+                    height=400
+                )
+            with gr.Tab("Batch Images"):
+                gallery = gr.Gallery(
+                    label="Generated X-rays",
+                    columns=3,
+                    height="auto",
+                    object_fit="contain"
+                )
+    
+    # Single image generation
+    condition.change(
+        fn=generate_single_image,
+        inputs=condition,
+        outputs=single_image
+    )
+    
+    # Batch image generation
+    submit_btn.click(
+        fn=generate_batch_images,
+        inputs=[condition, num_images],
+        outputs=gallery
+    )
+    
+    cancel_btn.click(
+        fn=cancel_generation,
+        outputs=None
+    )
+
+    demo.css = """
+    .gradio-container {
+        background: linear-gradient(135deg, #f5f7fa 0%, #e4e8f0 100%);
+    }
+    .gallery-container {
+        background-color: white !important;
+    }
+    """
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)