Update app.py

app.py

@@ -13,16 +13,16 @@ import onnxruntime
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 SWAPPER_MODEL_PATH = "models/inswapper_128.onnx"
-GFPGAN_MODEL_PATH = "models/gfpgan_1.4.onnx" # Or your chosen restoration model
+RESTORATION_MODEL_PATH = "models/gfpgan_1.4.onnx" # Path to your GFPGAN or chosen restoration ONNX model
 
-FACE_ANALYZER_NAME = 'buffalo_l'
-DETECTION_SIZE = (640, 640)
-EXECUTION_PROVIDERS = ['CPUExecutionProvider']
+FACE_ANALYZER_NAME = 'buffalo_l' # InsightFace detection model
+DETECTION_SIZE = (640, 640) # Input size for face detection
+EXECUTION_PROVIDERS = ['CPUExecutionProvider'] # Or ['CUDAExecutionProvider', 'CPUExecutionProvider'] for GPU
 
-# --- Global Variables (Lazy Loaded) ---
+# --- Global Variables (Lazy Loaded by initialize_models) ---
 face_analyzer = None
 swapper = None
-face_restorer = None # This will be updated by initialize_models
+face_restorer = None
 
 # --- Initialization Functions ---
 def initialize_models():
@@ -32,197 +32,298 @@ def initialize_models():
         if face_analyzer is None:
             logging.info(f"Initializing FaceAnalysis model: {FACE_ANALYZER_NAME}")
             face_analyzer = FaceAnalysis(name=FACE_ANALYZER_NAME, providers=EXECUTION_PROVIDERS)
-            face_analyzer.prepare(ctx_id=0, det_size=DETECTION_SIZE)
+            face_analyzer.prepare(ctx_id=0, det_size=DETECTION_SIZE) # ctx_id=0 for CPU, -1 for auto
             logging.info("FaceAnalysis model initialized successfully.")
 
         # Initialize Swapper model
         if swapper is None:
             if not os.path.exists(SWAPPER_MODEL_PATH):
                 logging.error(f"Swapper model FILE NOT FOUND at {SWAPPER_MODEL_PATH}. Swapping will fail.")
-                # Keep swapper as None, subsequent checks should handle this
             else:
                 logging.info(f"Loading swapper model from: {SWAPPER_MODEL_PATH}")
                 try:
+                    # Pass providers to get_model if supported
                     swapper = get_model(SWAPPER_MODEL_PATH, download=False, providers=EXECUTION_PROVIDERS)
                 except TypeError:
                     logging.warning(f"Failed to pass 'providers' to swapper model {SWAPPER_MODEL_PATH}. Retrying without 'providers' argument.")
-                    swapper = get_model(SWAPPER_MODEL_PATH, download=False)
+                    swapper = get_model(SWAPPER_MODEL_PATH, download=False) # Fallback
                 logging.info("Swapper model loaded successfully.")
 
         # Initialize Face Restoration Model
-        # This assignment happens once. If it fails, face_restorer remains None.
-        if face_restorer is None: # Check added to ensure it only tries to load once if called multiple times
-            if not os.path.exists(GFPGAN_MODEL_PATH):
-                # This log is critical for the user to see
-                logging.error(f"Face restoration model FILE NOT FOUND at: {GFPGAN_MODEL_PATH}. Enhancement feature will be disabled.")
-                face_restorer = None # Explicitly ensure it's None
+        if face_restorer is None: # Only attempt to load if not already loaded/failed
+            if not os.path.exists(RESTORATION_MODEL_PATH):
+                logging.error(f"Face restoration model FILE NOT FOUND at: {RESTORATION_MODEL_PATH}. Enhancement feature will be disabled.")
             else:
-                logging.info(f"Attempting to load face restoration model from: {GFPGAN_MODEL_PATH}")
+                logging.info(f"Attempting to load face restoration model from: {RESTORATION_MODEL_PATH}")
                 try:
-                    face_restorer = onnxruntime.InferenceSession(GFPGAN_MODEL_PATH, providers=EXECUTION_PROVIDERS)
+                    face_restorer = onnxruntime.InferenceSession(RESTORATION_MODEL_PATH, providers=EXECUTION_PROVIDERS)
                     logging.info("Face restoration model loaded successfully.")
                 except Exception as e:
-                    # This log is also critical
-                    logging.error(f"Error loading face restoration model from {GFPGAN_MODEL_PATH}: {e}. Enhancement feature will be disabled.", exc_info=True)
+                    logging.error(f"Error loading face restoration model from {RESTORATION_MODEL_PATH}: {e}. Enhancement feature will be disabled.", exc_info=True)
                     face_restorer = None # Ensure it's None if loading failed
-
-    except Exception as e: # Catch other unexpected errors during initialization
+    except Exception as e:
         logging.error(f"A critical error occurred during model initialization: {e}", exc_info=True)
-        # Depending on what failed, face_analyzer, swapper or face_restorer might be None
-        # The application might be in a partially usable or unusable state.
+        # face_analyzer, swapper, or face_restorer might be None. Subsequent checks will handle this.
 
 # --- Call Initialization Early ---
-# This ensures models are loaded (or attempted to be loaded) before the UI is built,
-# allowing the UI to reflect the status of optional models like the face_restorer.
 initialize_models()
-# Check if critical models loaded, otherwise, it might not make sense to proceed for UI.
-if face_analyzer is None or swapper is None:
-    logging.critical("Core models (Face Analyzer or Swapper) failed to initialize. Application functionality will be severely limited or non-functional.")
-    # Optionally raise an error here to stop the script if you deem these absolutely essential for any UI interaction.
-    # raise RuntimeError("Core models failed to load. Cannot start application.")
+# Global flags based on model loading status, used for UI conditional rendering
+core_models_loaded_successfully = face_analyzer is not None and swapper is not None
+restoration_model_loaded_successfully = face_restorer is not None
 
-
-# (Image Utility Functions, Core AI Functions, enhance_cropped_face - remain the same as your last provided version)
-# ... (ensure these functions are present here from your previous code) ...
-def convert_pil_to_cv2(pil_image: Image.Image) -> np.ndarray:
+# --- Image Utility Functions ---
+def convert_pil_to_cv2(pil_image: Image.Image) -> np.ndarray | None:
     if pil_image is None: return None
-    return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
+    try:
+        return cv2.cvtColor(np.array(pil_image.convert('RGB')), cv2.COLOR_RGB2BGR)
+    except Exception as e:
+        logging.error(f"Error converting PIL to CV2: {e}")
+        return None
 
-def convert_cv2_to_pil(cv2_image: np.ndarray) -> Image.Image:
+def convert_cv2_to_pil(cv2_image: np.ndarray) -> Image.Image | None:
     if cv2_image is None: return None
-    return Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))
-
-def get_faces_from_image(img_np: np.ndarray):
-    if face_analyzer is None:
-        gr.Warning("Face analyzer not initialized. Detection might fail.") # Changed to gr.Warning
+    try:
+        return Image.fromarray(cv2.cvtColor(cv2_image, cv2.COLOR_BGR2RGB))
+    except Exception as e:
+        logging.error(f"Error converting CV2 to PIL: {e}")
+        return None
+
+# --- Core AI & Image Processing Functions ---
+def get_faces_from_image(img_np: np.ndarray) -> list:
+    if not core_models_loaded_successfully or face_analyzer is None:
+        # This condition should ideally be caught before calling this function if core models failed.
+        logging.error("Face analyzer not available for get_faces_from_image.")
         return []
     if img_np is None: return []
-    return face_analyzer.get(img_np)
+    try:
+        return face_analyzer.get(img_np)
+    except Exception as e:
+        logging.error(f"Error during face detection: {e}", exc_info=True)
+        return []
 
-def draw_detected_faces(img_np: np.ndarray, faces: list):
+def draw_detected_faces(img_np: np.ndarray, faces: list) -> np.ndarray:
     img_with_boxes = img_np.copy()
     for i, face in enumerate(faces):
-        box = face.bbox.astype(int)
+        box = face.bbox.astype(int) # InsightFace bbox is [x1, y1, x2, y2]
         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))
+        label_position = (x1, max(0, y1 - 10)) # Ensure label is not drawn outside top
         cv2.putText(img_with_boxes, f"Face {i}", label_position,
                     cv2.FONT_HERSHEY_SIMPLEX, 0.6, (36, 255, 12), 2)
     return img_with_boxes
 
 def enhance_cropped_face(face_crop_bgr: np.ndarray) -> np.ndarray:
-    if face_restorer is None:
+    if not restoration_model_loaded_successfully or face_restorer is None:
         logging.warning("Face restorer model not available. Skipping enhancement for crop.")
         return face_crop_bgr
-    if face_crop_bgr.shape[0] == 0 or face_crop_bgr.shape[1] == 0:
-        logging.warning("Received empty face crop for enhancement.")
-        return face_crop_bgr
+    if face_crop_bgr is None or face_crop_bgr.shape[0] == 0 or face_crop_bgr.shape[1] == 0:
+        logging.warning("Received empty or invalid face crop for enhancement.")
+        return face_crop_bgr if face_crop_bgr is not None else np.array([])
+
 
     logging.info(f"Applying face restoration to crop of size {face_crop_bgr.shape[:2]}...")
     crop_height, crop_width = face_crop_bgr.shape[:2]
-    img_rgb = cv2.cvtColor(face_crop_bgr, cv2.COLOR_BGR2RGB)
-    restorer_input_size = (512, 512)
-    img_resized_for_model = cv2.resize(img_rgb, restorer_input_size, interpolation=cv2.INTER_AREA)
-    img_normalized = (img_resized_for_model / 255.0).astype(np.float32)
-    img_chw = np.transpose(img_normalized, (2, 0, 1))
-    input_tensor = np.expand_dims(img_chw, axis=0)
-    input_name = face_restorer.get_inputs()[0].name
-    output_name = face_restorer.get_outputs()[0].name
+
     try:
+        img_rgb = cv2.cvtColor(face_crop_bgr, cv2.COLOR_BGR2RGB)
+        
+        # GFPGAN typically expects 512x512 input.
+        restorer_input_size = (512, 512) # Check your specific model's expected input size
+        img_resized_for_model = cv2.resize(img_rgb, restorer_input_size, interpolation=cv2.INTER_AREA)
+        
+        img_normalized = (img_resized_for_model / 255.0).astype(np.float32) # Normalize to [0, 1]
+        img_chw = np.transpose(img_normalized, (2, 0, 1)) # HWC to CHW
+        input_tensor = np.expand_dims(img_chw, axis=0) # Add batch dimension
+
+        input_name = face_restorer.get_inputs()[0].name
+        output_name = face_restorer.get_outputs()[0].name
+        
         restored_output_model_size = face_restorer.run([output_name], {input_name: input_tensor})[0]
+
+        restored_img_chw = np.squeeze(restored_output_model_size, axis=0)
+        restored_img_hwc_model_size = np.transpose(restored_img_chw, (1, 2, 0)) # CHW to HWC
+        restored_img_uint8_model_size = np.clip(restored_img_hwc_model_size * 255.0, 0, 255).astype(np.uint8)
+
+        # Resize back to the original crop's dimensions
+        restored_crop_rgb = cv2.resize(restored_img_uint8_model_size, (crop_width, crop_height), interpolation=cv2.INTER_LANCZOS4)
+        restored_crop_bgr = cv2.cvtColor(restored_crop_rgb, cv2.COLOR_RGB2BGR)
+        
+        logging.info("Cropped face restoration complete.")
+        return restored_crop_bgr
     except Exception as e:
-        logging.error(f"Error during face restoration inference for crop: {e}")
-        return face_crop_bgr
-    restored_img_chw = np.squeeze(restored_output_model_size, axis=0)
-    restored_img_hwc_model_size = np.transpose(restored_img_chw, (1, 2, 0))
-    restored_img_uint8_model_size = np.clip(restored_img_hwc_model_size * 255.0, 0, 255).astype(np.uint8)
-    restored_crop_rgb = cv2.resize(restored_img_uint8_model_size, (crop_width, crop_height), interpolation=cv2.INTER_LANCZOS4)
-    restored_crop_bgr = cv2.cvtColor(restored_crop_rgb, cv2.COLOR_RGB2BGR)
-    logging.info("Cropped face restoration complete.")
-    return restored_crop_bgr
+        logging.error(f"Error during face restoration for crop: {e}", exc_info=True)
+        return face_crop_bgr # Return original crop on error
+
+def histogram_match_channel(source_channel: np.ndarray, target_channel: np.ndarray) -> np.ndarray:
+    """Matches histogram of a single source channel to a target channel."""
+    source_shape = source_channel.shape
+    source_channel_flat = source_channel.flatten()
+    target_channel_flat = target_channel.flatten()
+
+    # Get the histogram and CDF of the source image
+    source_hist, bins = np.histogram(source_channel_flat, 256, [0,256])
+    source_cdf = source_hist.cumsum()
+
+    # Get the histogram and CDF of the target image
+    target_hist, bins = np.histogram(target_channel_flat, 256, [0,256])
+    target_cdf = target_hist.cumsum()
+
+    # Normalize CDFs
+    source_cdf_norm = source_cdf * float(target_hist.max()) / source_hist.max() # Avoid divide by zero if source_hist.max() is 0
+    target_cdf_norm = target_cdf # No change needed for target if using its own max
+
+    # Create a lookup table
+    lookup_table = np.zeros(256, dtype='uint8')
+    gj = 0
+    for gi in range(256):
+        while gj < 256 and target_cdf_norm[gj] < source_cdf_norm[gi]:
+            gj += 1
+        if gj == 256: # Safety for out of bounds
+            gj = 255
+        lookup_table[gi] = gj
+    
+    matched_channel_flat = cv2.LUT(source_channel, lookup_table) # Apply lookup table
+    return matched_channel_flat.reshape(source_shape)
+
+
+def histogram_match_color(source_img: np.ndarray, target_img: np.ndarray) -> np.ndarray:
+    """
+    Performs histogram matching on color images (BGR).
+    """
+    if source_img is None or target_img is None or source_img.size == 0 or target_img.size == 0:
+        logging.warning("Cannot perform histogram matching on empty or None images.")
+        return source_img if source_img is not None else np.array([])
+
+    matched_img = np.zeros_like(source_img)
+    try:
+        for i in range(source_img.shape[2]): # Iterate over B, G, R channels
+            matched_img[:,:,i] = histogram_match_channel(source_img[:,:,i], target_img[:,:,i])
+        return matched_img
+    except Exception as e:
+        logging.error(f"Error during histogram matching: {e}", exc_info=True)
+        return source_img # Return original on error
+
 
 # --- Gradio Interface Functions ---
 def process_face_swap(source_pil_img: Image.Image, target_pil_img: Image.Image,
-                        target_face_index: int, apply_enhancement: bool,
+                        target_face_index: int, apply_enhancement: bool, apply_color_correction: bool,
                         progress=gr.Progress(track_tqdm=True)):
-    progress(0, desc="Initializing...")
-    if swapper is None or face_analyzer is None:
-        # This is a fallback; ideally, the app shouldn't fully start or should show a persistent error.
-        gr.Error("Core models (Swapper or Analyzer) not initialized. Please check logs and restart the application.")
-        # Return empty results or raise the error to stop processing
-        # For safety, let's return empty results for Gradio components.
-        blank_pil = Image.new('RGB', (100,100), color='grey')
-        return blank_pil, None # For output_image and download_output
+    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 # Placeholder for output
 
     if source_pil_img is None: raise gr.Error("Source image not provided.")
     if target_pil_img is None: raise gr.Error("Target image not provided.")
 
-    progress(0.1, desc="Converting images...")
+    progress(0.05, desc="Converting images...")
     source_np = convert_pil_to_cv2(source_pil_img)
     target_np = convert_pil_to_cv2(target_pil_img)
+
+    if source_np is None or target_np is None:
+        raise gr.Error("Image conversion failed. Please try different images.")
+    
     target_h, target_w = target_np.shape[:2]
 
-    progress(0.2, desc="Detecting face in source image...")
+    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.")
-    source_face = source_faces[0]
+    if not source_faces: raise gr.Error("No face found in the source image. Please use a clear image of a single face.")
+    source_face = source_faces[0] # Assuming the first detected face is the one to use
 
-    progress(0.3, desc="Detecting faces in target image...")
+    progress(0.25, desc="Detecting faces in target image...")
     target_faces = get_faces_from_image(target_np)
     if not target_faces: raise gr.Error("No faces found in the target image.")
     if not (0 <= target_face_index < len(target_faces)):
-        raise gr.Error(f"Selected face index ({target_face_index}) is out of range.")
-    target_face_to_swap_info = target_faces[int(target_face_index)]
+        # This case should ideally be prevented by the slider's dynamic range
+        raise gr.Error(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}).")
+    
+    target_face_to_swap_info = target_faces[int(target_face_index)] # Face object from InsightFace
 
-    swapped_bgr_img = target_np # Initialize with target_np to have a fallback
+    swapped_bgr_img = target_np.copy() # Start with a copy of target
     try:
-        progress(0.5, desc="Performing face swap...")
+        progress(0.4, desc="Performing face swap...")
+        # swapper.get returns the modified target_np with the face pasted back
         swapped_bgr_img = swapper.get(target_np, target_face_to_swap_info, source_face, paste_back=True)
 
+        # Define bounding box for post-processing (enhancement and color correction)
+        # Use the bbox of the target face where the swap occurred.
+        # InsightFace bbox is [x1, y1, x2, y2]
+        bbox_coords = target_face_to_swap_info.bbox.astype(int)
+        
         if apply_enhancement:
-            if face_restorer is not None: # Check again, in case it was somehow re-enabled
-                progress(0.7, desc="Applying selective face enhancement...")
-                bbox = target_face_to_swap_info.bbox.astype(int)
-                padding = 20
-                x1 = max(0, bbox[0] - padding)
-                y1 = max(0, bbox[1] - padding)
-                x2 = min(target_w, bbox[2] + padding)
-                y2 = min(target_h, bbox[3] + padding)
-
-                if x1 < x2 and y1 < y2:
-                    face_crop_to_enhance = swapped_bgr_img[y1:y2, x1:x2]
-                    if face_crop_to_enhance.size > 0:
-                        enhanced_crop = enhance_cropped_face(face_crop_to_enhance.copy())
-                        swapped_bgr_img[y1:y2, x1:x2] = enhanced_crop
+            if restoration_model_loaded_successfully:
+                progress(0.6, desc="Applying selective face enhancement...")
+                padding_enh = 20 # Pixels of padding around the bbox for enhancement context
+                
+                enh_x1 = max(0, bbox_coords[0] - padding_enh)
+                enh_y1 = max(0, bbox_coords[1] - padding_enh)
+                enh_x2 = min(target_w, bbox_coords[2] + padding_enh)
+                enh_y2 = min(target_h, bbox_coords[3] + padding_enh)
+
+                if enh_x1 < enh_x2 and enh_y1 < enh_y2: # Check if the crop is valid
+                    face_crop_to_enhance = swapped_bgr_img[enh_y1:enh_y2, enh_x1:enh_x2]
+                    enhanced_crop = enhance_cropped_face(face_crop_to_enhance.copy()) # Use .copy()
+                    if enhanced_crop.shape == face_crop_to_enhance.shape: # Ensure enhanced crop is valid and same size
+                        swapped_bgr_img[enh_y1:enh_y2, enh_x1:enh_x2] = enhanced_crop
                     else:
-                        logging.warning("Skipping enhancement, face crop region is empty.")
+                        logging.warning("Enhanced crop size mismatch. Skipping paste-back for enhancement.")
                 else:
                     logging.warning("Skipping enhancement, invalid crop dimensions after padding.")
             else:
-                # This message is key, as seen in your logs.
-                gr.Info("Face restoration model not available, enhancement skipped by process_face_swap.")
+                gr.Info("Face restoration model not available, enhancement skipped by process.")
                 logging.warning("Enhancement requested at runtime but face restorer is not available.")
+        
+        if apply_color_correction:
+            progress(0.75, desc="Applying color correction...")
+            # For color correction, we use the original target face region from `target_np` as the reference.
+            # The region to correct is the same bounding box in `swapped_bgr_img`.
+            cc_x1 = bbox_coords[0]
+            cc_y1 = bbox_coords[1]
+            cc_x2 = bbox_coords[2]
+            cc_y2 = bbox_coords[3]
+
+            if cc_x1 < cc_x2 and cc_y1 < cc_y2: # Valid bbox
+                target_face_region_for_color = target_np[cc_y1:cc_y2, cc_x1:cc_x2]
+                swapped_face_region_to_correct = swapped_bgr_img[cc_y1:cc_y2, cc_x1:cc_x2]
+
+                if target_face_region_for_color.size > 0 and swapped_face_region_to_correct.size > 0:
+                    corrected_swapped_region = histogram_match_color(swapped_face_region_to_correct.copy(), target_face_region_for_color.copy())
+                    if corrected_swapped_region.shape == swapped_face_region_to_correct.shape:
+                         swapped_bgr_img[cc_y1:cc_y2, cc_x1:cc_x2] = corrected_swapped_region
+                    else:
+                        logging.warning("Color corrected region size mismatch. Skipping paste-back for color correction.")
+                else:
+                    logging.warning("Skipping color correction, empty face region(s) for matching.")
+            else:
+                logging.warning("Skipping color correction, invalid bounding box for region extraction.")
+
 
     except Exception as e:
-        logging.error(f"Error during face swapping or enhancement: {e}", exc_info=True)
-        # Return the last valid state of swapped_bgr_img if it exists, or original target.
-        # This helps avoid a complete UI crash if swap partially succeeded.
-        swapped_pil_img_on_error = convert_cv2_to_pil(swapped_bgr_img if 'swapped_bgr_img' in locals() else target_np)
+        logging.error(f"Error during face swapping or post-processing: {e}", exc_info=True)
+        # Return the current state of swapped_bgr_img to avoid losing partial work if possible
+        swapped_pil_img_on_error = convert_cv2_to_pil(swapped_bgr_img)
+        if swapped_pil_img_on_error is None: # Fallback if conversion itself fails
+             swapped_pil_img_on_error = Image.new('RGB', (target_w, target_h), color='lightgrey')
         raise gr.Error(f"An error occurred: {str(e)}")
-        # return swapped_pil_img_on_error, None # Alternative to raising error for UI
+        # return swapped_pil_img_on_error, None # Alternative for UI stability
 
     progress(0.9, desc="Finalizing image...")
     swapped_pil_img = convert_cv2_to_pil(swapped_bgr_img)
+    if swapped_pil_img is None: # Handle case where final conversion fails
+        gr.Error("Failed to convert final image to display format.")
+        swapped_pil_img = Image.new('RGB', (target_w, target_h), color='lightgrey') # Placeholder
 
     temp_file_path = None
     try:
         # Using 'with' ensures the file is closed before Gradio tries to use it
-        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
-            swapped_pil_img.save(tmp_file.name)
+        with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False, mode="wb") as tmp_file:
+            swapped_pil_img.save(tmp_file, format="JPEG") # Specify format for PIL save
             temp_file_path = tmp_file.name
         logging.info(f"Swapped image saved to temporary file: {temp_file_path}")
     except Exception as e:
-        logging.error(f"Error saving to temporary file: {e}")
+        logging.error(f"Error saving to temporary file: {e}", exc_info=True)
         gr.Warning("Could not save the swapped image for download.")
     
     progress(1.0, desc="Processing complete!")
@@ -234,87 +335,104 @@ def preview_target_faces(target_pil_img: Image.Image):
         return blank_image_pil, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
 
     target_np = convert_pil_to_cv2(target_pil_img)
-    if target_np is None:
+    if target_np is None: # If conversion failed
         blank_image_pil = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
         return blank_image_pil, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
         
     faces = get_faces_from_image(target_np)
+    if not faces: # If no faces are detected
+        preview_pil_img = convert_cv2_to_pil(target_np) # Show original image if no faces
+        if preview_pil_img is None: # Fallback
+            preview_pil_img = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
+        return preview_pil_img, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
+
     preview_np_img = draw_detected_faces(target_np, faces)
     preview_pil_img = convert_cv2_to_pil(preview_np_img)
+    if preview_pil_img is None: # Fallback
+        preview_pil_img = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
+
 
     num_faces = len(faces)
-    slider_update = gr.Slider(minimum=0, maximum=max(0, num_faces - 1), value=0, step=1, interactive=num_faces > 0)
+    slider_update = gr.Slider(minimum=0, maximum=max(0, num_faces - 1), value=0, step=1, interactive=(num_faces > 0))
     return preview_pil_img, slider_update
 
 
 # --- Gradio UI Definition ---
-with gr.Blocks(title="Ultimate Face Swap AI 💎 v2", theme=gr.themes.Glass()) as demo:
+with gr.Blocks(title="Ultimate Face Swap AI 💎 v3", theme=gr.themes.Soft()) as demo: # Changed theme for variety
     gr.Markdown(
         """
-        # 💎 Ultimate Face Swap AI 🌠 v2
-        Upload a source face, a target image, and let the AI work its magic!
-        Optionally enhance the swapped face for stunning realism.
+        <div style="text-align: center;">
+            <h1>💎 Ultimate Face Swap AI 🌠 v3</h1>
+            <p>Upload a source face, a target image, and let the AI work its magic!</p>
+            <p>Optionally enhance with face restoration and color correction for stunning realism.</p>
+        </div>
         """
     )
 
-    # Check status of restoration model for UI feedback
-    restoration_model_loaded = face_restorer is not None
-    core_models_loaded = face_analyzer is not None and swapper is not None
-
-    if not core_models_loaded:
-        gr.Error("CRITICAL ERROR: Core models (Face Analyzer or Swapper) failed to load. Application will not function correctly. Please check the console logs for details (e.g., model file paths).")
+    if not core_models_loaded_successfully:
+        gr.Error("CRITICAL ERROR: Core models (Face Analyzer or Swapper) failed to load. Application will not function correctly. Please check the console logs for details (e.g., model file paths) and restart the application after fixing.")
 
     with gr.Row():
         with gr.Column(scale=1):
-            source_image_input = gr.Image(label="👤 Source Face Image", type="pil", sources=["upload", "clipboard"], height=400)
+            source_image_input = gr.Image(label="👤 Source Face Image (Clear, single face)", type="pil", sources=["upload", "clipboard"], height=350)
         with gr.Column(scale=1):
-            target_image_input = gr.Image(label="🖼️ Target Scene Image", type="pil", sources=["upload", "clipboard"], height=400)
+            target_image_input = gr.Image(label="🖼️ Target Scene Image", type="pil", sources=["upload", "clipboard"], height=350)
 
     with gr.Row(equal_height=True):
-        preview_button = gr.Button("🔍 Preview Detected Faces & Select", variant="secondary")
+        preview_button = gr.Button("🔍 Preview & Select Target Face", variant="secondary")
         face_index_slider = gr.Slider(
             label="🎯 Select Target Face (0-indexed)",
             minimum=0, maximum=0, step=1, value=0, interactive=False
         )
     
-    target_faces_preview_output = gr.Image(label="👀 Detected Faces in Target", interactive=False, height=400)
+    target_faces_preview_output = gr.Image(label="👀 Detected Faces in Target", interactive=False, height=350)
     
-    gr.HTML("<hr style='margin-top: 20px; margin-bottom: 20px;'>")
+    gr.HTML("<hr style='margin-top: 15px; margin-bottom: 15px;'>") # Styled HR
 
-    enhance_checkbox_label = "✨ Apply Selective Face Restoration (GFPGAN)"
-    if not restoration_model_loaded:
-        enhance_checkbox_label += " (Model not loaded)"
+    with gr.Row():
+        with gr.Column(scale=1):
+            enhance_checkbox_label = "✨ Apply Selective Face Restoration"
+            if not restoration_model_loaded_successfully:
+                enhance_checkbox_label += " (Model N/A)"
+            enhance_checkbox = gr.Checkbox(
+                label=enhance_checkbox_label,
+                value=restoration_model_loaded_successfully, # Default to True only if available
+                interactive=restoration_model_loaded_successfully # Disable if not available
+            )
+            if not restoration_model_loaded_successfully:
+                gr.Markdown("<p style='color: orange; font-size:0.8em;'>⚠️ Face restoration model not loaded. Feature disabled. Check logs.</p>")
+        
+        with gr.Column(scale=1):
+            color_correction_checkbox = gr.Checkbox(label="🎨 Apply Color Correction (Histogram Matching)", value=True) # Default to True
 
-    enhance_checkbox = gr.Checkbox(
-        label=enhance_checkbox_label,
-        value=restoration_model_loaded, # Default to True only if available and loaded
-        interactive=restoration_model_loaded # Disable if not available/loaded
-    )
-    if not restoration_model_loaded:
-        gr.Markdown("<p style='color: orange; font-size:0.9em;'>⚠️ Face restoration model could not be loaded. Enhancement feature is disabled. Please check console logs for model path errors (e.g., `models/gfpgan_1.4.onnx`).</p>")
     
     with gr.Row():
-        swap_button = gr.Button("🚀 SWAP & ENHANCE NOW!", variant="primary", scale=3, interactive=core_models_loaded) # Disable if core models failed
+        swap_button = gr.Button("🚀 GENERATE SWAP!", variant="primary", scale=3, interactive=core_models_loaded_successfully) # Disable if core models failed
         clear_button = gr.Button("🧹 Clear All", variant="stop", scale=1)
 
+
     with gr.Row():
-        swapped_image_output = gr.Image(label="✨ Swapped Result", interactive=False, height=500)
+        swapped_image_output = gr.Image(label="✨ Swapped Result", interactive=False, height=450)
         download_output_file = gr.File(label="⬇️ Download Swapped Image")
 
-    # Event Handlers (on_target_image_change_or_clear, clear_all_inputs_outputs, preview_button.click, swap_button.click, gr.Examples - same as your previous code)
-    def on_target_image_change_or_clear(target_img):
-        if target_img is None: # Cleared
+    # --- Event Handlers ---
+    def on_target_image_change_or_clear(target_img_pil):
+        if target_img_pil is None: # Cleared
             blank_image_pil = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color = 'lightgray')
             return blank_image_pil, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
-        return target_faces_preview_output.value, face_index_slider.value 
+        # If an image is uploaded, preview_target_faces will be called by the button click.
+        # This function is mainly for handling the clear/reset case or an initial state.
+        # For safety, return current values if not cleared, or trigger preview.
+        # Let's just return blank on clear, and let preview button handle new image.
+        return target_faces_preview_output.value, face_index_slider.value # Keep current state on new image upload until preview
 
     target_image_input.change(
         fn=on_target_image_change_or_clear,
         inputs=[target_image_input],
         outputs=[target_faces_preview_output, face_index_slider],
+        queue=False
     )
-    source_image_input.change(fn=None)
-
+    
     preview_button.click(
         fn=preview_target_faces,
         inputs=[target_image_input],
@@ -323,59 +441,72 @@ with gr.Blocks(title="Ultimate Face Swap AI 💎 v2", theme=gr.themes.Glass()) a
 
     swap_button.click(
         fn=process_face_swap,
-        inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox],
+        inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox, color_correction_checkbox],
         outputs=[swapped_image_output, download_output_file]
     )
     
     def clear_all_inputs_outputs():
         blank_preview = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color = 'lightgray')
-        return (None, None, 
-                blank_preview, 
-                gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False), 
-                None, None) 
+        # Also reset the slider and preview output properly
+        return (
+            None, # source_image_input
+            None, # target_image_input
+            blank_preview, # target_faces_preview_output
+            gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False), # face_index_slider
+            None, # swapped_image_output
+            None  # download_output_file
+        )
 
     clear_button.click(
         fn=clear_all_inputs_outputs,
-        inputs=None,
-        outputs=[source_image_input, target_image_input, 
-                 target_faces_preview_output, face_index_slider, 
-                 swapped_image_output, download_output_file],
-        queue=False 
+        inputs=None, # No inputs for this function
+        outputs=[
+            source_image_input, target_image_input, 
+            target_faces_preview_output, face_index_slider, 
+            swapped_image_output, download_output_file
+        ],
+        queue=False # No need to queue a simple clear
     )
 
     gr.Examples(
         examples=[
-            ["examples/source_face.jpg", "examples/target_group.jpg", 0, True],
-            ["examples/source_actor.png", "examples/target_scene.png", 1, True],
-            ["examples/source_face.jpg", "examples/target_group.jpg", 0, False],
+            ["examples/source_face.jpg", "examples/target_group.jpg", 0, True, True],
+            ["examples/source_actor.png", "examples/target_scene.png", 1, True, True],
+            ["examples/source_face.jpg", "examples/target_group.jpg", 0, False, True], # No enhancement, with color correction
+            ["examples/source_actor.png", "examples/target_scene.png", 0, True, False], # With enhancement, no color correction
         ],
-        inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox],
+        inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox, color_correction_checkbox],
         outputs=[swapped_image_output, download_output_file],
-        fn=process_face_swap,
-        cache_examples="lazy",
-        label="Example Face Swaps (Click to run with progress)"
+        fn=process_face_swap, # This will also show progress for examples
+        cache_examples=False, # Set to "lazy" or True if processing is slow and examples are static
+        label="Example Face Swaps (Click to run)"
     )
 
-
+# --- Main Execution Block ---
 if __name__ == "__main__":
     os.makedirs("models", exist_ok=True)
-    os.makedirs("examples", exist_ok=True)
-
-    # Startup messages are now part of initialize_models() or the UI definition.
-    # We can add a final check before launch.
-    if not core_models_loaded:
-        print("\n" + "="*50)
-        print("CRITICAL ERROR: Core models (Face Analyzer or Swapper) failed to load.")
-        print("The application UI will load but will be non-functional.")
-        print(f"Please check: \n  - Face Analyzer: '{FACE_ANALYZER_NAME}' \n  - Swapper Model: '{SWAPPER_MODEL_PATH}'")
-        print("Ensure the model names are correct and files exist.")
-        print("="*50 + "\n")
-    elif not restoration_model_loaded:
-        print("\n" + "="*50)
-        print("INFO: Face Restoration model not loaded.")
-        print(f"The enhancement feature will be disabled. Check path: '{GFPGAN_MODEL_PATH}'")
-        print("="*50 + "\n")
+    os.makedirs("examples", exist_ok=True) # Create if you intend to use example images
+
+    # Console messages about model loading status
+    print("\n" + "="*60)
+    print("🚀 ULTIMATE FACE SWAP AI - STARTUP STATUS 🚀")
+    print("="*60)
+    if not core_models_loaded_successfully:
+        print("🔴 CRITICAL ERROR: Core models (Face Analyzer or Swapper) FAILED to load.")
+        print(f"   - Face Analyzer: '{FACE_ANALYZER_NAME}' (Status: {'Loaded' if face_analyzer else 'Failed'})")
+        print(f"   - Swapper Model: '{SWAPPER_MODEL_PATH}' (Status: {'Loaded' if swapper else 'Failed'})")
+        print("   The application UI will load but will be NON-FUNCTIONAL.")
+        print("   Please check model paths and ensure files exist and are not corrupted.")
     else:
-        print("\nAll models loaded successfully or optional models handled.\n")
+        print("🟢 Core models (Face Analyzer & Swapper) loaded successfully.")
 
+    if not restoration_model_loaded_successfully:
+        print(f"🟡 INFO: Face Restoration model ('{RESTORATION_MODEL_PATH}') not loaded.")
+        print("   The 'Face Restoration' feature will be disabled.")
+        print("   To enable, ensure the model file exists at the specified path and is valid.")
+    else:
+        print("🟢 Face Restoration model loaded successfully.")
+    print("="*60 + "\n")
+    
+    print("Launching Gradio Interface...")
     demo.launch()