Update app.py

app.py

@@ -6,6 +6,7 @@ from insightface.app import FaceAnalysis
 import onnxruntime
 from PIL import Image
 import tempfile
+import math # Import math for debugging if needed, but not directly used in the fixed slider logic
 
 # --- Global Configurations and Model Loading ---
 # Determine the appropriate provider for ONNX Runtime.
@@ -29,17 +30,21 @@ SIMSWAP_INPUT_SIZE = 256 # SimSwap models typically expect 256x256 input
 
 if not os.path.exists(SIMSWAP_MODEL_PATH):
     raise FileNotFoundError(f"SimSwap model not found at: {SIMSWAP_MODEL_PATH}. "
-                            "Please place 'simswap_256.onnx' in the 'models/' directory.")
+                            "Please place 'simswap_256.onnx' in the 'models/' directory relative to this script.")
 try:
     simswap_session = onnxruntime.InferenceSession(SIMSWAP_MODEL_PATH, providers=providers)
-    simswap_input_name = simswap_session.get_inputs()[0].name
+    # Get input and output names from the ONNX model
+    simswap_input_name_0 = simswap_session.get_inputs()[0].name
+    simswap_input_name_1 = simswap_session.get_inputs()[1].name
     simswap_output_name = simswap_session.get_outputs()[0].name
     print(f"SimSwap model '{SIMSWAP_MODEL_PATH}' loaded successfully.")
-    print(f"SimSwap expected input: {simswap_session.get_inputs()[0].shape}")
-    print(f"SimSwap output shape: {simswap_session.get_outputs()[0].shape}")
+    print(f"SimSwap expected input 0: {simswap_session.get_inputs()[0].shape} (name: {simswap_input_name_0})")
+    print(f"SimSwap expected input 1: {simswap_session.get_inputs()[1].shape} (name: {simswap_input_name_1})")
+    print(f"SimSwap output shape: {simswap_session.get_outputs()[0].shape} (name: {simswap_output_name})")
 except Exception as e:
     raise RuntimeError(f"Failed to load SimSwap model from {SIMSWAP_MODEL_PATH}: {e}. "
-                       "Check the model file's integrity and ONNX Runtime compatibility.")
+                       "Check the model file's integrity and ONNX Runtime compatibility. "
+                       "Also ensure it's a 2-input SimSwap model.")
 
 
 # --- Helper Functions ---
@@ -120,6 +125,7 @@ def update_faces_preview(target_img_pil: Image.Image):
     """
     if target_img_pil is None:
         # Return default states for all outputs when input is None
+        # When no image, no faces, so slider remains at 0-0 non-interactive
         return None, gr.Slider(minimum=0, maximum=0, value=0, interactive=False), gr.File(value=None, interactive=False), "Please upload a target image first."
 
     target_np_rgb = np.array(target_img_pil)
@@ -129,13 +135,17 @@ def update_faces_preview(target_img_pil: Image.Image):
     num_faces = len(faces)
 
     if num_faces == 0:
+        # If no faces are detected, set slider range to 0 to 0 and make it non-interactive.
+        # This prevents maximum < minimum, which causes math domain error in Gradio Slider.
         return None, gr.Slider(minimum=0, maximum=0, value=0, interactive=False), gr.File(value=None, interactive=False), "No faces detected in the target image. Please try another image."
+    else:
+        # If faces are detected, set the slider range appropriately
+        preview_img_bgr = draw_faces(target_np_bgr, faces)
+        preview_img_rgb = cv2.cvtColor(preview_img_bgr, cv2.COLOR_BGR2RGB)
 
-    preview_img_bgr = draw_faces(target_np_bgr, faces)
-    preview_img_rgb = cv2.cvtColor(preview_img_bgr, cv2.COLOR_BGR2RGB)
-
-    # Set the max value of the slider to (number of faces - 1) as indices are 0-based
-    return preview_img_rgb, gr.Slider(minimum=0, maximum=num_faces - 1, value=0, interactive=True), gr.File(value=None, interactive=False), f"Detected {num_faces} face(s). Select the face to swap using the slider (0-indexed)."
+        # Set the max value of the slider to (number of faces - 1) as indices are 0-based
+        # This ensures minimum=0 and maximum >= 0, avoiding the math domain error.
+        return preview_img_rgb, gr.Slider(minimum=0, maximum=num_faces - 1, value=0, interactive=True), gr.File(value=None, interactive=False), f"Detected {num_faces} face(s). Select the face to swap using the slider (0-indexed)."
 
 def face_swap_simswap(source_img_pil: Image.Image, target_img_pil: Image.Image, face_index: int):
     """
@@ -164,18 +174,13 @@ def face_swap_simswap(source_img_pil: Image.Image, target_img_pil: Image.Image,
                        "You might need to click 'Preview Detected Faces' again to update the slider range.")
     target_face_info = target_faces[face_index]
 
-    # Crop and align faces using InsightFace's bbox/kps (for robust preprocessing)
-    # InsightFace's face.embedding is typically derived from an aligned face.
-    # To get the aligned face itself, you might need to use its ArcFace models or a custom aligner.
-    # For simplicity, we'll extract the bounding box and then resize/align
-    
-    # Extract source face region
+    # Crop source face region using bounding box
     x1s, y1s, x2s, y2s = source_face_info.bbox.astype(int)
     source_face_crop = source_np_bgr[y1s:y2s, x1s:x2s]
     if source_face_crop.size == 0:
         raise gr.Error("Could not crop source face properly. Image might be too small or face too close to edge.")
 
-    # Extract target face region
+    # Crop target face region using bounding box
     x1t, y1t, x2t, y2t = target_face_info.bbox.astype(int)
     target_face_crop = target_np_bgr[y1t:y2t, x1t:x2t]
     if target_face_crop.size == 0:
@@ -187,39 +192,31 @@ def face_swap_simswap(source_img_pil: Image.Image, target_img_pil: Image.Image,
 
     # --- 3. Run SimSwap Inference ---
     try:
-        # SimSwap typically takes two inputs: source face and target face
-        # Input names might vary based on the specific ONNX export. Common names are 'src' and 'dst' or 'source' and 'target'.
-        # Let's assume input names are 'src' and 'dst' based on common SimSwap ONNX exports.
-        # If your model has different names, you'll need to inspect it using:
-        # for inp in simswap_session.get_inputs(): print(inp.name)
+        # SimSwap typically takes two inputs: source face and target face.
+        # Ensure the input names match your specific ONNX model (e.g., 'src', 'dst', 'input.1', 'input.2').
+        # The variables `simswap_input_name_0` and `simswap_input_name_1` were retrieved from the model dynamically.
         simswap_inputs = {
-            simswap_session.get_inputs()[0].name: source_input_tensor, # This is often the driving (source) identity
-            simswap_session.get_inputs()[1].name: target_input_tensor # This is often the target pose/expression
+            simswap_input_name_0: source_input_tensor, # Often the driving (source) identity
+            simswap_input_name_1: target_input_tensor  # Often the target pose/expression
         }
         
         simswap_raw_output = simswap_session.run([simswap_output_name], simswap_inputs)[0]
     except Exception as e:
-        raise gr.Error(f"SimSwap inference failed: {e}. Check model inputs/outputs and ensure images are suitable.")
+        raise gr.Error(f"SimSwap inference failed: {e}. Check model inputs/outputs and ensure images are suitable. "
+                       "Common issues include incorrect input tensor names or shapes.")
 
     # --- 4. Postprocess SimSwap Output ---
     swapped_face_simswap_output = postprocess_simswap_output(simswap_raw_output)
 
     # --- 5. Blend Swapped Face back into Original Target Image ---
     # We'll use OpenCV's seamlessClone for a natural blend.
-    # This requires:
-    # 1. The original target image (target_np_bgr)
-    # 2. The swapped face image (swapped_face_simswap_output)
-    # 3. A mask for the swapped face (often a simple ellipse or a full white mask if blending is good)
-    # 4. The center point where the swapped face should be placed
 
     # Resize swapped face output to original target face bounding box dimensions
     target_face_width = x2t - x1t
     target_face_height = y2t - y1t
     swapped_face_resized = cv2.resize(swapped_face_simswap_output, (target_face_width, target_face_height))
 
-    # Create a mask for seamless cloning
-    # A simple white rectangle covering the area is often sufficient for seamlessClone
-    # For more advanced blending, a precise face parsing mask is ideal.
+    # Create a mask for seamless cloning. A simple white rectangle covering the area is often sufficient.
     mask = np.full(swapped_face_resized.shape[:2], 255, dtype=np.uint8) # White mask
 
     # Calculate the center of the target face bounding box
@@ -228,20 +225,19 @@ def face_swap_simswap(source_img_pil: Image.Image, target_img_pil: Image.Image,
     center_point = (center_x, center_y)
 
     try:
-        # Perform seamless cloning
-        # FLAG_NORMAL_CLONE often gives the best results for face swapping
+        # Perform seamless cloning. MIXED_CLONE often works well for blending.
         final_swapped_img_bgr = cv2.seamlessClone(
             swapped_face_resized,
             target_np_bgr,
             mask,
             center_point,
-            cv2.MIXED_CLONE # or cv2.NORMAL_CLONE
+            cv2.MIXED_CLONE # or cv2.NORMAL_CLONE for sharper edges
         )
     except Exception as e:
-        # Fallback to simple paste if seamlessClone fails (e.g., due to size mismatch issues)
-        print(f"Seamless cloning failed: {e}. Attempting simple paste.")
+        # Fallback to simple paste if seamlessClone fails (e.g., due to size mismatch issues or OpenCV errors)
+        print(f"Seamless cloning failed: {e}. Attempting simple paste as a fallback.")
         final_swapped_img_bgr = target_np_bgr.copy()
-        # Simple paste (less visually appealing than seamlessClone)
+        # Simple paste (less visually appealing than seamlessClone, but avoids crash)
         final_swapped_img_bgr[y1t:y2t, x1t:x2t] = swapped_face_resized
 
 
@@ -314,7 +310,7 @@ with gr.Blocks(title="Face Swap App (SimSwap)") as demo:
     )
 
     swap_button.click(
-        fn=face_swap_simswap, # Changed to the new SimSwap function
+        fn=face_swap_simswap, # The function using SimSwap
         inputs=[source_image, target_image, face_index_slider],
         outputs=[output_image, download_output, status_message]
     )