Update app.py

app.py

@@ -8,7 +8,8 @@ from PIL import Image
 import tempfile
 import logging
 import onnxruntime
-from typing import List, Optional, Tuple # For improved type hinting
+from typing import List, Optional, Tuple, Any # For improved type hinting
+import shutil # For cleaning up temporary directories
 
 # --- Configuration & Setup ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -40,7 +41,7 @@ logging.info(f"Using Execution Providers: {EXECUTION_PROVIDERS}")
 logging.info(f"InsightFace Context ID for preparation: {INSIGHTFACE_CTX_ID}")
 
 face_analyzer: Optional[FaceAnalysis] = None
-swapper = None
+swapper: Optional[Any] = None # Type Any because insightface.model_zoo.get_model doesn't have precise type hints readily available
 face_restorer: Optional[onnxruntime.InferenceSession] = None
 
 def initialize_models():
@@ -58,10 +59,11 @@ def initialize_models():
             else:
                 logging.info(f"Loading swapper model from: {SWAPPER_MODEL_PATH} with providers: {EXECUTION_PROVIDERS}")
                 try:
+                    # Attempt to pass providers if the model supports it
                     swapper = get_model(SWAPPER_MODEL_PATH, download=False, providers=EXECUTION_PROVIDERS)
-                except TypeError:
+                except TypeError: # Some versions of get_model might not accept 'providers'
                     logging.warning(f"Failed to pass 'providers' argument to swapper model {SWAPPER_MODEL_PATH}. Retrying without 'providers' argument. Model will use its default execution provider.")
-                    swapper = get_model(SWAPPER_MODEL_PATH, download=False)
+                    swapper = get_model(SWAPPER_MODEL_PATH, download=False) # Fallback
                 logging.info("Swapper model loaded successfully.")
 
         if face_restorer is None:
@@ -82,6 +84,7 @@ initialize_models()
 core_models_loaded_successfully = face_analyzer is not None and swapper is not None
 restoration_model_loaded_successfully = face_restorer is not None
 
+# --- Utility Functions (Image Conversion, Face Detection, Drawing) ---
 def convert_pil_to_cv2(pil_image: Optional[Image.Image]) -> Optional[np.ndarray]:
     if pil_image is None: return None
     try:
@@ -98,7 +101,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:
+def get_faces_from_image(img_np: np.ndarray) -> List[Any]: # Using Any for insightface.app.common.Face
     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 []
@@ -112,7 +115,7 @@ def get_faces_from_image(img_np: np.ndarray) -> List:
         logging.error(f"Error during face detection: {e}", exc_info=True)
         return []
 
-def draw_detected_faces(img_np: np.ndarray, faces: List) -> np.ndarray:
+def draw_detected_faces(img_np: np.ndarray, faces: List[Any]) -> np.ndarray:
     img_with_boxes = img_np.copy()
     for i, face in enumerate(faces):
         box = face.bbox.astype(int)
@@ -120,9 +123,10 @@ def draw_detected_faces(img_np: np.ndarray, faces: List) -> np.ndarray:
         cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (0, 255, 0), 2)
         label_position = (x1, max(0, y1 - 10))
         cv2.putText(img_with_boxes, f"Face {i}", label_position,
-                      cv2.FONT_HERSHEY_SIMPLEX, 0.7, (36, 255, 12), 2, cv2.LINE_AA)
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.7, (36, 255, 12), 2, cv2.LINE_AA)
     return img_with_boxes
 
+# --- Enhancement & Color Correction ---
 def enhance_cropped_face(face_crop_bgr: np.ndarray) -> np.ndarray:
     if not restoration_model_loaded_successfully or face_restorer is None:
         logging.warning("Face restorer model not available. Skipping enhancement for crop.")
@@ -131,7 +135,7 @@ def enhance_cropped_face(face_crop_bgr: np.ndarray) -> np.ndarray:
         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]}...")
+    logging.debug(f"Applying face restoration to crop of size {face_crop_bgr.shape[:2]}...")
     original_crop_height, original_crop_width = face_crop_bgr.shape[:2]
 
     try:
@@ -139,16 +143,16 @@ def enhance_cropped_face(face_crop_bgr: np.ndarray) -> np.ndarray:
         restorer_input_node = face_restorer.get_inputs()[0]
         try:
             shape = restorer_input_node.shape
-            if isinstance(shape[2], int) and isinstance(shape[3], int):
-                 restorer_input_size_hw = (shape[2], shape[3])
-                 logging.info(f"Restoration model expects input HxW: {restorer_input_size_hw}")
-            else:
+            if isinstance(shape[2], int) and isinstance(shape[3], int) and shape[2] > 0 and shape[3] > 0:
+                restorer_input_size_hw = (shape[2], shape[3])
+            else: # Handle dynamic axes or unexpected shapes
                 restorer_input_size_hw = (512, 512)
-                logging.warning(f"Could not determine restorer input size from model, defaulting to {restorer_input_size_hw}.")
-        except:
+                logging.warning(f"Could not reliably determine restorer input size from model shape {shape}, defaulting to {restorer_input_size_hw}.")
+        except Exception as e_shape:
             restorer_input_size_hw = (512, 512)
-            logging.warning(f"Error determining restorer input size, defaulting to {restorer_input_size_hw}.")
-
+            logging.warning(f"Error determining restorer input size (error: {e_shape}), defaulting to {restorer_input_size_hw}.")
+        
+        logging.debug(f"Restoration model expects input HxW: {restorer_input_size_hw}")
         img_resized_for_model = cv2.resize(img_rgb, (restorer_input_size_hw[1], restorer_input_size_hw[0]), interpolation=cv2.INTER_AREA)
         img_normalized = (img_resized_for_model / 255.0).astype(np.float32)
         img_chw = np.transpose(img_normalized, (2, 0, 1))
@@ -161,7 +165,7 @@ def enhance_cropped_face(face_crop_bgr: np.ndarray) -> np.ndarray:
         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, (original_crop_width, original_crop_height), interpolation=cv2.INTER_LANCZOS4)
         restored_crop_bgr = cv2.cvtColor(restored_crop_rgb, cv2.COLOR_RGB2BGR)
-        logging.info("Cropped face restoration complete.")
+        logging.debug("Cropped face restoration complete.")
         return restored_crop_bgr
     except Exception as e:
         logging.error(f"Error during face restoration for crop: {e}", exc_info=True)
@@ -171,129 +175,137 @@ def histogram_match_channel(source_channel: np.ndarray, target_channel: np.ndarr
     source_shape = source_channel.shape
     source_channel_flat = source_channel.flatten()
     target_channel_flat = target_channel.flatten()
+
     source_hist, bins = np.histogram(source_channel_flat, 256, [0, 256])
     source_cdf = source_hist.cumsum()
     if source_cdf[-1] == 0:
-        logging.warning("Source channel for histogram matching is effectively empty or flat. Returning original channel.")
+        logging.warning("Source channel for histogram matching is effectively empty. Returning original channel.")
         return source_channel
+
     target_hist, bins = np.histogram(target_channel_flat, 256, [0, 256])
     target_cdf = target_hist.cumsum()
     if target_cdf[-1] == 0:
         logging.warning("Target channel for histogram matching is effectively empty. No matching possible. Returning original source channel.")
         return source_channel
+
+    # Normalize CDFs
     source_cdf_normalized = (source_cdf * 255 / source_cdf[-1]).astype(np.uint8)
     target_cdf_normalized = (target_cdf * 255 / target_cdf[-1]).astype(np.uint8)
+
     lookup_table = np.zeros(256, dtype=np.uint8)
     gj = 0
     for gi in range(256):
         while gj < 255 and target_cdf_normalized[gj] < source_cdf_normalized[gi]:
             gj += 1
         lookup_table[gi] = gj
-    matched_channel_flat = cv2.LUT(source_channel, lookup_table)
-    return matched_channel_flat
+    
+    matched_channel = cv2.LUT(source_channel, lookup_table)
+    return matched_channel.reshape(source_shape)
+
 
 def histogram_match_color(source_img_bgr: np.ndarray, target_img_bgr: np.ndarray) -> np.ndarray:
     if source_img_bgr is None or target_img_bgr is None or source_img_bgr.size == 0 or target_img_bgr.size == 0:
         logging.warning("Cannot perform histogram matching on empty or None images.")
         return source_img_bgr if source_img_bgr is not None else np.array([])
+    
     matched_img = np.zeros_like(source_img_bgr)
     try:
-        for i in range(source_img_bgr.shape[2]):
+        for i in range(source_img_bgr.shape[2]): # For each channel (B, G, R)
             matched_img[:, :, i] = histogram_match_channel(source_img_bgr[:, :, i], target_img_bgr[:, :, i])
         return matched_img
     except Exception as e:
         logging.error(f"Error during color histogram matching: {e}", exc_info=True)
         return source_img_bgr
 
+# --- Image Face Swap Processing ---
 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.")
+    progress(0, desc="Initializing image swap...")
+    if not core_models_loaded_successfully or face_analyzer is None or swapper is None:
+        gr.Error("CRITICAL: Core models (Face Analyzer or Swapper) not loaded. Cannot proceed with image swap.")
         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 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.05, desc="Converting images to OpenCV format...")
+    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. Ensure images are valid (e.g., JPG, PNG) and not corrupted.")
+        raise gr.Error("Image conversion failed.")
     target_h, target_w = target_np.shape[:2]
 
-    progress(0.15, desc="Detecting face in source image...")
+    progress(0.15, desc="Detecting source face...")
     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.")
+        raise gr.Error("No face found in source image.")
     source_face = source_faces[0]
 
-    progress(0.25, desc="Detecting faces in target image...")
+    progress(0.25, desc="Detecting target faces...")
     target_faces = get_faces_from_image(target_np)
     if not target_faces:
-        raise gr.Error("No faces found in the target image. Please upload a target image containing at least one face.")
+        raise gr.Error("No faces found in target image.")
     if not (0 <= target_face_index < len(target_faces)):
-        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)
+        raise gr.Error(f"Target face index {target_face_index} out of range ({len(target_faces)} faces found).")
+    
     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...")
+        # Ensure swapper.get is called correctly. Assuming it modifies swapped_bgr_img or returns the modified image.
+        # Based on typical insightface usage, it modifies the input image if paste_back=True
         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:
-                progress(0.6, desc="Applying selective face enhancement...")
-                padding_enh = 20
+                progress(0.6, desc="Applying face enhancement...")
+                padding_enh = 20 
                 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:
                     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())
+                    enhanced_crop = enhance_cropped_face(face_crop_to_enhance.copy()) # Pass a copy
                     if enhanced_crop.shape == face_crop_to_enhance.shape:
                         swapped_bgr_img[enh_y1:enh_y2, enh_x1:enh_x2] = enhanced_crop
-                    else:
-                        logging.warning("Enhanced crop size mismatch. Skipping paste-back for enhancement.")
-                else:
-                    logging.warning("Skipping enhancement due to invalid crop dimensions after padding.")
+                    else: logging.warning("Enhanced crop size mismatch.")
+                else: logging.warning("Invalid crop for enhancement.")
             else:
-                logging.warning("Enhancement requested but face restorer model is not available. Enhancement skipped.")
-                gr.Info("Face restoration model not available, enhancement was skipped.")
+                logging.warning("Enhancement requested but model not available.")
+                gr.Info("Face restoration model not available, enhancement skipped.")
         
         if apply_color_correction:
-            progress(0.75, desc="Applying color correction (histogram matching)...")
+            progress(0.75, desc="Applying color correction...")
             cc_x1, cc_y1, cc_x2, cc_y2 = bbox_coords[0], bbox_coords[1], bbox_coords[2], bbox_coords[3]
             if cc_x1 < cc_x2 and cc_y1 < cc_y2:
                 original_target_face_region = 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 original_target_face_region.size > 0 and swapped_face_region_to_correct.size > 0:
                     corrected_swapped_region = histogram_match_color(
-                        swapped_face_region_to_correct.copy(),
-                        original_target_face_region.copy()
+                        swapped_face_region_to_correct.copy(), # Pass a copy
+                        original_target_face_region.copy()    # Pass a 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 due to empty face region(s) for matching.")
-            else:
-                logging.warning("Skipping color correction due to invalid bounding box for region extraction.")
+                         swapped_bgr_img[cc_y1:cc_y2, cc_x1:cc_x2] = corrected_swapped_region
+                    else: logging.warning("Color corrected region size mismatch.")
+                else: logging.warning("Empty regions for color correction.")
+            else: logging.warning("Invalid bbox for color correction.")
+
     except Exception as e:
-        logging.error(f"Error during face swapping or post-processing: {e}", exc_info=True)
+        logging.error(f"Error during image face swapping or post-processing: {e}", exc_info=True)
+        # Fallback to returning the (partially) swapped image if an error occurs mid-process
         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.")
+        if swapped_pil_img_on_error is None: # If conversion also fails
+             swapped_pil_img_on_error = Image.new('RGB', (target_w, target_h), color='lightgrey') # Placeholder
+        raise gr.Error(f"An error occurred during image processing: {str(e)}.")
 
-    progress(0.9, desc="Finalizing image conversion...")
+    progress(0.9, desc="Finalizing image...")
     swapped_pil_img = convert_cv2_to_pil(swapped_bgr_img)
     if swapped_pil_img is None:
-        gr.Error("Failed to convert final image to display format.")
+        gr.Error("Failed to convert final image.")
         return Image.new('RGB', (target_w, target_h), color='darkred'), None, "Failed to convert final image."
 
     temp_file_path = None
@@ -301,57 +313,316 @@ def process_face_swap(source_pil_img: Optional[Image.Image], target_pil_img: Opt
         with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False, mode="wb") as tmp_file:
             swapped_pil_img.save(tmp_file, format="JPEG", quality=95)
             temp_file_path = tmp_file.name
-        logging.info(f"Swapped image temporarily saved to: {temp_file_path} for download.")
     except Exception as e:
         logging.error(f"Error saving swapped image to temporary file: {e}", exc_info=True)
-        gr.Warning("Could not save the swapped image for download due to a server-side error.")
+        gr.Warning("Could not save image for download.")
+    
+    progress(1.0, desc="🎉 Image swap complete!")
+    return swapped_pil_img, temp_file_path, "🎉 Image swap complete!"
+
+# --- Video Face Swap Processing ---
+def process_video_face_swap(
+    source_pil_img: Optional[Image.Image],
+    target_video_path: Optional[str], # Path from gr.Video component
+    target_face_index_hint: int,
+    apply_enhancement: bool,
+    apply_color_correction: bool,
+    progress=gr.Progress(track_tqdm=True)
+):
+    progress(0, desc="Initializing video swap...")
+    if not core_models_loaded_successfully or face_analyzer is None or swapper is None:
+        gr.Error("CRITICAL: Core models not loaded. Cannot proceed with video swap.")
+        return None, "Core models failed."
+    if source_pil_img is None: raise gr.Error("Source image for video swap not provided.")
+    if target_video_path is None: raise gr.Error("Target video file not provided.")
+
+    source_np = convert_pil_to_cv2(source_pil_img)
+    if source_np is None: raise gr.Error("Source image conversion failed (video swap).")
+
+    source_faces = get_faces_from_image(source_np)
+    if not source_faces: raise gr.Error("No face found in source image (video swap).")
+    source_face = source_faces[0]
+
+    cap = cv2.VideoCapture(target_video_path)
+    if not cap.isOpened(): raise gr.Error(f"Could not open video file: {target_video_path}")
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    if total_frames <= 0: # Fallback if total_frames is not correctly read
+        logging.warning(f"Could not accurately determine total frames for {target_video_path}. Processing may not show correct progress percentage.")
+        # Estimate a large number or process until cap.read() fails
+        total_frames = 1 # Placeholder to avoid division by zero if loop is entered
+
+    temp_process_dir = tempfile.mkdtemp() # Create a temporary directory for all intermediate files
+    output_video_basename = f"swapped_{os.path.splitext(os.path.basename(target_video_path))[0]}"
+    output_video_path_nosound = os.path.join(temp_process_dir, output_video_basename + "_nosound.mp4")
+    final_output_video_path = os.path.join(temp_process_dir, output_video_basename + ".mp4")
     
-    progress(1.0, desc="🎉 Processing complete!")
-    return swapped_pil_img, temp_file_path, "🎉 Processing complete!"
+    # Using 'mp4v' as it's broadly available. For H.264 (libx264), moviepy is more reliable.
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out_nosound = cv2.VideoWriter(output_video_path_nosound, fourcc, fps if fps > 0 else 30.0, (frame_width, frame_height))
+
+    logging.info(f"Processing video: {target_video_path} ({total_frames if total_frames > 1 else 'unknown'} frames, {fps} FPS)")
+    frames_processed_count = 0
+    audio_extracted_path: Optional[str] = None
 
+    try: # Main processing block
+        # Attempt to extract audio first
+        progress(0.01, desc="Extracting audio (if any)...")
+        try:
+            from moviepy.editor import VideoFileClip
+            with VideoFileClip(target_video_path) as video_clip:
+                if video_clip.audio:
+                    audio_extracted_path = os.path.join(temp_process_dir, "original_audio.mp3")
+                    video_clip.audio.write_audiofile(audio_extracted_path, logger=None) # Suppress moviepy stdout
+                    logging.info(f"Original audio extracted to: {audio_extracted_path}")
+                else:
+                    logging.info("No audio track found in the original video.")
+        except ImportError:
+            logging.warning("moviepy library not found. Video will be processed without audio. Install with 'pip install moviepy'.")
+            gr.Info("moviepy not found. Video will be silent.")
+        except Exception as e_audio_extract:
+            logging.warning(f"Could not extract audio using moviepy: {e_audio_extract}. Proceeding without audio.", exc_info=True)
+            gr.Warning(f"Audio extraction failed: {e_audio_extract}. Video will be silent.")
+
+        # Frame-by-frame processing
+        # Use a manual loop if total_frames is unreliable
+        frame_num_iterator = range(total_frames) if total_frames > 1 else iter(lambda: cap.isOpened(), False)
+
+        for frame_idx in progress.tqdm(frame_num_iterator, desc="Swapping video frames", total=total_frames if total_frames > 1 else None):
+            ret, frame = cap.read()
+            if not ret:
+                logging.info(f"End of video or cannot read frame at index {frames_processed_count}.")
+                break
+            
+            current_frame_h, current_frame_w = frame.shape[:2]
+            target_faces_in_frame = get_faces_from_image(frame)
+            processed_frame = frame.copy()
+
+            if target_faces_in_frame and swapper:
+                target_face_to_swap_info_frame = None
+                if 0 <= target_face_index_hint < len(target_faces_in_frame):
+                    target_face_to_swap_info_frame = target_faces_in_frame[target_face_index_hint]
+                elif target_faces_in_frame: # Fallback: use the first detected face
+                    target_face_to_swap_info_frame = target_faces_in_frame[0]
+                
+                if target_face_to_swap_info_frame:
+                    try:
+                        processed_frame = swapper.get(processed_frame, target_face_to_swap_info_frame, source_face, paste_back=True)
+                        bbox_coords_frame = target_face_to_swap_info_frame.bbox.astype(int)
+
+                        if apply_enhancement and restoration_model_loaded_successfully and face_restorer:
+                            padding_enh = 20
+                            enh_x1 = max(0, bbox_coords_frame[0] - padding_enh)
+                            enh_y1 = max(0, bbox_coords_frame[1] - padding_enh)
+                            enh_x2 = min(current_frame_w, bbox_coords_frame[2] + padding_enh)
+                            enh_y2 = min(current_frame_h, bbox_coords_frame[3] + padding_enh)
+                            if enh_x1 < enh_x2 and enh_y1 < enh_y2:
+                                crop = processed_frame[enh_y1:enh_y2, enh_x1:enh_x2]
+                                enhanced = enhance_cropped_face(crop.copy())
+                                if enhanced.shape == crop.shape:
+                                    processed_frame[enh_y1:enh_y2, enh_x1:enh_x2] = enhanced
+                                # else: logging.debug(f"Frame {frames_processed_count}: Enhanced crop mismatch")
+                            # else: logging.debug(f"Frame {frames_processed_count}: Invalid enhancement crop")
+                        
+                        if apply_color_correction:
+                            cc_x1, cc_y1, cc_x2, cc_y2 = bbox_coords_frame
+                            if cc_x1 < cc_x2 and cc_y1 < cc_y2:
+                                original_region = frame[cc_y1:cc_y2, cc_x1:cc_x2]
+                                swapped_region = processed_frame[cc_y1:cc_y2, cc_x1:cc_x2]
+                                if original_region.size > 0 and swapped_region.size > 0:
+                                    corrected = histogram_match_color(swapped_region.copy(), original_region.copy())
+                                    if corrected.shape == swapped_region.shape:
+                                        processed_frame[cc_y1:cc_y2, cc_x1:cc_x2] = corrected
+                                # else: logging.debug(f"Frame {frames_processed_count}: Color correction region mismatch")
+                            # else: logging.debug(f"Frame {frames_processed_count}: Invalid color correction bbox")
+                    except Exception as e_frame_swap:
+                        logging.warning(f"Error swapping face in frame {frames_processed_count}: {e_frame_swap}", exc_info=False)
+                        # Continue with the (partially) processed frame or original if error is too severe
+
+            out_nosound.write(processed_frame)
+            frames_processed_count += 1
+            if total_frames <=1 and frames_processed_count % int(fps if fps > 0 else 30) == 0 : # Update progress for unknown total
+                 progress(frames_processed_count / (frames_processed_count +1), desc=f"Swapping video frames ({frames_processed_count} so far)")
+
+
+        cap.release()
+        out_nosound.release()
+        logging.info(f"Finished processing {frames_processed_count} frames. Silent video at: {output_video_path_nosound}")
+
+        # Re-mux audio using moviepy if audio was extracted
+        if audio_extracted_path and os.path.exists(output_video_path_nosound):
+            progress(0.95, desc="Combining video with audio...")
+            try:
+                from moviepy.editor import VideoFileClip, AudioFileClip
+                video_clip_nosound = VideoFileClip(output_video_path_nosound)
+                audio_clip = AudioFileClip(audio_extracted_path)
+                final_video_clip = video_clip_nosound.set_audio(audio_clip)
+                # Use a reliable codec like libx264 for video and aac for audio
+                final_video_clip.write_videofile(final_output_video_path, codec="libx264", audio_codec="aac", temp_audiofile_path=temp_process_dir, logger=None) # Suppress moviepy stdout
+                video_clip_nosound.close()
+                audio_clip.close()
+                logging.info(f"Video with audio successfully saved to: {final_output_video_path}")
+            except ImportError:
+                logging.warning("moviepy not found, cannot add audio back. Returning silent video.")
+                if os.path.exists(output_video_path_nosound) and not os.path.exists(final_output_video_path):
+                     os.rename(output_video_path_nosound, final_output_video_path) # Rename for consistency
+            except Exception as e_audio_mux:
+                logging.error(f"Error adding audio back to video: {e_audio_mux}. Returning silent video.", exc_info=True)
+                gr.Warning(f"Audio muxing failed: {e_audio_mux}. Resulting video may be silent.")
+                if os.path.exists(output_video_path_nosound) and not os.path.exists(final_output_video_path):
+                     os.rename(output_video_path_nosound, final_output_video_path)
+        elif os.path.exists(output_video_path_nosound) and not os.path.exists(final_output_video_path): # No audio to begin with or extraction failed
+            os.rename(output_video_path_nosound, final_output_video_path)
+            logging.info(f"No audio to add or audio extraction failed. Final video (silent) at: {final_output_video_path}")
+        
+        if os.path.exists(final_output_video_path):
+             progress(1.0, desc="🎉 Video swap complete!")
+             return final_output_video_path, f"🎉 Video processing complete! ({frames_processed_count} frames)"
+        else:
+            raise gr.Error("Video processing failed to produce a final output file.")
+
+    except Exception as e_video_proc:
+        cap.release() # Ensure cap is released on error
+        if out_nosound.isOpened(): out_nosound.release() # Ensure writer is released
+        logging.error(f"Unhandled error during video processing: {e_video_proc}", exc_info=True)
+        raise gr.Error(f"Video processing error: {str(e_video_proc)}")
+    finally:
+        # Clean up the temporary directory ONLY if final_output_video_path is NOT inside it
+        # OR if Gradio handles the file by copying it.
+        # Gradio usually copies output files. So, cleaning temp_process_dir should be okay.
+        # To be safe, one might copy the final_output_video_path to a Gradio managed temp location
+        # before deleting temp_process_dir. For now, we rely on Gradio's handling.
+        # If final_output_video_path is to be returned, it should not be deleted yet.
+        # Let Gradio handle the temp file. The `temp_process_dir` can be problematic if returned file is in it.
+        # A better approach for temp files given to Gradio is to use `tempfile.NamedTemporaryFile` and pass its name.
+        # For now, the video file will remain in temp_process_dir, and Gradio will serve it from there.
+        # It might get deleted when the app closes or based on OS temp cleanup.
+        # For robust cleanup, you'd manage this more carefully, possibly by copying final output
+        # to a location Gradio fully controls or a specific output directory.
+        # Let's assume Gradio copies it, and we can clean our main temp_process_dir later if needed.
+        # The current return path IS inside temp_process_dir. Gradio needs to access it.
+        # A solution: copy final_output_video_path to a NamedTemporaryFile before returning.
+        if os.path.exists(final_output_video_path):
+            final_gradio_video_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
+            shutil.copy2(final_output_video_path, final_gradio_video_path)
+            shutil.rmtree(temp_process_dir) # Now clean up the directory
+            logging.info(f"Temporary processing directory {temp_process_dir} cleaned up.")
+            return final_gradio_video_path, f"🎉 Video processing complete! ({frames_processed_count} frames)"
+        else: # If no final path, cleanup might still be useful if dir exists
+            if os.path.exists(temp_process_dir):
+                shutil.rmtree(temp_process_dir)
+                logging.info(f"Temporary processing directory {temp_process_dir} cleaned up after error.")
+            return None, "Video processing failed."
+
+
+# --- Gradio UI Functions ---
 def preview_target_faces(target_pil_img: Optional[Image.Image]):
     blank_preview_placeholder = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
     reset_slider = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False, label="🎯 Select Target Face")
-    if not core_models_loaded_successfully:
+    status_msg = ""
+
+    if not core_models_loaded_successfully or face_analyzer is None:
         gr.Warning("Face detection models not loaded. Preview cannot be generated.")
-        return blank_preview_placeholder, reset_slider
+        return blank_preview_placeholder, reset_slider, "Core models not loaded."
     if target_pil_img is None:
-        return blank_preview_placeholder, reset_slider
+        return blank_preview_placeholder, reset_slider, "Upload target image first."
+    
     target_np = convert_pil_to_cv2(target_pil_img)
     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
+        gr.Warning("Could not process target image for preview.")
+        return blank_preview_placeholder, reset_slider, "Error processing image."
+
     faces = get_faces_from_image(target_np)
     if not faces:
-        max_h, max_w = DETECTION_SIZE[1], DETECTION_SIZE[0]
+        # Resize image for preview if it's too large, even if no faces are found
+        max_h, max_w = DETECTION_SIZE[1], DETECTION_SIZE[0] # Assuming (width, height) for DETECTION_SIZE
         h, w = target_np.shape[:2]
+        preview_display_np = target_np
         if h > max_h or w > max_w:
             scale = min(max_h/h, max_w/w)
-            target_np_resized = cv2.resize(target_np, (int(w*scale), int(h*scale)))
-            preview_pil_img = convert_cv2_to_pil(target_np_resized)
-        else:
-            preview_pil_img = convert_cv2_to_pil(target_np)
+            preview_display_np = cv2.resize(target_np, (int(w*scale), int(h*scale)))
+        
+        preview_pil_img = convert_cv2_to_pil(preview_display_np)
         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)
+        return preview_pil_img, reset_slider, "No faces detected in image."
+
+    preview_np_with_boxes = draw_detected_faces(target_np.copy(), 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
         gr.Warning("Error generating preview image with face boxes.")
+        status_msg = "Error generating preview."
+    
     num_faces = len(faces)
     slider_update = gr.Slider(minimum=0, maximum=max(0, num_faces - 1), value=0, step=1,
                               label=f"🎯 Select Target Face (0 to {num_faces-1})", interactive=(num_faces > 0))
-    return preview_pil_img_with_boxes, slider_update
+    status_msg = status_msg if status_msg else f"Target preview updated. {num_faces} face(s) found."
+    return preview_pil_img_with_boxes, slider_update, status_msg
+
+
+def preview_video_first_frame(target_video_path: Optional[str]): # Path from gr.Video
+    blank_preview_placeholder = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
+    reset_slider = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False, label="🎯 Select Target Face")
+    status_msg = ""
+
+    if not core_models_loaded_successfully or face_analyzer is None:
+        gr.Warning("Face detection models not loaded. Video preview cannot be generated.")
+        return blank_preview_placeholder, reset_slider, "Core models not loaded."
+    if target_video_path is None:
+        return blank_preview_placeholder, reset_slider, "Upload target video first."
+
+    cap = cv2.VideoCapture(target_video_path)
+    if not cap.isOpened():
+        gr.Warning("Could not open video file for preview.")
+        return blank_preview_placeholder, reset_slider, "Error opening video."
+
+    ret, first_frame_np = cap.read()
+    cap.release()
+
+    if not ret or first_frame_np is None:
+        gr.Warning("Could not read the first frame from the video.")
+        return blank_preview_placeholder, reset_slider, "Error reading first frame."
+
+    faces = get_faces_from_image(first_frame_np)
+    if not faces:
+        max_h, max_w = DETECTION_SIZE[1], DETECTION_SIZE[0]
+        h, w = first_frame_np.shape[:2]
+        preview_display_np = first_frame_np
+        if h > max_h or w > max_w: # Resize for display
+            scale = min(max_h/h, max_w/w)
+            preview_display_np = cv2.resize(first_frame_np, (int(w*scale), int(h*scale)))
+        
+        preview_pil_img = convert_cv2_to_pil(preview_display_np)
+        if preview_pil_img is None: preview_pil_img = blank_preview_placeholder
+        gr.Info("No faces detected in the first frame of the video.")
+        return preview_pil_img, reset_slider, "No faces in first frame of video."
+
+    preview_np_with_boxes = draw_detected_faces(first_frame_np.copy(), 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
+        gr.Warning("Error generating preview image for video's first frame.")
+        status_msg = "Error generating video preview."
+
+    num_faces = len(faces)
+    slider_update = gr.Slider(minimum=0, maximum=max(0, num_faces - 1), value=0, step=1,
+                              label=f"🎯 Select Target Face in First Frame (0 to {num_faces-1})", interactive=(num_faces > 0))
+    status_msg = status_msg if status_msg else f"Video first frame preview updated. {num_faces} face(s) found."
+    return preview_pil_img_with_boxes, slider_update, status_msg
+
 
 # --- Gradio UI Definition ---
-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:
+with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.3 Video", 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.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>
+            <h1>💎 Ultimate Face Swap AI <span style="font-size:0.8em; color: #555;">v3.3 Video</span> ✨</h1>
+            <p>Swap faces in <strong>images</strong> or <strong>videos</strong>!</p>
+            <p style="font-size:0.9em;">Optionally enhance with <strong>face restoration</strong> and <strong>color correction</strong>.</p>
         </div>
         """
     )
@@ -360,137 +631,181 @@ with gr.Blocks(title="Ultimate Face Swap AI 💎 v3.2 Lite", theme=gr.themes.Sof
         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 (e.g., model file paths, ONNX Runtime issues) "
-            "and restart the application after fixing the underlying problem."
+            "Please check console logs and restart after fixing."
         )
 
-    with gr.Row():
-        with gr.Column(scale=1, min_width=300):
-            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)
-
-    gr.HTML("<hr style='margin-top: 10px; margin-bottom: 15px; border-style: dashed; border-color: #ccc;'>")
-
-    with gr.Row(equal_height=False):
-        with gr.Column(scale=2):
-            target_faces_preview_output = gr.Image(label="👀 Target Image Preview & Detected Faces", interactive=False, height=380, show_download_button=False)
-        with gr.Column(scale=1, min_width=250):
-            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
-            )
-            gr.Markdown("<p style='font-size:0.8em; color:#555;'>After uploading target, click 'Preview'. Then use slider to select face.</p>")
+    with gr.Tabs() as tabs:
+        # --- IMAGE SWAP TAB ---
+        with gr.TabItem("🖼️ Image Face Swap", id="image_tab"):
+            with gr.Row():
+                with gr.Column(scale=1, min_width=300):
+                    img_source_image_input = gr.Image(label="👤 Source Face Image (Clear, single face)", type="pil", sources=["upload", "clipboard"], height=380)
+                with gr.Column(scale=1, min_width=300):
+                    img_target_image_input = gr.Image(label="🖼️ Target Scene Image", type="pil", sources=["upload", "clipboard"], height=380)
             
-            # 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;'>")
+            gr.HTML("<hr style='margin-top: 10px; margin-bottom: 15px; border-style: dashed; border-color: #ccc;'>")
+
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=2):
+                    img_target_faces_preview_output = gr.Image(label="👀 Target Image Preview & Detected Faces", interactive=False, height=380, show_download_button=False)
+                with gr.Column(scale=1, min_width=250):
+                    img_preview_button = gr.Button("🔍 Preview & Select Target Face (Image)", variant="secondary")
+                    img_face_index_slider = gr.Slider(
+                        label="🎯 Select Target Face (0-indexed)", minimum=0, maximum=0, step=1, value=0, interactive=False
+                    )
+                    gr.Markdown("<p style='font-size:0.8em; color:#555;'>Upload target, click 'Preview', then use slider.</p>")
+                    gr.HTML("<hr style='margin-top: 5px; margin-bottom: 5px; border-style: solid; border-width: thin; border-color: #ddd;'>")
+                    
+                    img_enhance_checkbox_label = "✨ Apply Face Restoration"
+                    if not restoration_model_loaded_successfully: img_enhance_checkbox_label += " (Model N/A)"
+                    img_enhance_checkbox = gr.Checkbox(label=img_enhance_checkbox_label, value=restoration_model_loaded_successfully, interactive=restoration_model_loaded_successfully)
+                    if not restoration_model_loaded_successfully: gr.Markdown("<p style='color: orange; font-size:0.8em;'>⚠️ Face restoration model N/A.</p>")
+                    else: gr.Markdown("<p style='font-size:0.8em; color:#555;'>Improves swapped face quality.</p>")
+
+                    img_color_correction_checkbox = gr.Checkbox(label="🎨 Apply Color Correction", value=True)
+                    gr.Markdown("<p style='font-size:0.8em; color:#555;'>Matches swapped face color to scene.</p>")
+
+            gr.HTML("<hr style='margin-top: 15px; margin-bottom: 15px;'>")
+            with gr.Row():
+                img_swap_button = gr.Button("🚀 GENERATE IMAGE SWAP!", variant="primary", scale=3, interactive=core_models_loaded_successfully)
+                img_clear_button = gr.Button("🧹 Clear Image Inputs", variant="stop", 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,
-                interactive=restoration_model_loaded_successfully
+            with gr.Row():
+                with gr.Column(scale=2):
+                    img_swapped_image_output = gr.Image(label="✨ Swapped Image Result", interactive=False, height=480, show_download_button=False)
+                with gr.Column(scale=1):
+                    img_download_output_file = gr.File(label="⬇️ Download Swapped Image", interactive=True)
+                    img_status_message_output = gr.Markdown(value="<p style='color: #555; font-style: italic;'>Image status messages.</p>")
+
+            # Image Tab Event Handlers
+            def on_img_target_image_upload_or_clear(target_img_pil: Optional[Image.Image]):
+                blank_preview = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='whitesmoke')
+                reset_slider = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False, label="🎯 Select Target Face")
+                status = "Target image changed. Click 'Preview' to update." if target_img_pil else "Target image cleared."
+                return blank_preview, reset_slider, status
+
+            img_target_image_input.change(
+                fn=on_img_target_image_upload_or_clear,
+                inputs=[img_target_image_input],
+                outputs=[img_target_faces_preview_output, img_face_index_slider, img_status_message_output],
+                queue=False
             )
-            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. 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. Requires restoration model.</p>")
-
-            color_correction_checkbox = gr.Checkbox(
-                label="🎨 Apply Color Correction",
-                value=True
+            img_preview_button.click(
+                fn=preview_target_faces, inputs=[img_target_image_input],
+                outputs=[img_target_faces_preview_output, img_face_index_slider, img_status_message_output],
+                show_progress="full"
+            )
+            img_swap_button.click(
+                fn=process_face_swap,
+                inputs=[img_source_image_input, img_target_image_input, img_face_index_slider, img_enhance_checkbox, img_color_correction_checkbox],
+                outputs=[img_swapped_image_output, img_download_output_file, img_status_message_output]
+            )
+            def clear_image_tab():
+                blank_preview = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
+                reset_slider = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
+                return None, None, blank_preview, reset_slider, None, None, "<p style='color: #555; font-style: italic;'>Image fields cleared.</p>"
+            img_clear_button.click(
+                fn=clear_image_tab, inputs=None,
+                outputs=[img_source_image_input, img_target_image_input, img_target_faces_preview_output, img_face_index_slider, img_swapped_image_output, img_download_output_file, img_status_message_output],
+                queue=False
             )
-            gr.Markdown("<p style='font-size:0.8em; color:#555;'>Matches color tone of swapped face to the target scene.</p>")
-
-
-    gr.HTML("<hr style='margin-top: 15px; margin-bottom: 15px;'>")
-
-    with gr.Row():
-        swap_button = gr.Button("🚀 GENERATE SWAP!", variant="primary", scale=3, interactive=core_models_loaded_successfully)
-        clear_button = gr.Button("🧹 Clear All", variant="stop", scale=1)
-
-    with gr.Row():
-        with gr.Column(scale=2):
-            swapped_image_output = gr.Image(label="✨ Swapped Result", interactive=False, height=480, show_download_button=False)
-        with gr.Column(scale=1):
-            download_output_file = gr.File(label="⬇️ Download Swapped Image", interactive=True)
-            status_message_output = gr.Markdown(value="<p style='color: #555; font-style: italic;'>Status messages will appear here.</p>")
-
-    def on_target_image_upload_or_clear(target_img_pil: Optional[Image.Image]):
-        blank_preview = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='whitesmoke')
-        reset_slider = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False, label="🎯 Select Target Face")
-        status_update = "Target image changed. Click 'Preview & Select Target Face' to proceed." if target_img_pil else "Target image cleared."
-        return blank_preview, reset_slider, status_update
-
-    target_image_input.change(
-        fn=on_target_image_upload_or_clear,
-        inputs=[target_image_input],
-        outputs=[target_faces_preview_output, face_index_slider, status_message_output],
-        queue=False
-    )
-    
-    preview_button.click(
-        fn=preview_target_faces,
-        inputs=[target_image_input],
-        outputs=[target_faces_preview_output, face_index_slider],
-        show_progress="full" 
-    ).then(lambda: "Target preview updated. Select a face if multiple are shown.", outputs=[status_message_output])
-
-    swap_button.click(
-        fn=process_face_swap,
-        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]
-    )
 
-    def clear_all_inputs_outputs():
-        blank_preview = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
-        reset_slider_dict = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
-        return (
-            None, None, blank_preview, reset_slider_dict, None, None,
-            "<p style='color: #555; font-style: italic;'>All fields cleared. Ready for new input.</p>"
-        )
+        # --- VIDEO SWAP TAB ---
+        with gr.TabItem("🎬 Video Face Swap", id="video_tab"):
+            gr.Markdown("<p style='font-size:0.9em; color:orange;'>⚠️ Video processing can be time-consuming. Audio is preserved if <strong>moviepy</strong> is installed (`pip install moviepy`).</p>")
+            with gr.Row():
+                with gr.Column(scale=1, min_width=300):
+                    vid_source_image_input = gr.Image(label="👤 Source Face Image (for video)", type="pil", sources=["upload", "clipboard"], height=380)
+                with gr.Column(scale=1, min_width=300):
+                    vid_target_video_input = gr.Video(label="🎬 Target Video File", sources=["upload"], height=380)
+            
+            gr.HTML("<hr style='margin-top: 10px; margin-bottom: 15px; border-style: dashed; border-color: #ccc;'>")
+
+            with gr.Row(equal_height=False):
+                with gr.Column(scale=2):
+                    vid_target_preview_output = gr.Image(label="👀 Target Video First Frame Preview", interactive=False, height=380, show_download_button=False)
+                with gr.Column(scale=1, min_width=250):
+                    vid_preview_button = gr.Button("🔍 Preview First Frame & Select Target Face (Video)", variant="secondary")
+                    vid_face_index_slider = gr.Slider(
+                        label="🎯 Select Target Face in First Frame (0-indexed)", minimum=0, maximum=0, step=1, value=0, interactive=False
+                    )
+                    gr.Markdown("<p style='font-size:0.8em; color:#555;'>Upload video, click 'Preview', then use slider for first frame.</p>")
+                    gr.HTML("<hr style='margin-top: 5px; margin-bottom: 5px; border-style: solid; border-width: thin; border-color: #ddd;'>")
+
+                    vid_enhance_checkbox_label = "✨ Apply Face Restoration (Video)"
+                    if not restoration_model_loaded_successfully: vid_enhance_checkbox_label += " (Model N/A)"
+                    vid_enhance_checkbox = gr.Checkbox(label=vid_enhance_checkbox_label, value=restoration_model_loaded_successfully, interactive=restoration_model_loaded_successfully)
+                    if not restoration_model_loaded_successfully: gr.Markdown("<p style='color: orange; font-size:0.8em;'>⚠️ Face restoration model N/A.</p>")
+                    
+                    vid_color_correction_checkbox = gr.Checkbox(label="🎨 Apply Color Correction (Video)", value=True)
+
+            gr.HTML("<hr style='margin-top: 15px; margin-bottom: 15px;'>")
+            with gr.Row():
+                vid_swap_button = gr.Button("🚀 GENERATE VIDEO SWAP!", variant="primary", scale=3, interactive=core_models_loaded_successfully)
+                vid_clear_button = gr.Button("🧹 Clear Video Inputs", variant="stop", scale=1)
+
+            with gr.Row():
+                with gr.Column(scale=2): # Changed scale for better layout
+                    vid_swapped_video_output = gr.Video(label="✨ Swapped Video Result", interactive=False, height=480, show_download_button=True)
+                with gr.Column(scale=1):
+                     vid_status_message_output = gr.Markdown(value="<p style='color: #555; font-style: italic;'>Video status messages.</p>")
+
+            # Video Tab Event Handlers
+            def on_vid_target_video_upload_or_clear(target_video_file: Optional[str]):
+                blank_preview = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='whitesmoke')
+                reset_slider = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False, label="🎯 Select Target Face")
+                status = "Target video changed. Click 'Preview First Frame' to update." if target_video_file else "Target video cleared."
+                return blank_preview, reset_slider, status
+
+            vid_target_video_input.change(
+                fn=on_vid_target_video_upload_or_clear,
+                inputs=[vid_target_video_input],
+                outputs=[vid_target_preview_output, vid_face_index_slider, vid_status_message_output],
+                queue=False # No need to queue for simple UI updates
+            )
+            vid_preview_button.click(
+                fn=preview_video_first_frame, inputs=[vid_target_video_input],
+                outputs=[vid_target_preview_output, vid_face_index_slider, vid_status_message_output],
+                show_progress="full"
+            )
+            vid_swap_button.click(
+                fn=process_video_face_swap,
+                inputs=[vid_source_image_input, vid_target_video_input, vid_face_index_slider, vid_enhance_checkbox, vid_color_correction_checkbox],
+                outputs=[vid_swapped_video_output, vid_status_message_output] # Output is video path and status
+            )
+            def clear_video_tab():
+                blank_preview = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
+                reset_slider = gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
+                return None, None, blank_preview, reset_slider, None, "<p style='color: #555; font-style: italic;'>Video fields cleared.</p>"
+            vid_clear_button.click(
+                fn=clear_video_tab, inputs=None,
+                outputs=[vid_source_image_input, vid_target_video_input, vid_target_preview_output, vid_face_index_slider, vid_swapped_video_output, vid_status_message_output],
+                queue=False
+            )
 
-    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,
-            status_message_output
-        ],
-        queue=False
-    )
-    
     gr.HTML("<hr style='margin-top: 20px; margin-bottom: 10px;'>")
-    gr.Markdown("### 📋 Example Usage")
-    
-    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.Markdown("### 📋 Example Usage (Image Tab)")
     gr.Examples(
         examples=[
             ["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 (Click to run - ensure example images exist in 'examples' folder)"
+        inputs=[img_source_image_input, img_target_image_input, img_face_index_slider, img_enhance_checkbox, img_color_correction_checkbox],
+        outputs=[img_swapped_image_output, img_download_output_file, img_status_message_output],
+        fn=process_face_swap, # Using the image swap function
+        cache_examples=False, # Re-run for dynamic results, set to True if examples are static
+        label="Image Swap Examples (Ensure 'examples' folder and images exist)"
     )
+    # Note: gr.Examples for video is more complex due to large file inputs.
+    # It's often omitted or handled by providing downloadable example video files and instructions.
 
 if __name__ == "__main__":
     os.makedirs("models", exist_ok=True)
     os.makedirs("examples", exist_ok=True)
+    # Consider adding example images if they don't exist for gr.Examples to work out of the box.
 
     print("\n" + "="*70)
-    print("🚀 ULTIMATE FACE SWAP AI - v3.2 Lite STARTUP STATUS 🚀") # Version bump
+    print("🚀 ULTIMATE FACE SWAP AI - v3.3 Video STARTUP STATUS 🚀")
     print("="*70)
     print(f"Execution Providers Selected: {EXECUTION_PROVIDERS}")
     
@@ -498,18 +813,25 @@ if __name__ == "__main__":
         print("\n🔴 CRITICAL ERROR: CORE MODELS FAILED TO LOAD.")
         print(f"  - Face Analyzer ('{FACE_ANALYZER_NAME}'): {'Loaded' if face_analyzer else 'FAILED'}")
         print(f"  - Swapper Model ('{SWAPPER_MODEL_PATH}'): {'Loaded' if swapper else 'FAILED'}")
-        print("  The application UI will be severely limited or non-functional.")
     else:
         print("\n🟢 Core models (Face Analyzer & Swapper) loaded successfully.")
+    
     if not restoration_model_loaded_successfully:
-        print(f"\n🟡 INFO: Face Restoration model ('{RESTORATION_MODEL_PATH}') NOT loaded.")
-        print("  The 'Apply Selective Face Restoration' feature will be DISABLED.")
+        print(f"\n🟡 INFO: Face Restoration model ('{RESTORATION_MODEL_PATH}') NOT loaded. Enhancement will be disabled.")
     else:
         print("\n🟢 Face Restoration model loaded successfully.")
+    
+    try:
+        import moviepy.editor
+        print("\n🟢 moviepy library found. Audio processing for videos is ENABLED.")
+    except ImportError:
+        print("\n🟡 WARNING: moviepy library NOT found. Audio will NOT be processed for videos.")
+        print("   To enable audio in video swaps, please install it: pip install moviepy")
     print("="*70)
+
     if not core_models_loaded_successfully:
-        print("\n👉 The Gradio interface will launch, but swapping functionality will be broken. Please address errors.")
+        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).")
+        print("\n👉 Launching Gradio Interface... Access it in your browser (usually at http://127.0.0.1:7860 or a public link if shared).")
     
     demo.launch()