Update app.py

app.py

@@ -8,12 +8,12 @@ import os
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Model parameters (must match your training)
+# Model parameters
 nz = 100
 ngf = 64
 num_classes = 10
 
-# Generator class (same as your training script)
+# Generator class
 class Generator(nn.Module):
     def __init__(self):
         super(Generator, self).__init__()
@@ -47,21 +47,19 @@ class Generator(nn.Module):
         output = self.resize(output)
         return output
 
-# Load the trained model
-@st.cache_resource
+# Load model function (NO @st.cache_resource decorator!)
 def load_model():
     generator = Generator().to(device)
     
-    # Load the saved model
-    if os.path.exists('mnist_gan_model.pth'):
+    try:
         checkpoint = torch.load('mnist_gan_model.pth', map_location=device)
         generator.load_state_dict(checkpoint['generator_state_dict'])
         generator.eval()
-        print("Model loaded successfully!")
-    else:
-        print("Warning: Model file not found!")
-    
-    return generator
+        print("✅ Model loaded successfully!")
+        return generator
+    except Exception as e:
+        print(f"❌ Error loading model: {e}")
+        return None
 
 # Initialize generator
 generator = load_model()
@@ -70,107 +68,75 @@ generator = load_model()
 def generate_digit_images(digit):
     """Generate 5 images of the specified digit"""
     
+    if generator is None:
+        return [Image.new('L', (112, 112), 128)] * 5
+    
     digit = int(digit)
     num_images = 5
     
     with torch.no_grad():
-        # Generate random noise
         noise = torch.randn(num_images, nz, 1, 1).to(device)
         labels = torch.full((num_images,), digit, dtype=torch.long).to(device)
         
-        # Generate images
         generated_images = generator(noise, labels)
         
-        # Convert to numpy and denormalize
         images = generated_images.cpu().numpy()
-        images = (images + 1) / 2.0  # Denormalize from [-1, 1] to [0, 1]
-        images = np.squeeze(images)  # Remove channel dimension
+        images = (images + 1) / 2.0
+        images = np.squeeze(images)
         
-        # Convert to PIL Images for Gradio
         pil_images = []
         for img in images:
-            # Convert to 0-255 range and uint8
             img_uint8 = (img * 255).astype(np.uint8)
             pil_img = Image.fromarray(img_uint8, mode='L')
-            # Resize for better visibility
-            pil_img = pil_img.resize((112, 112), Image.NEAREST)  # 4x upscale
+            pil_img = pil_img.resize((112, 112), Image.NEAREST)
             pil_images.append(pil_img)
         
         return pil_images
 
-# Create Gradio interface
+# Gradio interface
 def create_app():
-    with gr.Blocks(
-        title="Handwritten Digit Generator", 
-        theme=gr.themes.Soft(),
-        css=".gradio-container {max-width: 700px; margin: auto;}"
-    ) as app:
+    with gr.Blocks(title="Handwritten Digit Generator", theme=gr.themes.Soft()) as app:
         
         gr.Markdown("# 🔢 Handwritten Digit Generator")
-        gr.Markdown("Generate synthetic MNIST-like digit images using a trained GAN model. Select a digit (0-9) to generate 5 unique images.")
+        gr.Markdown("Generate synthetic MNIST-like digit images using a trained GAN model.")
         
         with gr.Row():
             with gr.Column(scale=1):
                 digit_input = gr.Dropdown(
                     choices=list(range(10)), 
                     value=2, 
-                    label="Choose a digit to generate (0-9)",
-                    interactive=True
-                )
-                generate_btn = gr.Button(
-                    "🎨 Generate Images", 
-                    variant="primary",
-                    size="lg"
+                    label="Choose a digit (0-9)"
                 )
+                generate_btn = gr.Button("🎨 Generate Images", variant="primary")
             
             with gr.Column(scale=2):
                 gr.Markdown("### Generated Images")
-                
-                # Gallery to display 5 images
                 image_gallery = gr.Gallery(
                     label="Generated Digit Images", 
                     show_label=False,
                     columns=5,
                     rows=1,
-                    height=200,
-                    object_fit="contain"
+                    height=200
                 )
         
-        # Example section
-        gr.Markdown("---")
-        gr.Markdown("### How it works")
-        gr.Markdown("""
-        1. **Select** a digit from the dropdown (0-9)
-        2. **Click** 'Generate Images' button  
-        3. **View** 5 unique generated images of your chosen digit
-        4. Each generation produces different variations of the same digit
-        """)
-        
-        # Connect button to generation function
         generate_btn.click(
             fn=generate_digit_images,
             inputs=[digit_input],
             outputs=[image_gallery]
         )
         
-        # Auto-generate on page load
+        # Auto-generate on load
         app.load(
-            fn=generate_digit_images,
-            inputs=[gr.Number(value=2, visible=False)],
+            fn=lambda: generate_digit_images(2),
             outputs=[image_gallery]
         )
         
-        # Footer
         gr.Markdown("---")
-        gr.Markdown("**🤖 Model**: Conditional GAN trained on MNIST | **⚡ Framework**: PyTorch + Gradio")
+        gr.Markdown("**🤖 Model**: Conditional GAN | **⚡ Framework**: PyTorch + Gradio")
     
     return app
 
-# Launch the app
+# Launch
 if __name__ == "__main__":
     app = create_app()
-    app.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False  # Set to False for Hugging Face deployment
-    )
+    app.launch()