Update app.py

app.py

@@ -1,12 +1,40 @@
-# Import needed libraries
+# Import required libraries
 import streamlit as st
-import torchvision
+import torch
 from medigan import Generators
 from torchvision.transforms.functional import to_pil_image
-from torchvision.utils import make_grid
 
-# Define the GAN models available in the app
-model_ids = [
+# Apply critical patch for PGGAN model compatibility
+def apply_pggan_patch():
+    try:
+        from model19.base_GAN import BaseGAN
+        
+        # Monkey-patch the BaseGAN class initialization
+        original_init = BaseGAN.__init__
+        
+        def patched_init(self, dimLatentVector, **kwargs):
+            original_init(self, dimLatentVector, **kwargs)
+            
+            # Force correct optimizer parameters
+            self.optimizerG = torch.optim.Adam(
+                self.netG.parameters(),
+                lr=0.0002,
+                betas=(0.5, 0.999)  # Explicit tuple of floats
+            )
+            
+            self.optimizerD = torch.optim.Adam(
+                self.netD.parameters(),
+                lr=0.0002,
+                betas=(0.5, 0.999)  # Explicit tuple of floats
+            )
+        
+        BaseGAN.__init__ = patched_init
+        st.success("Applied PGGAN compatibility patch!")
+    except Exception as e:
+        st.warning(f"Patching failed: {str(e)}")
+
+# Model configuration
+MODEL_IDS = [
     "00001_DCGAN_MMG_CALC_ROI",
     "00002_DCGAN_MMG_MASS_ROI",
     "00003_CYCLEGAN_MMG_DENSITY_FULL",
@@ -15,75 +43,48 @@ model_ids = [
 ]
 
 def main():
-    # Setup page configuration
     st.set_page_config(page_title="MEDIGAN Generator", layout="wide")
-    
-    # Main page title and description
-    st.title("🧠 MEDIGAN Medical Image Generator")
-    st.markdown("""
-    **Generate synthetic medical images using GAN models.**  
-    🔍 Select model and parameters in the sidebar → Click **Generate Images**
-    """)
+    st.title("🧠 Medical Image Generator")
     
     # Sidebar controls
     with st.sidebar:
         st.header("⚙️ Settings")
-        model_id = st.selectbox("Select GAN Model", model_ids)
-        num_images = st.number_input("Number of Images", 1, 7, 1)
+        model_id = st.selectbox("Select Model", MODEL_IDS)
+        num_images = st.slider("Number of Images", 1, 8, 4)
         generate_btn = st.button("✨ Generate Images")
     
     # Main content area
     if generate_btn:
-        with st.spinner(f"Generating {num_images} image(s) using {model_id}..."):
-            generate_images(num_images, model_id)
-
-def torch_images(num_images, model_id):
-    generators = Generators()
-    dataloader = generators.get_as_torch_dataloader(
-        model_id=model_id,
-        install_dependencies=True,
-        num_samples=num_images,
-        num_workers=0,
-        prefetch_factor=None
-    )
-
-    images = []
-    for _, data_dict in enumerate(dataloader):
-        batch_images = []
-        for tensor in data_dict.values():
-            if tensor.dim() == 4:
-                tensor = tensor.squeeze(0).permute(2, 0, 1)
-            img = to_pil_image(tensor).convert("RGB")
-            batch_images.append(img)
-        
-        # Create image grid
-        grid_tensor = make_grid(
-            [torchvision.transforms.ToTensor()(img) for img in batch_images],
-            nrow=2 if len(batch_images) > 1 else 1
-        )
-        grid_img = to_pil_image(grid_tensor)
-        images.append(grid_img)
-    
-    return images
+        with st.spinner(f"Generating {num_images} images..."):
+            try:
+                generate_images(num_images, model_id)
+            except Exception as e:
+                st.error(f"Generation failed: {str(e)}")
 
 def generate_images(num_images, model_id):
-    # Clear previous results
-    st.empty()
+    # Apply model-specific patches
+    if "00019" in model_id:
+        apply_pggan_patch()
     
-    # Generate and display new images
-    images = torch_images(num_images, model_id)
-    
-    # Create columns for responsive layout
-    cols = st.columns(len(images))
+    # Initialize generator
+    generators = Generators()
     
-    for col, img in zip(cols, images):
-        with col:
-            st.image(
-                img,
-                caption=f"Generated by: {model_id}",
-                width=300
-            )
-            st.markdown("---")
+    # Generate and display images
+    cols = st.columns(4)
+    for i in range(num_images):
+        with cols[i % 4]:
+            try:
+                # Generate single sample
+                sample = generators.generate(
+                    model_id=model_id,
+                    num_samples=1,
+                    install_dependencies=True
+                )
+                img = to_pil_image(sample[0]).convert("RGB")
+                st.image(img, caption=f"Image {i+1}", use_column_width=True)
+                st.markdown("---")
+            except Exception as e:
+                st.error(f"Error generating image {i+1}: {str(e)}")
 
 if __name__ == "__main__":
     main()