Update app.py

app.py

@@ -17,6 +17,8 @@ model_id = "Arrcttacsrks/netrunner-exe_Insight-Swap-models-onnx"
 model_file = hf_hub_download(repo_id=model_id, filename="simswap_512_unoff.onnx", token=huggingface_token)
 
 def load_and_preprocess_image(image):
+    if image is None or not isinstance(image, np.ndarray):
+        raise ValueError("Input image is not valid.")
     img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
     img = cv2.resize(img, (512, 512))
     img = img / 255.0  # Normalize to [0, 1]
@@ -26,24 +28,37 @@ def swap_faces(source_image, target_image):
     # Load the ONNX model
     session = ort.InferenceSession(model_file)
 
+    # Print model input names for debugging
+    input_names = [input.name for input in session.get_inputs()]
+    print("Model input names:", input_names)
+
     # Preprocess the images
     source_img = load_and_preprocess_image(source_image)
     target_img = load_and_preprocess_image(target_image)
 
     # Prepare input data for the model
     input_data = np.array([source_img, target_img]).transpose(0, 3, 1, 2)
-    
-    # Run inference
-    result = session.run(None, {'input': input_data})[0]
+
+    # Validate input shape
+    expected_shape = (1, 3, 512, 512)
+    if input_data.shape != expected_shape:
+        raise ValueError(f"Input data shape {input_data.shape} does not match expected shape {expected_shape}.")
+
+    # Run inference using the correct input name
+    result = session.run(None, {input_names[0]: input_data})[0]  # Use the first input name
 
     # Post-process the result
     return np.clip(result * 255, 0, 255).astype(np.uint8)
 
 # Create Gradio interface
-interface = gr.Interface(fn=swap_faces,
-                         inputs=["image", "image"],
-                         outputs="image",
-                         title="Face Swap using SimSwap",
-                         description="Upload source and target images to swap faces.")
-# Launch the interface
-interface.launch()
+interface = gr.Interface(
+    fn=swap_faces,
+    inputs=["image", "image"],
+    outputs="image",
+    title="Face Swap using SimSwap",
+    description="Upload source and target images to swap faces.",
+    allow_flagging="never"  # Prevent flagging to keep the interface clean
+)
+
+# Launch the interface with share=True for public access
+interface.launch(share=True)