Update app.py

app.py

@@ -2,8 +2,8 @@ import os
 import cv2
 import numpy as np
 import gradio as gr
-from insightface.app import FaceAnalysis
-from insightface.model_zoo import get_model
+from insightface.app import FaceAnalysis # Still needed for detection and embeddings
+# from insightface.model_zoo import get_model # No longer using this for the swapper
 from PIL import Image
 import tempfile
 import logging
@@ -12,66 +12,71 @@ import onnxruntime
 # --- Configuration & Setup ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-SWAPPER_MODEL_PATH = "models/inswapper_128.onnx"
-RESTORATION_MODEL_PATH = "models/gfpgan_1.4.onnx" # Path to your GFPGAN or chosen restoration ONNX model
+# MODIFIED: Path to the new swapper model
+SWAPPER_MODEL_PATH = "models/reswapper_256.onnx" # <--- USING RESWAPPER_256
+RESTORATION_MODEL_PATH = "models/gfpgan_1.4.onnx" # Or your chosen restorer, e.g., restoreformer++.onnx
+# UPSCALER_MODEL_PATH = "models/RealESRGANx4plus.onnx" # Path for AI Upscaler if you were to add it
 
-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
+FACE_ANALYZER_NAME = 'buffalo_l' # InsightFace detection and embedding model
+DETECTION_SIZE = (640, 640)
+EXECUTION_PROVIDERS = ['CPUExecutionProvider'] # Or ['CUDAExecutionProvider', 'CPUExecutionProvider']
 
 # --- Global Variables (Lazy Loaded by initialize_models) ---
 face_analyzer = None
-swapper = None
+reswapper_session = None # MODIFIED: Will hold the onnxruntime session for reswapper
 face_restorer = None
+# face_upscaler = None # If you add AI upscaling
 
 # --- Initialization Functions ---
 def initialize_models():
-    global face_analyzer, swapper, face_restorer
+    global face_analyzer, reswapper_session, face_restorer # MODIFIED: reswapper_session
     try:
-        # Initialize FaceAnalysis model
+        # Initialize FaceAnalysis model (for detection and embeddings)
         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) # ctx_id=0 for CPU, -1 for auto
+            face_analyzer.prepare(ctx_id=0, det_size=DETECTION_SIZE)
             logging.info("FaceAnalysis model initialized successfully.")
 
-        # Initialize Swapper model
-        if swapper is None:
+        # MODIFIED: Initialize ReSwapper model directly with onnxruntime
+        if reswapper_session 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.")
+                logging.error(f"ReSwapper model FILE NOT FOUND at {SWAPPER_MODEL_PATH}. Swapping will fail.")
             else:
-                logging.info(f"Loading swapper model from: {SWAPPER_MODEL_PATH}")
+                logging.info(f"Loading ReSwapper 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) # Fallback
-                logging.info("Swapper model loaded successfully.")
-
-        # Initialize Face Restoration Model
-        if face_restorer is None: # Only attempt to load if not already loaded/failed
+                    reswapper_session = onnxruntime.InferenceSession(SWAPPER_MODEL_PATH, providers=EXECUTION_PROVIDERS)
+                    logging.info(f"ReSwapper model ({SWAPPER_MODEL_PATH}) loaded successfully with onnxruntime.")
+                except Exception as e:
+                    logging.error(f"Error loading ReSwapper model {SWAPPER_MODEL_PATH} with onnxruntime: {e}", exc_info=True)
+                    reswapper_session = None # Ensure it's None if loading failed
+        
+        # Initialize Face Restoration Model (same as before)
+        if face_restorer is None:
             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.")
+                logging.error(f"Face restoration model FILE NOT FOUND at: {RESTORATION_MODEL_PATH}. Enhancement will be disabled.")
             else:
+                # (loading logic for face_restorer as in previous complete code)
                 logging.info(f"Attempting to load face restoration model from: {RESTORATION_MODEL_PATH}")
                 try:
                     face_restorer = onnxruntime.InferenceSession(RESTORATION_MODEL_PATH, providers=EXECUTION_PROVIDERS)
                     logging.info("Face restoration model loaded successfully.")
                 except Exception as e:
                     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
+                    face_restorer = None
+        
+        # (Initialize AI Face Upscaler Model - if you were adding it)
+
     except Exception as e:
         logging.error(f"A critical error occurred during model initialization: {e}", exc_info=True)
-        # face_analyzer, swapper, or face_restorer might be None. Subsequent checks will handle this.
 
 # --- Call Initialization Early ---
 initialize_models()
-# 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
+core_models_loaded_successfully = face_analyzer is not None and reswapper_session is not None # MODIFIED
 restoration_model_loaded_successfully = face_restorer is not None
+# upscaler_model_loaded_successfully = face_upscaler is not None # If adding upscaler
 
-# --- Image Utility Functions ---
+# --- Image Utility Functions (convert_pil_to_cv2, convert_cv2_to_pil - same as before) ---
 def convert_pil_to_cv2(pil_image: Image.Image) -> np.ndarray | None:
     if pil_image is None: return None
     try:
@@ -88,139 +93,132 @@ def convert_cv2_to_pil(cv2_image: np.ndarray) -> Image.Image | None:
         logging.error(f"Error converting CV2 to PIL: {e}")
         return None
 
-# --- Core AI & Image Processing Functions ---
+# --- Core AI & Image Processing Functions (get_faces_from_image, draw_detected_faces, enhance_cropped_face, histogram_match_color, apply_naturalness_filters - mostly same) ---
+# These functions (except the main swap logic) should largely remain the same.
+# Make sure they are copied from your previous complete version.
+# For brevity here, I'll assume they are present.
+# Example of get_faces_from_image:
 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.
+    if face_analyzer is None: # Check against the global face_analyzer
         logging.error("Face analyzer not available for get_faces_from_image.")
         return []
     if img_np is None: return []
     try:
-        return face_analyzer.get(img_np)
+        return face_analyzer.get(img_np) # This provides Face objects with embeddings
     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) -> np.ndarray:
+# ... (draw_detected_faces, enhance_cropped_face, histogram_match_color, apply_naturalness_filters functions from previous full code)
+def draw_detected_faces(img_np: np.ndarray, faces: list) -> np.ndarray: # Copied for completeness
     img_with_boxes = img_np.copy()
     for i, face in enumerate(faces):
-        box = face.bbox.astype(int) # InsightFace bbox is [x1, y1, x2, y2]
+        box = face.bbox.astype(int)
         x1, y1, x2, y2 = box[0], box[1], box[2], box[3]
         cv2.rectangle(img_with_boxes, (x1, y1), (x2, y2), (0, 255, 0), 2)
-        label_position = (x1, max(0, y1 - 10)) # Ensure label is not drawn outside top
+        label_position = (x1, max(0, y1 - 10))
         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:
+def enhance_cropped_face(face_crop_bgr: np.ndarray) -> np.ndarray: # Copied for completeness
     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 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]
-
     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
+        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) # 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
-
+        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
-        
         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_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)
-
-        # 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 for crop: {e}", exc_info=True)
-        return face_crop_bgr # Return original crop on error
+        return face_crop_bgr
 
-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."""
+def histogram_match_channel(source_channel: np.ndarray, target_channel: np.ndarray) -> np.ndarray: # Copied
     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_hist, bins = np.histogram(source_channel_flat, 256, [0,256], density=True) # Use density for CDF
+    target_hist, bins = np.histogram(target_channel_flat, 256, [0,256], density=True)
     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
+    # Normalize CDFs to range [0, 255] for lookup table
+    source_cdf_norm = (source_cdf * 255 / source_cdf[-1]).astype(np.uint8) # Ensure last val is 255
+    target_cdf_norm = (target_cdf * 255 / target_cdf[-1]).astype(np.uint8)
     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
+    idx = 0
+    for i in range(256):
+        while idx < 255 and target_cdf_norm[idx] < source_cdf_norm[i]:
+            idx += 1
+        lookup_table[i] = idx
+    matched_channel_flat = cv2.LUT(source_channel, 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).
-    """
+def histogram_match_color(source_img: np.ndarray, target_img: np.ndarray) -> np.ndarray: # Copied
     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
+        for i in range(source_img.shape[2]):
             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
+        return source_img
 
+def apply_naturalness_filters(face_region_bgr: np.ndarray, noise_level: int = 3) -> np.ndarray: # Copied
+    if face_region_bgr is None or face_region_bgr.size == 0:
+        logging.warning("Cannot apply naturalness filters to empty or None region.")
+        return face_region_bgr if face_region_bgr is not None else np.array([])
+    processed_region = face_region_bgr.copy()
+    try:
+        processed_region = cv2.medianBlur(processed_region, 3)
+        if noise_level > 0:
+            gaussian_noise = np.random.normal(0, noise_level, processed_region.shape).astype(np.int16)
+            noisy_region_int16 = processed_region.astype(np.int16) + gaussian_noise
+            processed_region = np.clip(noisy_region_int16, 0, 255).astype(np.uint8)
+        logging.info("Applied naturalness filters to face region.")
+    except Exception as e:
+        logging.error(f"Error applying naturalness filters: {e}", exc_info=True)
+        return face_region_bgr
+    return processed_region
 
-# --- Gradio Interface Functions ---
+# --- Main Processing Function ---
 def process_face_swap(source_pil_img: Image.Image, target_pil_img: Image.Image,
-                        target_face_index: int, apply_enhancement: bool, apply_color_correction: bool,
+                        target_face_index: int, # apply_ai_upscaling: bool, (Removed for this example to focus on ReSwapper)
+                        apply_enhancement: bool, apply_color_correction: bool, 
+                        apply_naturalness: bool,
                         progress=gr.Progress(track_tqdm=True)):
+    
     progress(0, desc="Initializing process...")
-
-    if not core_models_loaded_successfully:
-        gr.Error("CRITICAL: Core models (Face Analyzer or Swapper) not loaded. Cannot proceed.")
-        return Image.new('RGB', (100,100), color='lightgrey'), None # Placeholder for output
+    if not core_models_loaded_successfully: # Checks face_analyzer and reswapper_session
+        gr.Error("CRITICAL: Core models (Face Analyzer or ReSwapper) not loaded. Cannot proceed.")
+        return Image.new('RGB', (100,100), color='lightgrey'), None
 
     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...")
-    source_np = convert_pil_to_cv2(source_pil_img)
-    target_np = convert_pil_to_cv2(target_pil_img)
-
+    source_np = convert_pil_to_cv2(source_pil_img) # BGR
+    target_np = convert_pil_to_cv2(target_pil_img) # BGR
     if source_np is None or target_np is None:
         raise gr.Error("Image conversion failed. Please try different images.")
     
@@ -228,98 +226,204 @@ def process_face_swap(source_pil_img: Image.Image, target_pil_img: Image.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. Please use a clear image of a single face.")
-    source_face = source_faces[0] # Assuming the first detected face is the one to use
+    if not source_faces: raise gr.Error("No face found in the source image.")
+    source_face = source_faces[0] # InsightFace Face object
 
     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)):
-        # 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
+        raise gr.Error(f"Selected target face index ({target_face_index}) is out of range.")
+    target_face_to_swap_info = target_faces[int(target_face_index)] # InsightFace Face object
+
+    swapped_bgr_img = target_np.copy() # Initialize with a copy
 
-    swapped_bgr_img = target_np.copy() # Start with a copy of target
     try:
-        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)
+        progress(0.4, desc="Preparing inputs for ReSwapper...")
         
-        if apply_enhancement:
-            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("Enhanced crop size mismatch. Skipping paste-back for enhancement.")
-                else:
-                    logging.warning("Skipping enhancement, invalid crop dimensions after padding.")
-            else:
-                gr.Info("Face restoration model not available, enhancement skipped by process.")
-                logging.warning("Enhancement requested at runtime but face restorer is not available.")
+        ########## USER VERIFICATION NEEDED FOR RESWAPPER_256.ONNX ##########
+        # 1. Source Face Embedding Preparation
+        source_embedding = source_face.embedding.astype(np.float32) # (512,)
+        # Most ONNX models expect a batch dimension.
+        source_embedding_tensor = np.expand_dims(source_embedding, axis=0) # (1, 512)
+
+        # 2. Target Image Preparation
+        # ASSUMPTION: reswapper_256.onnx might expect a fixed size input, e.g., 256x256.
+        # If it does, you MUST resize target_np to that size.
+        # For this example, let's assume it takes the target image and processes internally
+        # or expects a specific format. A common pattern is RGB, normalized, CHW.
+        # Let's assume input size of 256x256 for the target image input to the model.
+        # We will process the whole target_np, let the model do its work, and it should return the modified whole target_np.
+        
+        target_img_for_onnx = cv2.cvtColor(target_np, cv2.COLOR_BGR2RGB) # Often RGB for ONNX
+        
+        # RESIZE THE TARGET IMAGE IF THE MODEL EXPECTS A FIXED INPUT SIZE
+        # E.g., if reswapper_256.onnx expects 256x256 input for the *target image itself*:
+        # target_img_for_onnx = cv2.resize(target_img_for_onnx, (256, 256), interpolation=cv2.INTER_AREA)
+        # This is a CRITICAL assumption. If it operates on full-res target, no resize needed here.
+        # If the "256" in reswapper_256 refers to an internal patch size, it might handle varied input.
+        # For now, let's NOT resize the full target image, assuming the model handles it.
+
+        target_img_normalized = (target_img_for_onnx / 255.0).astype(np.float32) # Normalize [0, 1]
+        # Or normalize to [-1, 1]: (target_img_for_onnx / 127.5 - 1.0).astype(np.float32)
+        
+        target_img_chw = np.transpose(target_img_normalized, (2, 0, 1)) # HWC to CHW
+        target_image_tensor = np.expand_dims(target_img_chw, axis=0) # Add batch dimension
+
+        # 3. Target Face Bounding Box (Optional, but many swappers need it)
+        # The InsightFace Face object (target_face_to_swap_info) contains .bbox
+        # bbox is [x1, y1, x2, y2]. Normalize if model expects normalized bbox.
+        target_bbox_tensor = target_face_to_swap_info.bbox.astype(np.float32) 
+        # Reshape if needed, e.g., (1, 4)
+        target_bbox_tensor = np.expand_dims(target_bbox_tensor, axis=0) 
+
+
+        # Placeholder names for ONNX model inputs. YOU MUST FIND THE REAL NAMES.
+        # Common names could be "target", "source_embed", "target_face_box"
+        # Check model's .get_inputs() if loaded, or use Netron app to view ONNX.
+        onnx_input_names = [inp.name for inp in reswapper_session.get_inputs()]
+        # Crude assumption based on common patterns (highly likely to be WRONG):
+        # This is where you absolutely need the model's documentation.
+        # Let's assume: input0=target_image, input1=source_embedding, input2=target_bbox (OPTIONAL)
+        # This is a GUESS.
+        
+        # For this example, let's try a simpler input signature if possible,
+        # like what an insightface swapper's .get() method might abstract.
+        # A common pattern for standalone ONNX swappers is:
+        # Input 1: Target image tensor (e.g., BCHW, float32, normalized)
+        # Input 2: Source embedding tensor (e.g., B x EmbDim, float32)
+        # Some might also take target face landmarks or a mask.
+
+        # Let's assume these input names based on typical conventions.
+        # YOU MUST VERIFY THESE using Netron or the model's documentation.
+        # input_feed = {
+        #     "target_image": target_image_tensor,      # Placeholder name
+        #     "source_embedding": source_embedding_tensor # Placeholder name
+        # }
+        # If target_bbox is also an input:
+        # input_feed["target_bbox"] = target_bbox_tensor # Placeholder name
+
+        # A more robust way is to check input names from the model:
+        # input_name_target_img = reswapper_session.get_inputs()[0].name
+        # input_name_source_emb = reswapper_session.get_inputs()[1].name
+        # input_feed = {
+        #     input_name_target_img: target_image_tensor,
+        #     input_name_source_emb: source_embedding_tensor
+        # }
+        # if len(reswapper_session.get_inputs()) > 2:
+        #     input_name_target_bbox = reswapper_session.get_inputs()[2].name
+        #     input_feed[input_name_target_bbox] = target_bbox_tensor
+            
+        # For this illustrative code, I will make up plausible input names.
+        # ***** REPLACE "actual_target_input_name", "actual_source_embed_input_name" etc. *****
+        # ***** WITH THE REAL TENSOR NAMES FROM reswapper_256.onnx. *****
+        # ***** Common names are often like 'input', 'target', 'source', 'embedding'.*****
+        # ***** Use Netron app to inspect your reswapper_256.onnx file! *****
+        input_feed = {
+             # Example, assuming first input is target image, second is source embedding
+            reswapper_session.get_inputs()[0].name: target_image_tensor,
+            reswapper_session.get_inputs()[1].name: source_embedding_tensor
+        }
+        # If your model takes a third input like target_bbox (very model specific):
+        if len(reswapper_session.get_inputs()) > 2: # Check if there's a third input
+             input_feed[reswapper_session.get_inputs()[2].name] = target_bbox_tensor
+
+
+        # Placeholder for output tensor name. YOU MUST FIND THE REAL NAME.
+        # output_name_placeholder = "output_image" 
+        # More robust:
+        output_name = reswapper_session.get_outputs()[0].name
+        ########## END OF USER VERIFICATION NEEDED ##########
+        
+        progress(0.5, desc="Running ReSwapper model...")
+        onnx_outputs = reswapper_session.run([output_name], input_feed)
+        swapped_image_tensor = onnx_outputs[0] # Assuming first output is the image
+
+        progress(0.6, desc="Post-processing ReSwapper output...")
+        # Post-process the output tensor back to a BGR image
+        # 1. Remove batch dimension if present
+        if swapped_image_tensor.ndim == 4:
+            swapped_image_tensor = np.squeeze(swapped_image_tensor, axis=0) # (C, H, W)
+        # 2. Transpose CHW to HWC
+        swapped_image_hwc = np.transpose(swapped_image_tensor, (1, 2, 0)) # (H, W, C)
+        # 3. Denormalize (inverse of the normalization applied earlier)
+        # If normalized to [0, 1]:
+        swapped_image_denormalized = swapped_image_hwc * 255.0
+        # If normalized to [-1, 1]:
+        # swapped_image_denormalized = (swapped_image_hwc + 1.0) * 127.5
+        
+        swapped_image_uint8 = np.clip(swapped_image_denormalized, 0, 255).astype(np.uint8)
+        
+        # 4. Convert RGB to BGR if model outputs RGB
+        swapped_bgr_img = cv2.cvtColor(swapped_image_uint8, cv2.COLOR_RGB2BGR)
+
+        # IF `reswapper_256.onnx` returned a resized image (e.g. 256x256),
+        # you might need to resize `swapped_bgr_img` back to original target_h, target_w.
+        if swapped_bgr_img.shape[0] != target_h or swapped_bgr_img.shape[1] != target_w:
+            logging.info(f"Resizing ReSwapper output from {swapped_bgr_img.shape[:2]} to {(target_h, target_w)}")
+            swapped_bgr_img = cv2.resize(swapped_bgr_img, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4)
         
-        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.")
+        # --- Post-processing pipeline (Restoration, Color Correction, Naturalness) ---
+        # These will now operate on `swapped_bgr_img` which came from reswapper.
+        current_processed_image = swapped_bgr_img.copy()
+        bbox_coords = target_face_to_swap_info.bbox.astype(int) # [x1, y1, x2, y2] for ROI processing
+        roi_x1, roi_y1, roi_x2, roi_y2 = bbox_coords[0], bbox_coords[1], bbox_coords[2], bbox_coords[3]
+        
+        face_roi_for_postprocessing = current_processed_image[roi_y1:roi_y2, roi_x1:roi_x2].copy()
+
+        if face_roi_for_postprocessing.size == 0:
+            logging.warning("ROI for post-processing is empty. Skipping post-processing steps.")
+        else:
+            if apply_enhancement:
+                if restoration_model_loaded_successfully:
+                    progress(0.7, desc="Applying face restoration...")
+                    face_roi_for_postprocessing = enhance_cropped_face(face_roi_for_postprocessing)
                 else:
-                    logging.warning("Skipping color correction, empty face region(s) for matching.")
+                    gr.Info("Face restoration model N/A, enhancement skipped.")
+            
+            if apply_color_correction:
+                progress(0.8, desc="Applying color correction...")
+                target_face_region_for_color = target_np[roi_y1:roi_y2, roi_x1:roi_x2]
+                if target_face_region_for_color.size > 0:
+                    face_roi_for_postprocessing = histogram_match_color(face_roi_for_postprocessing, target_face_region_for_color.copy())
+
+            if apply_naturalness:
+                progress(0.85, desc="Applying naturalness filters...")
+                face_roi_for_postprocessing = apply_naturalness_filters(face_roi_for_postprocessing)
+            
+            # Paste the processed ROI back
+            if face_roi_for_postprocessing.shape[0] == (roi_y2 - roi_y1) and face_roi_for_postprocessing.shape[1] == (roi_x2 - roi_x1):
+                 current_processed_image[roi_y1:roi_y2, roi_x1:roi_x2] = face_roi_for_postprocessing
             else:
-                logging.warning("Skipping color correction, invalid bounding box for region extraction.")
+                logging.warning("Processed ROI size mismatch after post-processing. Attempting resize for paste.")
+                try:
+                    resized_processed_roi = cv2.resize(face_roi_for_postprocessing, (roi_x2-roi_x1, roi_y2-roi_y1), interpolation=cv2.INTER_LANCZOS4)
+                    current_processed_image[roi_y1:roi_y2, roi_x1:roi_x2] = resized_processed_roi
+                except Exception as e_resize:
+                    logging.error(f"Failed to resize and paste processed ROI: {e_resize}")
 
 
+        swapped_bgr_img = current_processed_image # Final image after all steps
+
     except Exception as e:
-        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
+        logging.error(f"Error during main processing pipeline: {e}", exc_info=True)
+        # Fallback to display the original target image or last good state if error occurs mid-process
+        img_to_show_on_error = swapped_bgr_img # Use current state which might be partially processed
+        swapped_pil_img_on_error = convert_cv2_to_pil(img_to_show_on_error)
+        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: {str(e)}")
-        # 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
+    if swapped_pil_img is None:
         gr.Error("Failed to convert final image to display format.")
-        swapped_pil_img = Image.new('RGB', (target_w, target_h), color='lightgrey') # Placeholder
+        swapped_pil_img = Image.new('RGB', (target_w, target_h), color='lightgrey')
 
     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, mode="wb") as tmp_file:
-            swapped_pil_img.save(tmp_file, format="JPEG") # Specify format for PIL save
+            swapped_pil_img.save(tmp_file, format="JPEG")
             temp_file_path = tmp_file.name
         logging.info(f"Swapped image saved to temporary file: {temp_file_path}")
     except Exception as e:
@@ -329,181 +433,131 @@ def process_face_swap(source_pil_img: Image.Image, target_pil_img: Image.Image,
     progress(1.0, desc="Processing complete!")
     return swapped_pil_img, temp_file_path
 
-def preview_target_faces(target_pil_img: Image.Image):
+# --- Gradio Preview Function (preview_target_faces - same as before) ---
+def preview_target_faces(target_pil_img: Image.Image): # Copied for completeness
     if target_pil_img is None:
         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)
-
     target_np = convert_pil_to_cv2(target_pil_img)
-    if target_np is None: # If conversion failed
+    if target_np is None:
         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
+    if not faces:
+        preview_pil_img = convert_cv2_to_pil(target_np)
+        if preview_pil_img is None:
             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
+    if preview_pil_img is None:
         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))
     return preview_pil_img, slider_update
 
-
-# --- Gradio UI Definition ---
-with gr.Blocks(title="Ultimate Face Swap AI 💎 v3", theme=gr.themes.Soft()) as demo: # Changed theme for variety
+# --- Gradio UI Definition (Largely same, ensure inputs to process_face_swap are correct) ---
+with gr.Blocks(title="ReSwapper AI Face Swap 💎", theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
         <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>
+            <h1>💎 ReSwapper AI Face Swap 🌠</h1>
+            <p>Using ReSwapper_256.onnx. Upload a source face, a target image, and let the AI work!</p>
         </div>
         """
     )
 
-    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.")
+    if not core_models_loaded_successfully: # This now checks face_analyzer and reswapper_session
+        gr.Error("CRITICAL ERROR: Core models (Face Analyzer or ReSwapper) failed to load. Application will not function correctly. Please check console logs and restart.")
 
     with gr.Row():
         with gr.Column(scale=1):
-            source_image_input = gr.Image(label="👤 Source Face Image (Clear, single face)", type="pil", sources=["upload", "clipboard"], height=350)
+            source_image_input = gr.Image(label="👤 Source Face Image", 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=350)
 
     with gr.Row(equal_height=True):
         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
-        )
+        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=350)
-    
-    gr.HTML("<hr style='margin-top: 15px; margin-bottom: 15px;'>") # Styled HR
+    gr.HTML("<hr style='margin-top: 15px; margin-bottom: 15px;'>")
 
     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>")
+            enhance_checkbox_label = "✨ Apply 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)
+            if not restoration_model_loaded_successfully: gr.Markdown("<p style='color: orange; font-size:0.8em;'>⚠️ Restoration model N/A.</p>")
+        
+        with gr.Column(scale=1):
+            color_correction_checkbox = gr.Checkbox(label="🎨 Apply Color Correction", value=True)
         
         with gr.Column(scale=1):
-            color_correction_checkbox = gr.Checkbox(label="🎨 Apply Color Correction (Histogram Matching)", value=True) # Default to True
+            naturalness_checkbox = gr.Checkbox(label="🌿 Apply Subtle Naturalness Filters", value=False)
 
-    
     with gr.Row():
-        swap_button = gr.Button("🚀 GENERATE SWAP!", variant="primary", scale=3, interactive=core_models_loaded_successfully) # Disable if core models failed
+        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():
         swapped_image_output = gr.Image(label="✨ Swapped Result", interactive=False, height=450)
         download_output_file = gr.File(label="⬇️ Download Swapped Image")
 
-    # --- Event Handlers ---
+    # 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')
+        if target_img_pil is None:
+            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)
-        # 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
-    )
-    
-    preview_button.click(
-        fn=preview_target_faces,
-        inputs=[target_image_input],
-        outputs=[target_faces_preview_output, face_index_slider]
-    )
+        return target_faces_preview_output.value, face_index_slider.value
+
+    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)
+    preview_button.click(fn=preview_target_faces, inputs=[target_image_input], outputs=[target_faces_preview_output, face_index_slider])
 
     swap_button.click(
         fn=process_face_swap,
-        inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox, color_correction_checkbox],
+        inputs=[
+            source_image_input, target_image_input, face_index_slider,
+            enhance_checkbox, color_correction_checkbox, naturalness_checkbox
+            # Note: Removed apply_ai_upscaling from inputs list for this example as we focused on ReSwapper integration
+        ],
         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')
-        # 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, # 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
-    )
+        return (None, None, blank_preview, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False), None, None)
+
+    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)
 
     gr.Examples(
         examples=[
-            ["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
+            ["examples/source_face.jpg", "examples/target_group.jpg", 0, True, True, True],
         ],
-        inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox, color_correction_checkbox],
+        inputs=[source_image_input, target_image_input, face_index_slider, enhance_checkbox, color_correction_checkbox, naturalness_checkbox],
         outputs=[swapped_image_output, download_output_file],
-        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
+        fn=process_face_swap, cache_examples=False,
         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) # Create if you intend to use example images
+    os.makedirs("examples", exist_ok=True)
 
-    # Console messages about model loading status
     print("\n" + "="*60)
-    print("🚀 ULTIMATE FACE SWAP AI - STARTUP STATUS 🚀")
+    print("🚀 ReSwapper AI FACE SWAP - STARTUP STATUS 🚀")
     print("="*60)
     if not core_models_loaded_successfully:
-        print("🔴 CRITICAL ERROR: Core models (Face Analyzer or Swapper) FAILED to load.")
+        print(f"🔴 CRITICAL ERROR: Core models 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(f"   - ReSwapper Model: '{SWAPPER_MODEL_PATH}' (Status: {'Loaded' if reswapper_session else 'Failed'})") # MODIFIED
         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("🟢 Core models (Face Analyzer & Swapper) loaded successfully.")
+        print("🟢 Core models (Face Analyzer & ReSwapper) 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.")
+        print(f"🟡 INFO: Face Restoration model ('{RESTORATION_MODEL_PATH}') not loaded. Enhancement feature disabled.")
     else:
         print("🟢 Face Restoration model loaded successfully.")
     print("="*60 + "\n")