Update app.py

app.py

@@ -1,40 +1,53 @@
 import gradio as gr
 import torch
 import torch.nn as nn
+import os
 from PIL import Image
 from torchvision.transforms import ToTensor, ToPILImage
 
-# Define the EDSR model architecture (simplified version)
+# Define the EDSR model architecture
 class EDSR(nn.Module):
     def __init__(self):
         super(EDSR, self).__init__()
-        # Simplified architecture - in practice you'd want the full EDSR
-        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
-        self.conv2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
-        self.conv3 = nn.Conv2d(64, 3, kernel_size=3, padding=1)
+        # Basic EDSR architecture
+        self.conv1 = nn.Conv2d(3, 256, kernel_size=3, padding=1)
+        self.resblocks = nn.Sequential(*[nn.Sequential(
+            nn.Conv2d(256, 256, kernel_size=3, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(256, 256, kernel_size=3, padding=1)
+        ) for _ in range(8)])
+        self.conv2 = nn.Conv2d(256, 3, kernel_size=3, padding=1)
         
     def forward(self, x):
         x = torch.relu(self.conv1(x))
-        x = torch.relu(self.conv2(x))
-        x = self.conv3(x)
+        residual = x
+        x = self.resblocks(x)
+        x += residual
+        x = self.conv2(x)
         return x
 
-# Load the model (we'll use a pretrained version from Hugging Face)
+# Initialize model
 model = EDSR()
 
-# Load pretrained weights (alternative approach)
+# Try loading pretrained weights from Hugging Face
 try:
     state_dict = torch.hub.load_state_dict_from_url(
         "https://huggingface.co/eugenesiow/edsr-base/resolve/main/pytorch_model.bin",
-        map_location="cpu"
+        map_location="cpu",
+        file_name="edsr_weights.pth"
     )
     model.load_state_dict(state_dict)
-except:
-    print("Couldn't load pretrained weights, using random initialization")
+    print("Successfully loaded pretrained weights")
+except Exception as e:
+    print(f"Couldn't load pretrained weights: {str(e)}")
+    print("Using randomly initialized model")
 
 model.eval()
 
 def enhance_image(input_img):
+    # Resize input to prevent memory issues
+    input_img = input_img.resize((256, 256))
+    
     # Convert to tensor and add batch dimension
     input_tensor = ToTensor()(input_img).unsqueeze(0)
     
@@ -46,13 +59,19 @@ def enhance_image(input_img):
     output_img = ToPILImage()(output_tensor.squeeze(0).clamp(0, 1))
     return output_img
 
+# Prepare examples
+example_images = []
+if os.path.exists("example_image.jpg"):
+    example_images = ["example_image.jpg"]
+
 # Gradio UI
 demo = gr.Interface(
     fn=enhance_image,
-    inputs=gr.Image(type="pil", label="Upload Image"),
+    inputs=gr.Image(type="pil", label="Input Image"),
     outputs=gr.Image(type="pil", label="Enhanced Image"),
-    title="Image Super-Resolution (EDSR)",
-    examples=["example_image.jpg"] if os.path.exists("example_image.jpg") else None,
+    title="EDSR Image Super-Resolution",
+    examples=example_images,
+    allow_flagging="never"
 )
 
 demo.launch()