Update app.py

app.py

@@ -5,63 +5,89 @@ from PIL import Image
 import cv2
 import os
 import requests
+import hashlib
 
-# Download the model if missing
+# Configuration
 MODEL_URL = "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.onnx"
 MODEL_PATH = "realesr.onnx"
+EXPECTED_MD5 = "8a628e89b1e4d9f5f174a3e8c0c7b3b1"  # MD5 hash of correct model file
 
-if not os.path.exists(MODEL_PATH):
+def verify_file(file_path, expected_md5):
+    """Verify file integrity using MD5 hash"""
+    if not os.path.exists(file_path):
+        return False
+    with open(file_path, "rb") as f:
+        file_hash = hashlib.md5(f.read()).hexdigest()
+    return file_hash == expected_md5
+
+def download_model():
+    """Download model with verification"""
     print("Downloading model...")
     try:
-        response = requests.get(MODEL_URL)
+        response = requests.get(MODEL_URL, stream=True)
+        response.raise_for_status()
+        
+        # Save in chunks to handle large files
         with open(MODEL_PATH, "wb") as f:
-            f.write(response.content)
-        print("Model downloaded successfully!")
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        
+        # Verify download
+        if not verify_file(MODEL_PATH, EXPECTED_MD5):
+            os.remove(MODEL_PATH)
+            raise ValueError("Downloaded file is corrupted")
+            
+        print("Model downloaded and verified successfully!")
     except Exception as e:
-        raise gr.Error(f"Failed to download model: {str(e)}")
+        raise gr.Error(f"Model download failed: {str(e)}")
+
+# Download model if missing or corrupted
+if not verify_file(MODEL_PATH, EXPECTED_MD5):
+    download_model()
 
-# Initialize ONNX Runtime
+# Initialize ONNX Runtime (CPU only for free tier)
 try:
     ort_session = ort.InferenceSession(
         MODEL_PATH,
-        providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
+        providers=['CPUExecutionProvider']  # Only use CPU on free tier
     )
     print("Model loaded successfully!")
 except Exception as e:
-    raise gr.Error(f"Model loading failed: {str(e)}")
+    raise gr.Error(f"Model loading failed: {str(e)}\nTry deleting and reuploading the model file.")
 
 def enhance(image):
-    # Convert to numpy array
-    img = np.array(image)
-    
-    # Resize if too large (free tier has limited memory)
-    max_size = 512
-    if max(img.shape) > max_size:
-        scale = max_size / max(img.shape)
-        img = cv2.resize(img, (0,0), fx=scale, fy=scale)
-    
-    # Convert to model input format
-    img = img.astype(np.float32) / 255.0
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
-    img = np.transpose(img, (2, 0, 1))
-    img = np.expand_dims(img, axis=0)
-    
-    # Run enhancement
-    output = ort_session.run(None, {'input': img})[0]
-    
-    # Convert back to PIL Image
-    output = output.squeeze().transpose(1, 2, 0)
-    output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
-    output = (output * 255).clip(0, 255).astype(np.uint8)
-    return Image.fromarray(output)
+    """Image enhancement function"""
+    try:
+        # Convert and resize
+        img = np.array(image)
+        if max(img.shape) > 512:  # Free tier memory limit
+            scale = 512 / max(img.shape)
+            img = cv2.resize(img, (0,0), fx=scale, fy=scale)
+        
+        # Preprocess
+        img = img.astype(np.float32) / 255.0
+        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+        img = np.transpose(img, (2, 0, 1))
+        img = np.expand_dims(img, axis=0)
+        
+        # Inference
+        output = ort_session.run(None, {'input': img})[0]
+        
+        # Postprocess
+        output = output.squeeze().transpose(1, 2, 0)
+        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        output = (output * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(output)
+    except Exception as e:
+        raise gr.Error(f"Enhancement failed: {str(e)}")
 
-# Create Gradio interface
+# Create interface
 demo = gr.Interface(
     fn=enhance,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Image(type="pil", label="Enhanced"),
-    title="Professional Image Enhancement",
-    examples=["example.jpg"] if os.path.exists("example.jpg") else None
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Image(type="pil"),
+    title="Image Enhancement",
+    allow_flagging="never"
 )
 
 demo.launch()