Update app.py

app.py

@@ -4,12 +4,11 @@ import numpy as np
 import gradio as gr
 from insightface.app import FaceAnalysis
 from insightface.model_zoo import get_model # For swapper
-from PIL import Image # Import PIL directly
+from PIL import Image
 import tempfile
 import logging
 import onnxruntime
-from typing import List, Optional, Tuple, Any # For improved type hinting
-import copy # Import for deepcopy
+from typing import List, Optional, Tuple # For improved type hinting
 
 # --- Configuration & Setup ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -44,24 +43,6 @@ face_analyzer: Optional[FaceAnalysis] = None
 swapper = None
 face_restorer: Optional[onnxruntime.InferenceSession] = None
 
-# --- Create Directories and Dummy Example Files (Moved Earlier) ---
-os.makedirs("models", exist_ok=True) 
-os.makedirs("examples", exist_ok=True)
-example_files_to_create = {
-    "examples/source_example_1.jpg": ('blue', (100,100)),
-    "examples/source_example_2.png": ('green', (100,100)),
-    "examples/target_example_1.jpg": ('yellow', (200,200)),
-    "examples/target_example_2.png": ('magenta', (200,200)),
-}
-for f_path, (color, size) in example_files_to_create.items():
-    if not os.path.exists(f_path):
-        try:
-            Image.new('RGB', size, color=color).save(f_path)
-            logging.info(f"Created dummy example file: {f_path}")
-        except Exception as e:
-            logging.warning(f"Could not create dummy example file {f_path}: {e}")
-
-
 def initialize_models():
     global face_analyzer, swapper, face_restorer
     try:
@@ -117,7 +98,7 @@ def convert_cv2_to_pil(cv2_image: Optional[np.ndarray]) -> Optional[Image.Image]
         logging.error(f"Error converting CV2 to PIL: {e}")
         return None
 
-def get_faces_from_image(img_np: np.ndarray) -> List[Any]: 
+def get_faces_from_image(img_np: np.ndarray) -> List:
     if not core_models_loaded_successfully or face_analyzer is None:
         logging.error("Face analyzer not available for get_faces_from_image. Core models might have failed to load.")
         return []
@@ -131,11 +112,10 @@ def get_faces_from_image(img_np: np.ndarray) -> List[Any]:
         logging.error(f"Error during face detection: {e}", exc_info=True)
         return []
 
-def draw_detected_faces(img_np: np.ndarray, faces: List[Any]) -> np.ndarray:
+def draw_detected_faces(img_np: np.ndarray, faces: List) -> np.ndarray:
     img_with_boxes = img_np.copy()
     for i, face in enumerate(faces):
-        # Face object from insightface usually supports attribute access for bbox
-        box = face.bbox.astype(int) 
+        box = face.bbox.astype(int)
         x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
         cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (0, 255, 0), 2)
         label_position = (x1, max(0, y1 - 10))
@@ -225,124 +205,29 @@ def histogram_match_color(source_img_bgr: np.ndarray, target_img_bgr: np.ndarray
         logging.error(f"Error during color histogram matching: {e}", exc_info=True)
         return source_img_bgr
 
-def process_face_swap(source_file_objects: Optional[List[Any]], 
-                      target_pil_img: Optional[Image.Image],
+def process_face_swap(source_pil_img: Optional[Image.Image], target_pil_img: Optional[Image.Image],
                       target_face_index: int, apply_enhancement: bool, apply_color_correction: bool,
                       progress=gr.Progress(track_tqdm=True)):
     progress(0, desc="Initializing process...")
     if not core_models_loaded_successfully:
         gr.Error("CRITICAL: Core models (Face Analyzer or Swapper) not loaded. Cannot proceed.")
         return Image.new('RGB', (100, 100), color='lightgrey'), None, "Core models failed to load."
+    if source_pil_img is None: raise gr.Error("Source image not provided. Please upload a source face image.")
+    if target_pil_img is None: raise gr.Error("Target image not provided. Please upload a target scene image.")
 
-    if not source_file_objects:
-        raise gr.Error("Source image(s) not provided. Please upload one or more source face images.")
-    if target_pil_img is None:
-        raise gr.Error("Target image not provided. Please upload a target scene image.")
-
-    progress(0.05, desc="Reading source images...")
-    source_pil_list: List[Image.Image] = []
-    for file_obj in source_file_objects:
-        try:
-            img = Image.open(file_obj.name).convert('RGB')
-            source_pil_list.append(img)
-        except Exception as e:
-            logging.warning(f"Could not read or convert one of the source images ({file_obj.name}): {e}")
-            gr.Info(f"Skipping one source image due to read error: {os.path.basename(file_obj.name)}")
-    
-    if not source_pil_list:
-        raise gr.Error("No valid source images could be read. Please check the uploaded files.")
-
-    progress(0.10, desc=f"Extracting features from {len(source_pil_list)} source image(s)...")
-    valid_source_face_objects: List[Any] = [] 
-    
-    for idx, pil_img in enumerate(source_pil_list):
-        source_np_item = convert_pil_to_cv2(pil_img)
-        if source_np_item is None:
-            logging.warning(f"Could not convert source image {idx+1} to CV2 format.")
-            gr.Info(f"Skipping source image {idx+1} due to conversion error.")
-            continue
-
-        current_source_faces = get_faces_from_image(source_np_item)
-        if not current_source_faces:
-            logging.warning(f"No face found in source image {idx+1}. It will be skipped.")
-            gr.Info(f"No face found in source image {idx+1}, skipping it.")
-            continue
-        
-        detected_face = current_source_faces[0]
-        # Check for embedding attribute first, then key
-        current_embedding = None
-        if hasattr(detected_face, 'embedding') and detected_face.embedding is not None:
-            current_embedding = detected_face.embedding
-        elif 'embedding' in detected_face and detected_face['embedding'] is not None:
-            current_embedding = detected_face['embedding']
-        
-        if current_embedding is not None:
-            valid_source_face_objects.append(detected_face)
-        else:
-            logging.warning(f"Valid embedding not found for face in source image {idx+1}.")
-            gr.Info(f"Valid embedding not found in source image {idx+1}, skipping.")
-
-    if not valid_source_face_objects:
-        raise gr.Error("No usable faces (with embeddings) found in any of the uploaded source images.")
-
-    source_face_for_swap: Optional[Any] = None
-    if len(valid_source_face_objects) == 1:
-        source_face_for_swap = valid_source_face_objects[0]
-        logging.info("Using single source face for swap.")
-        gr.Info("Using features from 1 source image.")
-    else: 
-        logging.info(f"Averaging embeddings from {len(valid_source_face_objects)} source faces.")
-        gr.Info(f"Averaging features from {len(valid_source_face_objects)} source images for potentially better results.")
-        
-        embeddings_to_average = []
-        for face_obj in valid_source_face_objects:
-            emb_val = None
-            if hasattr(face_obj, 'embedding') and face_obj.embedding is not None:
-                emb_val = face_obj.embedding
-            elif 'embedding' in face_obj and face_obj['embedding'] is not None:
-                emb_val = face_obj['embedding']
-            
-            if emb_val is None: # Should have been caught, but defensive
-                logging.error("A face object in valid_source_face_objects lacked an embedding unexpectedly.")
-                raise gr.Error("Internal error: Missing embedding in a validated source face.")
-            embeddings_to_average.append(emb_val)
-
-        first_emb_shape = embeddings_to_average[0].shape
-        for i, emb in enumerate(embeddings_to_average):
-            if emb.shape != first_emb_shape:
-                logging.error(f"Embeddings have inconsistent shapes. Cannot average. Emb {i} shape: {emb.shape}, Expected: {first_emb_shape}")
-                raise gr.Error("Source face embeddings have inconsistent structures. Cannot average.")
-
-        avg_embedding = np.mean(np.array(embeddings_to_average), axis=0)
-        avg_embedding = avg_embedding / np.linalg.norm(avg_embedding)
-        
-        source_face_for_swap = copy.deepcopy(valid_source_face_objects[0]) # Deepcopy to preserve class and attributes
-        
-        # Set 'embedding' (primary)
-        if hasattr(source_face_for_swap, 'embedding'):
-            source_face_for_swap.embedding = avg_embedding
-        elif 'embedding' in source_face_for_swap:
-            source_face_for_swap['embedding'] = avg_embedding
-        else: # If somehow neither exists, create it (less likely for Face objects)
-            source_face_for_swap['embedding'] = avg_embedding
-            
-        # Set 'normed_embedding' if it exists, as this was the original error trigger
-        if hasattr(source_face_for_swap, 'normed_embedding'):
-            source_face_for_swap.normed_embedding = avg_embedding
-        elif 'normed_embedding' in source_face_for_swap:
-            source_face_for_swap['normed_embedding'] = avg_embedding
-        # If 'normed_embedding' does not exist, we assume the swapper might fall back to 'embedding'
-        # or that the original Face object structure usually has both as the same if already normalized.
-    
-    if source_face_for_swap is None: 
-        raise gr.Error("Failed to prepare a source face for swapping.")
-
-    progress(0.20, desc="Converting target image...")
+    progress(0.05, desc="Converting images to OpenCV format...")
+    source_np = convert_pil_to_cv2(source_pil_img)
     target_np = convert_pil_to_cv2(target_pil_img)
-    if target_np is None:
-        raise gr.Error("Target image conversion failed. Ensure the image is valid and not corrupted.")
+    if source_np is None or target_np is None:
+        raise gr.Error("Image conversion failed. Ensure images are valid (e.g., JPG, PNG) and not corrupted.")
     target_h, target_w = target_np.shape[:2]
 
+    progress(0.15, desc="Detecting face in source image...")
+    source_faces = get_faces_from_image(source_np)
+    if not source_faces:
+        raise gr.Error("No face found in the source image. Please use a clear, well-lit image of a single, front-facing face.")
+    source_face = source_faces[0]
+
     progress(0.25, desc="Detecting faces in target image...")
     target_faces = get_faces_from_image(target_np)
     if not target_faces:
@@ -351,15 +236,12 @@ def process_face_swap(source_file_objects: Optional[List[Any]],
         err_msg = (f"Selected target face index ({target_face_index}) is out of range. "
                    f"Detected {len(target_faces)} faces (indices 0 to {len(target_faces)-1}). Please re-preview and select a valid face.")
         raise gr.Error(err_msg)
-    target_face_to_swap_info = target_faces[int(target_face_index)] 
-    
+    target_face_to_swap_info = target_faces[int(target_face_index)]
     swapped_bgr_img = target_np.copy()
     try:
         progress(0.4, desc="Performing face swap...")
-        swapped_bgr_img = swapper.get(swapped_bgr_img, target_face_to_swap_info, source_face_for_swap, paste_back=True)
-        
-        # target_face_to_swap_info is directly from insightface, .bbox should be safe
-        bbox_coords = target_face_to_swap_info.bbox.astype(int) 
+        swapped_bgr_img = swapper.get(swapped_bgr_img, target_face_to_swap_info, source_face, paste_back=True)
+        bbox_coords = target_face_to_swap_info.bbox.astype(int)
         
         if apply_enhancement:
             if restoration_model_loaded_successfully and face_restorer is not None:
@@ -403,6 +285,9 @@ def process_face_swap(source_file_objects: Optional[List[Any]],
                 logging.warning("Skipping color correction due to invalid bounding box for region extraction.")
     except Exception as e:
         logging.error(f"Error during face swapping or post-processing: {e}", exc_info=True)
+        swapped_pil_img_on_error = convert_cv2_to_pil(swapped_bgr_img)
+        if swapped_pil_img_on_error is None:
+            swapped_pil_img_on_error = Image.new('RGB', (target_w, target_h), color='lightgrey')
         raise gr.Error(f"An error occurred during processing: {str(e)}. Check logs for details.")
 
     progress(0.9, desc="Finalizing image conversion...")
@@ -436,7 +321,7 @@ def preview_target_faces(target_pil_img: Optional[Image.Image]):
     if target_np is None:
         gr.Warning("Could not process target image for preview. Please try a different image.")
         return blank_preview_placeholder, reset_slider
-    faces = get_faces_from_image(target_np) 
+    faces = get_faces_from_image(target_np)
     if not faces:
         max_h, max_w = DETECTION_SIZE[1], DETECTION_SIZE[0]
         h, w = target_np.shape[:2]
@@ -449,8 +334,7 @@ def preview_target_faces(target_pil_img: Optional[Image.Image]):
         if preview_pil_img is None: preview_pil_img = blank_preview_placeholder
         gr.Info("No faces were detected in the target image.")
         return preview_pil_img, reset_slider
-    
-    preview_np_with_boxes = draw_detected_faces(target_np, faces) 
+    preview_np_with_boxes = draw_detected_faces(target_np, faces)
     preview_pil_img_with_boxes = convert_cv2_to_pil(preview_np_with_boxes)
     if preview_pil_img_with_boxes is None:
         preview_pil_img_with_boxes = blank_preview_placeholder
@@ -461,13 +345,13 @@ def preview_target_faces(target_pil_img: Optional[Image.Image]):
     return preview_pil_img_with_boxes, slider_update
 
 # --- Gradio UI Definition ---
-with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.5 Multi-Source", theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky)) as demo: #Version bump
+with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.2 Lite", theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky)) as demo:
     gr.Markdown(
         """
         <div style="text-align: center; border-bottom: 1px solid #eee; padding-bottom:10px;">
-            <h1>💎 Ultimate Face Swap AI <span style="font-size:0.8em; color: #555;">v3.5 Multi-Source</span> ✨</h1>
-            <p>Upload <strong>one or more source face images</strong>, a <strong>target image</strong>, and let the AI work its magic!</p>
-            <p style="font-size:0.9em;">If multiple source images are provided, their facial features will be averaged for the swap.</p>
+            <h1>💎 Ultimate Face Swap AI <span style="font-size:0.8em; color: #555;">v3.2 Lite</span> ✨</h1>
+            <p>Upload a <strong>source face</strong>, a <strong>target image</strong>, and let the AI work its magic!</p>
+            <p style="font-size:0.9em;">Optionally enhance with <strong>face restoration</strong> and <strong>color correction</strong> for improved realism.</p>
         </div>
         """
     )
@@ -476,15 +360,13 @@ with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.5 Multi-Source", theme=gr.th
         gr.Error(
             "🔴 CRITICAL ERROR: Core models (Face Analyzer or Swapper) failed to load. "
             "The application UI will load but WILL NOT FUNCTION correctly. "
-            "Please check the console logs for details and restart after fixing."
+            "Please check the console logs for details (e.g., model file paths, ONNX Runtime issues) "
+            "and restart the application after fixing the underlying problem."
         )
 
     with gr.Row():
         with gr.Column(scale=1, min_width=300):
-            source_image_input = gr.Files(
-                label="👤 Source Face Images (Upload one or more clear face images)",
-                file_types=["image"],
-            )
+            source_image_input = gr.Image(label="👤 Source Face Image (Clear, single face recommended)", type="pil", sources=["upload", "clipboard"], height=380)
         with gr.Column(scale=1, min_width=300):
             target_image_input = gr.Image(label="🖼️ Target Scene Image (Upload image to swap a face into)", type="pil", sources=["upload", "clipboard"], height=380)
 
@@ -501,6 +383,7 @@ with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.5 Multi-Source", theme=gr.th
             )
             gr.Markdown("<p style='font-size:0.8em; color:#555;'>After uploading target, click 'Preview'. Then use slider to select face.</p>")
             
+            # Changed from gr.Markdown("---", style="...") to gr.HTML for wider compatibility
             gr.HTML("<hr style='margin-top: 5px; margin-bottom: 5px; border-style: solid; border-width: thin; border-color: #ddd;'>")
             
             enhance_checkbox_label = "✨ Apply Selective Face Restoration"
@@ -512,9 +395,9 @@ with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.5 Multi-Source", theme=gr.th
                 interactive=restoration_model_loaded_successfully
             )
             if not restoration_model_loaded_successfully:
-                gr.Markdown("<p style='color: orange; font-size:0.8em;'>⚠️ Face restoration model not loaded. This feature is disabled.</p>")
+                gr.Markdown("<p style='color: orange; font-size:0.8em;'>⚠️ Face restoration model not loaded. This feature is disabled. Check console logs for details on `RESTORATION_MODEL_PATH`.</p>")
             else:
-                gr.Markdown("<p style='font-size:0.8em; color:#555;'>Improves quality of the swapped face region.</p>")
+                gr.Markdown("<p style='font-size:0.8em; color:#555;'>Improves quality of the swapped face region. Requires restoration model.</p>")
 
             color_correction_checkbox = gr.Checkbox(
                 label="🎨 Apply Color Correction",
@@ -585,57 +468,29 @@ with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.5 Multi-Source", theme=gr.th
     gr.HTML("<hr style='margin-top: 20px; margin-bottom: 10px;'>")
     gr.Markdown("### 📋 Example Usage")
     
-    def process_face_swap_for_examples(source_file_paths_list: List[str],
-                                     target_img_path: str, idx: int, enhance: bool, color_correct: bool,
-                                     progress=gr.Progress(track_tqdm=True)):
-        mock_source_file_objects = []
-        if source_file_paths_list: 
-            for path in source_file_paths_list:
-                if not os.path.exists(path): 
-                    logging.error(f"Example source file not found: {path}. Skipping this example or part of it.")
-                    continue 
-                mock_file_obj = type('MockFile', (object,), {'name': path})()
-                mock_source_file_objects.append(mock_file_obj)
-        
-        target_pil = None
-        if target_img_path and os.path.exists(target_img_path): 
-            try:
-                target_pil = Image.open(target_img_path)
-            except Exception as e:
-                logging.error(f"Failed to load example target image {target_img_path}: {e}")
-                return Image.new('RGB', (100,100), 'red'), None, f"Error loading example target: {os.path.basename(target_img_path)}"
-        elif target_img_path and not os.path.exists(target_img_path):
-            logging.error(f"Example target file not found: {target_img_path}")
-            return Image.new('RGB', (100,100), 'red'), None, f"Example target missing: {os.path.basename(target_img_path)}"
-
-        if not mock_source_file_objects and not target_pil: # Both sources and target failed or were not provided
-             return Image.new('RGB', (100,100), 'red'), None, "Error loading example images."
-        if not mock_source_file_objects and target_pil : # Sources failed/missing but target is ok
-             # This case means process_face_swap will raise an error about no source images.
-             # Let's try to call it anyway and let its internal error handling manage the gr.Error.
-             # Or, we can return a specific message here.
-             # For consistency, let process_face_swap handle the "no source" error.
-             pass
-
-
-        return process_face_swap(mock_source_file_objects, target_pil, idx, enhance, color_correct, progress)
+    def process_face_swap_for_examples(src_img, tgt_img, idx, enhance, color_correct, progress=gr.Progress(track_tqdm=True)):
+        img, file_path, status = process_face_swap(src_img, tgt_img, idx, enhance, color_correct, progress)
+        return img, file_path, status
 
     gr.Examples(
         examples=[
-            [["examples/source_example_1.jpg"], "examples/target_example_1.jpg", 0, True, True],
-            [["examples/source_example_1.jpg", "examples/source_example_2.png"], "examples/target_example_2.png", 1, True, False],
-            [["examples/source_example_1.jpg"], "examples/target_example_1.jpg", 0, False, True],
+            ["examples/source_example_1.jpg", "examples/target_example_1.jpg", 0, True, True],
+            ["examples/source_example_2.png", "examples/target_example_2.png", 1, True, False],
+            ["examples/source_example_1.jpg", "examples/target_example_1.jpg", 0, False, True],
         ],
         inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox, color_correction_checkbox],
         outputs=[swapped_image_output, download_output_file, status_message_output],
         fn=process_face_swap_for_examples, 
         cache_examples=False,
-        label="Example Face Swaps (Ensure example images exist in 'examples' folder)"
+        label="Example Face Swaps (Click to run - ensure example images exist in 'examples' folder)"
     )
 
 if __name__ == "__main__":
+    os.makedirs("models", exist_ok=True)
+    os.makedirs("examples", exist_ok=True)
+
     print("\n" + "="*70)
-    print("🚀 ULTIMATE FACE SWAP AI - v3.5 Multi-Source STARTUP STATUS 🚀")
+    print("🚀 ULTIMATE FACE SWAP AI - v3.2 Lite STARTUP STATUS 🚀") # Version bump
     print("="*70)
     print(f"Execution Providers Selected: {EXECUTION_PROVIDERS}")
     
@@ -656,5 +511,5 @@ if __name__ == "__main__":
         print("\n👉 The Gradio interface will launch, but swapping functionality will be broken. Please address errors.")
     else:
         print("\n👉 Launching Gradio Interface... Access it in your browser (usually at http://127.0.0.1:7860).")
-        
+    
     demo.launch()