Upload app.py

app.py

@@ -4,49 +4,47 @@ import torch
 from RealESRGAN import RealESRGAN
 from io import BytesIO
 
-# Define the target size for the image (used for initial resizing before enhancement)
-TARGET_SIZE = (240, 240)
-
 # Function to load the model based on scale and anime toggle
 def load_model(scale, anime=False):
-    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = RealESRGAN(device, scale=scale, anime=anime)
-    model_path = {
-        (2, False): 'model/RealESRGAN_x2.pth',
-        (4, False): 'model/RealESRGAN_x4plus.pth',
-        (8, False): 'model/RealESRGAN_x8.pth',
-        (4, True): 'model/RealESRGAN_x4plus_anime_6B.pth'
-    }[(scale, anime)]
-    model.load_weights(model_path)
-    return model
+    try:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        model = RealESRGAN(device, scale=scale, anime=anime)
+        model_path = {
+            (2, False): 'model/RealESRGAN_x2.pth',
+            (4, False): 'model/RealESRGAN_x4plus.pth',
+            (8, False): 'model/RealESRGAN_x8.pth',
+            (4, True): 'model/RealESRGAN_x4plus_anime_6B.pth'
+        }[(scale, anime)]
+        model.load_weights(model_path)
+        return model
+    except Exception as e:
+        st.error(f"Error loading the model: {e}")
+        return None
 
 def enhance_image(image, scale, anime):
+    model = load_model(scale, anime=anime)
+    if model is None:
+        return None, None
+
     try:
-        model = load_model(scale, anime=anime)
-        
         # Convert image to RGB if it has an alpha channel
         if image.mode != 'RGB':
             image = image.convert('RGB')
         
-        # Store original image size
-        original_size = image.size
-        
-        # Resize image to target dimensions for processing
-        image = image.resize(TARGET_SIZE)
-        
-        # Perform image enhancement
+        # Process the image with the model
         sr_image = model.predict(image)
         
-        # Resize the enhanced image back to the original size
-        sr_image = sr_image.resize(original_size)
+        # Ensure the enhanced image has the same size as the original
+        sr_image = sr_image.resize(image.size)
         
+        # Save enhanced image to buffer
         buffer = BytesIO()
         sr_image.save(buffer, format="PNG")
         buffer.seek(0)
-        
-        return sr_image, buffer, None
+        return sr_image, buffer
     except Exception as e:
-        return None, None, str(e)
+        st.error(f"Error enhancing the image: {e}")
+        return None, None
 
 def main():
     st.title("Generative AI Image Restoration")
@@ -70,11 +68,9 @@ def main():
         
         # Enhance button
         if st.button("Restore Image"):
-            enhanced_image, buffer, error_message = enhance_image(image, scale_value, anime)
+            enhanced_image, buffer = enhance_image(image, scale_value, anime)
             
-            if error_message:
-                st.error(f"An error occurred: {error_message}")
-            else:
+            if enhanced_image:
                 # Show images side by side
                 col1, col2 = st.columns(2)
                 with col1: