Update app.py

app.py

@@ -7,6 +7,7 @@ import sys
 sys.path.append('CodeFormer')
 import os
 import cv2
+import numpy as np
 import torch
 import torch.nn.functional as F
 import uuid, threading, time, glob
@@ -62,10 +63,11 @@ torch.hub.download_url_to_file(
     'https://raw.githubusercontent.com/sczhou/CodeFormer/master/inputs/cropped_faces/0729.png',
     '06.png')
 
-def imread(img_path):
-    img = cv2.imread(img_path)
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    return img
+
+def imread_unicode_safe(path):
+    with open(path, "rb") as f:
+        data = np.frombuffer(f.read(), dtype=np.uint8)
+    return cv2.imdecode(data, cv2.IMREAD_COLOR)
 
 def delayed_remove(path, delay=60):
     time.sleep(delay)
@@ -124,7 +126,6 @@ def inference(image, face_align, background_enhance, face_upsample, upscale, cod
         draw_box = False
         detection_model = "retinaface_resnet50"
 
-        print('Inp:', image.shape, background_enhance, face_upsample, upscale, codeformer_fidelity)
         face_align = face_align if face_align is not None else True
         background_enhance = background_enhance if background_enhance is not None else True
         face_upsample = face_upsample if face_upsample is not None else True
@@ -132,13 +133,26 @@ def inference(image, face_align, background_enhance, face_upsample, upscale, cod
 
         has_aligned = not face_align
         upscale = 1 if has_aligned else upscale
+        
+        if isinstance(image, dict):
+            image_path = image.get("name")
+        elif isinstance(image, str):
+            image_path = image
+        else:
+            image_path = None
+            raise gr.Error("Invalid input image.")
+        
+        if not os.path.exists(image_path):
+            raise gr.Error("Invalid input image.")
 
-        img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+        print('Inp:', image_path, background_enhance, face_upsample, upscale, codeformer_fidelity)
         
+        img = imread_unicode_safe(image_path)
+
         if img is None:
-            raise ValueError("Input image is None")
-    
-        # print('\timage size:', img.shape)
+            raise gr.Error("Failed to read input image.")
+        
+        print('\timage size:', img.shape)
 
         upscale = int(upscale) # convert type to int
         if upscale > 4: # avoid memory exceeded due to too large upscale
@@ -184,12 +198,12 @@ def inference(image, face_align, background_enhance, face_upsample, upscale, cod
             # align and warp each face
             face_helper.align_warp_face()
 
-        if min(img.shape[:2]) > 1000 and num_det_faces > 15:
-            raise gr.Error(
-                "Too many faces detected (>15) in a high-resolution image. "
-                "To keep the demo responsive and avoid long queue times, this case is skipped. "
-                "For such inputs, please deploy this demo locally and remove this limit."
-            )
+            if min(img.shape[:2]) > 1000 and num_det_faces > 15:
+                raise gr.Error(
+                    "Too many faces detected (>15) in a high-resolution image. "
+                    "To keep the demo responsive and avoid long queue times, this case is skipped. "
+                    "For such inputs, please deploy this demo locally and remove this limit."
+                )
 
             
         # face restoration for each cropped face
@@ -257,11 +271,14 @@ def inference(image, face_align, background_enhance, face_upsample, upscale, cod
         ).start()
     
         return save_path, None
+
+    except gr.Error:
+        raise
         
     except Exception as error:
-        print('Global exception', error)
-        return None, None
-
+        print('[UNEXPECTED ERROR]', error)
+        raise gr.Error("Unexpected error. Please try another image.")
+        
 
 title = "CodeFormer: Robust Face Restoration and Enhancement Network"
 
@@ -324,7 +341,7 @@ td {
 
 demo = gr.Interface(
     inference, [
-        gr.Image(type="numpy", label="Input"),
+        gr.Image(type="filepath", label="Input"),
         gr.Checkbox(value=True, label="Pre_Face_Align"),
         gr.Checkbox(value=True, label="Background_Enhance"),
         gr.Checkbox(value=True, label="Face_Upsample"),