Update app.py

app.py

@@ -1,39 +1,40 @@
-# Import required libraries
 import streamlit as st
 import torch
+import os
 from medigan import Generators
 from torchvision.transforms.functional import to_pil_image
 
-# Apply critical patch for PGGAN model compatibility
-def apply_pggan_patch():
+def install_and_patch_model(model_id):
     try:
-        from model19.base_GAN import BaseGAN
+        # Install model if not exists
+        generators = Generators()
+        if model_id not in generators.get_model_ids():
+            generators.download_and_install_model(model_id)
         
-        # Monkey-patch the BaseGAN class initialization
-        original_init = BaseGAN.__init__
+        # Locate model files
+        model_path = os.path.join(os.path.expanduser("~"), ".medigan", "models", model_id)
         
-        def patched_init(self, dimLatentVector, **kwargs):
-            original_init(self, dimLatentVector, **kwargs)
+        # Patch base_GAN.py
+        base_gan_path = os.path.join(model_path, "model19", "base_GAN.py")
+        if os.path.exists(base_gan_path):
+            with open(base_gan_path, "r") as f:
+                content = f.read()
             
-            # Force correct optimizer parameters
-            self.optimizerG = torch.optim.Adam(
-                self.netG.parameters(),
-                lr=0.0002,
-                betas=(0.5, 0.999)  # Explicit tuple of floats
+            # Replace problematic optimizer lines
+            content = content.replace(
+                "betas=(self.beta1, self.beta2)",
+                "betas=(0.5, 0.999)"  # Force correct beta values
             )
             
-            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!")
+            with open(base_gan_path, "w") as f:
+                f.write(content)
+            
+            st.success("Model successfully patched!")
+        return True
     except Exception as e:
-        st.warning(f"Patching failed: {str(e)}")
+        st.error(f"Patching failed: {str(e)}")
+        return False
 
-# Model configuration
 MODEL_IDS = [
     "00001_DCGAN_MMG_CALC_ROI",
     "00002_DCGAN_MMG_MASS_ROI",
@@ -46,45 +47,40 @@ def main():
     st.set_page_config(page_title="MEDIGAN Generator", layout="wide")
     st.title("🧠 Medical Image Generator")
     
-    # Sidebar controls
     with st.sidebar:
         st.header("⚙️ Settings")
         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} images..."):
-            try:
+        
+        if st.button("✨ Generate Images"):
+            if "00019" in model_id and not install_and_patch_model(model_id):
+                return
+                
+            with st.spinner("Generating images..."):
                 generate_images(num_images, model_id)
-            except Exception as e:
-                st.error(f"Generation failed: {str(e)}")
 
 def generate_images(num_images, model_id):
-    # Apply model-specific patches
-    if "00019" in model_id:
-        apply_pggan_patch()
-    
-    # Initialize generator
-    generators = Generators()
-    
-    # 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)
+    try:
+        generators = Generators()
+        images = []
+        
+        for i in range(num_images):
+            sample = generators.generate(
+                model_id=model_id,
+                num_samples=1,
+                install_dependencies=False
+            )
+            img = to_pil_image(sample[0]).convert("RGB")
+            images.append(img)
+        
+        cols = st.columns(4)
+        for idx, img in enumerate(images):
+            with cols[idx % 4]:
+                st.image(img, caption=f"Image {idx+1}", use_column_width=True)
                 st.markdown("---")
-            except Exception as e:
-                st.error(f"Error generating image {i+1}: {str(e)}")
+                
+    except Exception as e:
+        st.error(f"Generation failed: {str(e)}")
 
 if __name__ == "__main__":
     main()