Upload app.py with huggingface_hub

app.py

@@ -1,46 +1,68 @@
 import spaces
 import gradio as gr
 import cv2
-import traceback
+import numpy as np
+
+
+def _to_bgr(img):
+    if img is None:
+        raise ValueError("no image received")
+    if isinstance(img, str):
+        im = cv2.imread(img)
+        if im is None:
+            raise ValueError("unreadable path: %s" % img)
+        return im
+    if hasattr(img, "convert"):
+        return cv2.cvtColor(np.array(img.convert("RGB")), cv2.COLOR_RGB2BGR)
+    a = np.array(img)
+    if a.ndim == 3 and a.shape[2] >= 3:
+        return cv2.cvtColor(a[:, :, :3], cv2.COLOR_RGB2BGR)
+    return a
+
+
 @spaces.GPU
 def swap_face(source_img, target_img):
+    import traceback
     import insightface
     from insightface.app import FaceAnalysis
     try:
-        app = FaceAnalysis(name='buffalo_l')
+        src = _to_bgr(source_img)
+        tgt = _to_bgr(target_img)
+        app = FaceAnalysis(name="buffalo_l")
         try:
-            app.prepare(ctx_id=0, det_size=(640,640))
-        except:
-            app.prepare(ctx_id=-1, det_size=(640,640))
+            app.prepare(ctx_id=0, det_size=(640, 640))
+        except Exception:
+            app.prepare(ctx_id=-1, det_size=(640, 640))
         import insightface.model_zoo
-        swapper = insightface.model_zoo.get_model('inswapper_128.onnx', download=True)
-        src = cv2.imread(source_img)
-        tgt = cv2.imread(target_img)
+        swapper = insightface.model_zoo.get_model("inswapper_128.onnx", download=True)
         s_faces = app.get(src)
         t_faces = app.get(tgt)
-        if len(s_faces)==0 or len(t_faces)==0:
-            app2 = FaceAnalysis(name='buffalo_l')
-            app2.prepare(ctx_id=-1, det_size=(1024,1024))
-            if len(s_faces)==0: s_faces = app2.get(src)
-            if len(t_faces)==0: t_faces = app2.get(tgt)
-        if len(s_faces)==0 or len(t_faces)==0:
-            cv2.putText(tgt, f"s:{len(s_faces)} t:{len(t_faces)}", (20,40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255),2)
-            out="/tmp/debug.jpg"
-            cv2.imwrite(out,tgt)
-            return out
+        if len(s_faces) == 0 or len(t_faces) == 0:
+            app2 = FaceAnalysis(name="buffalo_l")
+            app2.prepare(ctx_id=-1, det_size=(1024, 1024))
+            if len(s_faces) == 0:
+                s_faces = app2.get(src)
+            if len(t_faces) == 0:
+                t_faces = app2.get(tgt)
+        if len(s_faces) == 0 or len(t_faces) == 0:
+            raise gr.Error("no faces detected s:%d t:%d" % (len(s_faces), len(t_faces)))
         res = swapper.get(tgt, t_faces[0], s_faces[0], paste_back=True)
-        out="/tmp/swapped.jpg"
-        cv2.imwrite(out,res)
+        out = "/tmp/swapped.jpg"
+        cv2.imwrite(out, res)
         return out
-    except Exception as e:
-        import traceback
-        raise gr.Error(traceback.format_exc())
+    except gr.Error:
+        raise
+    except Exception:
+        raise gr.Error(traceback.format_exc()[-1500:])
+
+
 with gr.Blocks() as demo:
     gr.Markdown("# FaceFusion Zero")
     with gr.Row():
-        src = gr.Image(type="filepath", label="Source")
-        tgt = gr.Image(type="filepath", label="Target")
-    out = gr.Image(type="filepath", label="Result")
+        src = gr.Image(label="Source")
+        tgt = gr.Image(label="Target")
+    out = gr.Image(label="Result")
     btn = gr.Button("Swap")
     btn.click(swap_face, inputs=[src, tgt], outputs=out)
+
 demo.launch()