Update app.py

app.py

@@ -3,7 +3,6 @@ import cv2
 import numpy as np
 import gradio as gr
 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,33 +11,28 @@ import onnxruntime
 # --- Configuration & Setup ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# MODIFIED: Path to the new swapper model
-SWAPPER_MODEL_PATH = "models/reswapper_256.onnx" # <--- USING RESWAPPER_256
+SWAPPER_MODEL_PATH = "models/reswapper_256.onnx" # Using ReSwapper
 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 and embedding model
+FACE_ANALYZER_NAME = 'buffalo_l'
 DETECTION_SIZE = (640, 640)
-EXECUTION_PROVIDERS = ['CPUExecutionProvider'] # Or ['CUDAExecutionProvider', 'CPUExecutionProvider']
+EXECUTION_PROVIDERS = ['CPUExecutionProvider']
 
-# --- Global Variables (Lazy Loaded by initialize_models) ---
+# --- Global Variables ---
 face_analyzer = None
-reswapper_session = None # MODIFIED: Will hold the onnxruntime session for reswapper
+reswapper_session = None # For reswapper_256.onnx
 face_restorer = None
-# face_upscaler = None # If you add AI upscaling
 
 # --- Initialization Functions ---
 def initialize_models():
-    global face_analyzer, reswapper_session, face_restorer # MODIFIED: reswapper_session
+    global face_analyzer, reswapper_session, face_restorer
     try:
-        # 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)
             logging.info("FaceAnalysis model initialized successfully.")
 
-        # MODIFIED: Initialize ReSwapper model directly with onnxruntime
         if reswapper_session is None:
             if not os.path.exists(SWAPPER_MODEL_PATH):
                 logging.error(f"ReSwapper model FILE NOT FOUND at {SWAPPER_MODEL_PATH}. Swapping will fail.")
@@ -49,14 +43,12 @@ def initialize_models():
                     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)
+                    reswapper_session = None
+
         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 will be disabled.")
+                logging.error(f"Face restoration model FILE NOT FOUND at: {RESTORATION_MODEL_PATH}. Enhancement feature 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)
@@ -64,19 +56,14 @@ def initialize_models():
                 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
-        
-        # (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)
 
-# --- Call Initialization Early ---
 initialize_models()
-core_models_loaded_successfully = face_analyzer is not None and reswapper_session is not None # MODIFIED
+core_models_loaded_successfully = face_analyzer is not None and reswapper_session is not None
 restoration_model_loaded_successfully = face_restorer is not None
-# upscaler_model_loaded_successfully = face_upscaler is not None # If adding upscaler
 
-# --- Image Utility Functions (convert_pil_to_cv2, convert_cv2_to_pil - same as before) ---
+# --- Image Utility Functions ---
 def convert_pil_to_cv2(pil_image: Image.Image) -> np.ndarray | None:
     if pil_image is None: return None
     try:
@@ -93,123 +80,109 @@ 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 (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:
+# --- Core AI & Image Processing Functions ---
 def get_faces_from_image(img_np: np.ndarray) -> list:
-    if face_analyzer is None: # Check against the global face_analyzer
+    if face_analyzer is None:
         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) # This provides Face objects with embeddings
+        return face_analyzer.get(img_np)
     except Exception as e:
         logging.error(f"Error during face detection: {e}", exc_info=True)
         return []
 
-# ... (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
+def draw_detected_faces(img_np: np.ndarray, faces: list) -> np.ndarray:
     img_with_boxes = img_np.copy()
     for i, face in enumerate(faces):
         box = face.bbox.astype(int)
         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))
-        cv2.putText(img_with_boxes, f"Face {i}", label_position,
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (36, 255, 12), 2)
+        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: # Copied for completeness
+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.")
         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([])
+    if face_crop_bgr is None or face_crop_bgr.size == 0: 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]
+    crop_h, crop_w = face_crop_bgr.shape[:2]
     try:
         img_rgb = cv2.cvtColor(face_crop_bgr, cv2.COLOR_BGR2RGB)
-        restorer_input_size = (512, 512) 
-        img_resized_for_model = cv2.resize(img_rgb, restorer_input_size, interpolation=cv2.INTER_AREA)
-        img_normalized = (img_resized_for_model / 255.0).astype(np.float32)
-        img_chw = np.transpose(img_normalized, (2, 0, 1))
+        input_size = (512, 512) # Common for GFPGAN-like models
+        img_resized = cv2.resize(img_rgb, input_size, interpolation=cv2.INTER_AREA)
+        img_norm = (img_resized / 255.0).astype(np.float32)
+        img_chw = np.transpose(img_norm, (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))
-        restored_img_uint8_model_size = np.clip(restored_img_hwc_model_size * 255.0, 0, 255).astype(np.uint8)
-        restored_crop_rgb = cv2.resize(restored_img_uint8_model_size, (crop_width, crop_height), interpolation=cv2.INTER_LANCZOS4)
-        restored_crop_bgr = cv2.cvtColor(restored_crop_rgb, cv2.COLOR_RGB2BGR)
-        logging.info("Cropped face restoration complete.")
-        return restored_crop_bgr
+        output = face_restorer.run([output_name], {input_name: input_tensor})[0]
+        out_chw = np.squeeze(output, axis=0)
+        out_hwc = np.transpose(out_chw, (1,2,0))
+        out_denorm = np.clip(out_hwc * 255.0, 0, 255).astype(np.uint8)
+        out_resized_rgb = cv2.resize(out_denorm, (crop_w, crop_h), interpolation=cv2.INTER_LANCZOS4)
+        return cv2.cvtColor(out_resized_rgb, cv2.COLOR_RGB2BGR)
     except Exception as e:
         logging.error(f"Error during face restoration for crop: {e}", exc_info=True)
         return face_crop_bgr
 
-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()
-    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()
-    target_cdf = target_hist.cumsum()
-    # 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')
-    idx = 0
+def histogram_match_channel(source_ch: np.ndarray, target_ch: np.ndarray) -> np.ndarray:
+    src_shape = source_ch.shape
+    src_flat = source_ch.flatten()
+    tgt_flat = target_ch.flatten()
+    src_hist, _ = np.histogram(src_flat, 256, [0,255])
+    tgt_hist, _ = np.histogram(tgt_flat, 256, [0,255])
+    src_cdf = src_hist.cumsum()
+    tgt_cdf = tgt_hist.cumsum()
+    if src_cdf[-1] == 0: return source_ch # Avoid division by zero for blank source
+    if tgt_cdf[-1] == 0: return source_ch # Avoid issues with blank target
+    src_cdf_norm = src_cdf * float(tgt_cdf[-1]) / src_cdf[-1] # Normalize src CDF to target range
+    lut = np.zeros(256, dtype='uint8')
+    j = 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)
+        while j < 255 and tgt_cdf[j] < src_cdf_norm[i]:
+            j += 1
+        lut[i] = j
+    return cv2.LUT(source_ch, lut).reshape(src_shape)
 
-def histogram_match_color(source_img: np.ndarray, target_img: np.ndarray) -> np.ndarray: # Copied
+def histogram_match_color(source_img: np.ndarray, target_img: np.ndarray) -> np.ndarray:
     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)
+    matched = np.zeros_like(source_img)
     try:
         for i in range(source_img.shape[2]):
-            matched_img[:,:,i] = histogram_match_channel(source_img[:,:,i], target_img[:,:,i])
-        return matched_img
+            matched[:,:,i] = histogram_match_channel(source_img[:,:,i], target_img[:,:,i])
+        return matched
     except Exception as e:
         logging.error(f"Error during histogram matching: {e}", exc_info=True)
         return source_img
 
-def apply_naturalness_filters(face_region_bgr: np.ndarray, noise_level: int = 3) -> np.ndarray: # Copied
+def apply_naturalness_filters(face_region_bgr: np.ndarray, noise_level: int = 3) -> np.ndarray:
     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()
+    proc = face_region_bgr.copy()
     try:
-        processed_region = cv2.medianBlur(processed_region, 3)
+        proc = cv2.medianBlur(proc, 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.")
+            noise = np.random.normal(0, noise_level, proc.shape).astype(np.int16)
+            proc_int16 = proc.astype(np.int16) + noise
+            proc = np.clip(proc_int16, 0, 255).astype(np.uint8)
+        logging.info("Applied naturalness filters.")
     except Exception as e:
         logging.error(f"Error applying naturalness filters: {e}", exc_info=True)
         return face_region_bgr
-    return processed_region
+    return proc
 
 # --- Main Processing Function ---
 def process_face_swap(source_pil_img: Image.Image, target_pil_img: Image.Image,
-                        target_face_index: int, # apply_ai_upscaling: bool, (Removed for this example to focus on ReSwapper)
+                        target_face_index: int,
                         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: # Checks face_analyzer and reswapper_session
+    if not core_models_loaded_successfully:
         gr.Error("CRITICAL: Core models (Face Analyzer or ReSwapper) not loaded. Cannot proceed.")
         return Image.new('RGB', (100,100), color='lightgrey'), None
 
@@ -217,350 +190,265 @@ def process_face_swap(source_pil_img: Image.Image, target_pil_img: Image.Image,
     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) # BGR
-    target_np = convert_pil_to_cv2(target_pil_img) # BGR
+    source_np = convert_pil_to_cv2(source_pil_img)
+    target_np = convert_pil_to_cv2(target_pil_img)
     if source_np is None or target_np is None:
-        raise gr.Error("Image conversion failed. Please try different images.")
+        raise gr.Error("Image conversion failed.")
     
-    target_h, target_w = target_np.shape[:2]
+    original_target_h, original_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.")
-    source_face = source_faces[0] # InsightFace Face object
+    if not source_faces: raise gr.Error("No face in source.")
+    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.")
+    if not target_faces: raise gr.Error("No faces in target.")
     if not (0 <= target_face_index < len(target_faces)):
-        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
+        raise gr.Error(f"Target face index out of range.")
+    target_face_info = target_faces[int(target_face_index)]
 
-    swapped_bgr_img = target_np.copy() # Initialize with a copy
+    swapped_bgr_img_final = target_np.copy() # Initialize for fallback
 
     try:
-        progress(0.4, desc="Preparing inputs for ReSwapper...")
+        progress(0.35, desc="Preparing inputs for ReSwapper...")
         
-        ########## 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.
+        # 1. Source Face Embedding
+        source_embedding = source_face.embedding.astype(np.float32)
         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_rgb = cv2.cvtColor(target_np, cv2.COLOR_BGR2RGB)
+        
+        # ***** FIX APPLIED HERE: Resize target image for ReSwapper *****
+        # The "256" in reswapper_256.onnx strongly suggests this expected input size.
+        # Your error indicated: Expected 256, Got 225 (or similar).
+        reswapper_input_size = (256, 256) # Expected H, W by the model
         
-        target_img_for_onnx = cv2.cvtColor(target_np, cv2.COLOR_BGR2RGB) # Often RGB for ONNX
+        # Store original size of the target image if it's not already stored for final resize
+        # original_target_h, original_target_w are already available
+
+        target_img_resized_rgb = cv2.resize(target_img_rgb, (reswapper_input_size[1], reswapper_input_size[0]), interpolation=cv2.INTER_AREA)
+                                        # cv2.resize takes (width, height)
         
-        # 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)
+        # --- VERIFY NORMALIZATION for reswapper_256.onnx ---
+        # Common options:
+        # Option A: [0, 1]
+        target_img_normalized = (target_img_resized_rgb / 255.0).astype(np.float32)
+        # Option B: [-1, 1]
+        # target_img_normalized = (target_img_resized_rgb / 127.5 - 1.0).astype(np.float32)
+        # Choose the one your reswapper_256.onnx model expects!
         
         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.
+        # 3. Target Face Bounding Box (Highly Model Specific!)
+        # --- VERIFY IF AND HOW reswapper_256.onnx USES THIS ---
+        # If it does, it needs to be scaled to the resized_target_image (256x256)
+        # Original bbox is target_face_info.bbox relative to original_target_w, original_target_h
+        # For now, we'll pass the original bbox, but this might be incorrect if the model
+        # expects a bbox relative to the (potentially resized) target_image_tensor.
+        # Some models might not need it if they perform their own detection on the target_image_tensor.
+        target_bbox_original = target_face_info.bbox.astype(np.float32)
+        target_bbox_tensor_input = np.expand_dims(target_bbox_original, axis=0)
+
+
+        # --- VERIFY TENSOR NAMES AND INPUT ORDER for reswapper_256.onnx ---
+        # Use Netron to inspect your .onnx file.
+        # The following are GUESSES and placeholders.
+        onnx_input_list = reswapper_session.get_inputs()
+        onnx_output_list = reswapper_session.get_outputs()
+
+        input_feed = {}
+        # Example GUESS:
+        # Assumed input order: target_image, source_embedding, (optional) target_bbox
+        if len(onnx_input_list) > 0: input_feed[onnx_input_list[0].name] = target_image_tensor
+        if len(onnx_input_list) > 1: input_feed[onnx_input_list[1].name] = source_embedding_tensor
+        if len(onnx_input_list) > 2: # If model expects a third input (e.g., bbox)
+            input_feed[onnx_input_list[2].name] = target_bbox_tensor_input # This bbox might need scaling/normalization!
         
-        # 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 ##########
+        output_name = onnx_output_list[0].name # Assuming first output is the image
         
         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
+        swapped_image_tensor_output = onnx_outputs[0]
 
         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 swapped_image_tensor_output.ndim == 4:
+            swapped_image_tensor_output = np.squeeze(swapped_image_tensor_output, axis=0)
+        
+        swapped_image_hwc = np.transpose(swapped_image_tensor_output, (1, 2, 0))
+        
+        # --- VERIFY DENORMALIZATION (inverse of normalization used above) ---
         # 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)
+        swapped_image_uint8_rgb = 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)
+        # The output of reswapper will be 256x256 (if input was 256x256 and model maintains size)
+        # Resize this output back to the original target image's dimensions
+        swapped_bgr_resized_to_original = cv2.resize(
+            swapped_image_uint8_rgb, 
+            (original_target_w, original_target_h), 
+            interpolation=cv2.INTER_LANCZOS4
+        )
+        # Convert to BGR as the rest of the pipeline expects BGR
+        swapped_bgr_img_final = cv2.cvtColor(swapped_bgr_resized_to_original, cv2.COLOR_RGB2BGR)
         
         # --- 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]
+        current_processed_image = swapped_bgr_img_final.copy()
+        # ROI for post-processing is based on original bbox coordinates
+        roi_x1, roi_y1, roi_x2, roi_y2 = target_face_info.bbox.astype(int)
         
         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:
-                    gr.Info("Face restoration model N/A, enhancement skipped.")
+        if face_roi_for_postprocessing.size > 0:
+            if apply_enhancement and restoration_model_loaded_successfully:
+                progress(0.7, desc="Applying face restoration...")
+                face_roi_for_postprocessing = enhance_cropped_face(face_roi_for_postprocessing)
             
             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())
+                original_target_roi_for_color = target_np[roi_y1:roi_y2, roi_x1:roi_x2] # Original target pixels
+                if original_target_roi_for_color.size > 0:
+                    face_roi_for_postprocessing = histogram_match_color(face_roi_for_postprocessing, original_target_roi_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):
+            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("Processed ROI size mismatch after post-processing. Attempting resize for paste.")
+            else: # If size changed during post-processing (should not happen with current funcs)
+                logging.warning("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}")
-
+                    resized_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_roi
+                except Exception as e_resize_paste:
+                    logging.error(f"Failed to resize/paste processed ROI: {e_resize_paste}")
+        else:
+            logging.warning("ROI for post-processing is empty. Skipping these steps.")
 
-        swapped_bgr_img = current_processed_image # Final image after all steps
+        swapped_bgr_img_final = current_processed_image
 
     except Exception as e:
         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)
+        swapped_pil_img_on_error = convert_cv2_to_pil(swapped_bgr_img_final) # Use last good state
         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)}")
+             swapped_pil_img_on_error = Image.new('RGB', (original_target_w, original_target_h), color='lightgrey')
+        raise gr.Error(f"An error occurred: {str(e)}. Check console for details, especially ONNX input/output errors.")
 
     progress(0.9, desc="Finalizing image...")
-    swapped_pil_img = convert_cv2_to_pil(swapped_bgr_img)
+    swapped_pil_img = convert_cv2_to_pil(swapped_bgr_img_final)
     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')
+        gr.Error("Failed to convert final image.")
+        swapped_pil_img = Image.new('RGB', (original_target_w, original_target_h), color='lightgrey')
 
     temp_file_path = None
     try:
         with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False, mode="wb") as tmp_file:
             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:
-        logging.error(f"Error saving to temporary file: {e}", exc_info=True)
-        gr.Warning("Could not save the swapped image for download.")
+        logging.error(f"Error saving temp file: {e}", exc_info=True)
+        gr.Warning("Could not save swapped image for download.")
     
     progress(1.0, desc="Processing complete!")
     return swapped_pil_img, temp_file_path
 
-# --- Gradio Preview Function (preview_target_faces - same as before) ---
-def preview_target_faces(target_pil_img: Image.Image): # Copied for completeness
+# --- Gradio Preview Function ---
+def preview_target_faces(target_pil_img: Image.Image):
     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)
+        blank_img = Image.new('RGB', DETECTION_SIZE, color='lightgray')
+        return blank_img, 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:
-        blank_image_pil = Image.new('RGB', (DETECTION_SIZE[0], DETECTION_SIZE[1]), color='lightgray')
-        return blank_image_pil, gr.Slider(minimum=0, maximum=0, value=0, step=1, interactive=False)
+        return Image.new('RGB', DETECTION_SIZE, color='lightgray'), gr.Slider(interactive=False)
     faces = get_faces_from_image(target_np)
     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:
-        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 (Largely same, ensure inputs to process_face_swap are correct) ---
-with gr.Blocks(title="ReSwapper AI Face Swap 💎", theme=gr.themes.Soft()) as demo:
+        return convert_cv2_to_pil(target_np) or Image.new('RGB', DETECTION_SIZE, color='lightgray'), gr.Slider(interactive=False)
+    
+    preview_np = draw_detected_faces(target_np, faces)
+    preview_pil = convert_cv2_to_pil(preview_np) or Image.new('RGB', DETECTION_SIZE, color='lightgray')
+    return preview_pil, gr.Slider(minimum=0, maximum=max(0, len(faces)-1), value=0, step=1, interactive=len(faces)>0)
+
+# --- Gradio UI Definition ---
+with gr.Blocks(title="ReSwapper AI Face Swap 💎 v2", theme=gr.themes.Soft()) as demo: # Changed title slightly
     gr.Markdown(
         """
         <div style="text-align: center;">
-            <h1>💎 ReSwapper AI Face Swap 🌠</h1>
-            <p>Using ReSwapper_256.onnx. Upload a source face, a target image, and let the AI work!</p>
+            <h1>💎 ReSwapper AI Face Swap 🌠 v2</h1>
+            <p>Utilizing <code>reswapper_256.onnx</code>. Please ensure model inputs/outputs are correctly configured in the code if issues arise.</p>
         </div>
         """
     )
-
-    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.")
+    if not core_models_loaded_successfully:
+        gr.Error("CRITICAL ERROR: Core models (Face Analyzer or ReSwapper) failed to load. Check console.")
 
     with gr.Row():
         with gr.Column(scale=1):
-            source_image_input = gr.Image(label="👤 Source Face Image", 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)
-
+            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", 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;'>")
-
+    gr.HTML("<hr style='margin:15px 0;'>")
     with gr.Row():
         with gr.Column(scale=1):
-            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>")
-        
+            enh_label = "✨ Face Restoration" + (" (Model N/A)" if not restoration_model_loaded_successfully else "")
+            enhance_checkbox = gr.Checkbox(label=enh_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 N/A</p>")
         with gr.Column(scale=1):
-            color_correction_checkbox = gr.Checkbox(label="🎨 Apply Color Correction", value=True)
-        
+            color_correction_checkbox = gr.Checkbox(label="🎨 Color Correction", value=True)
         with gr.Column(scale=1):
-            naturalness_checkbox = gr.Checkbox(label="🌿 Apply Subtle Naturalness Filters", value=False)
-
+            naturalness_checkbox = gr.Checkbox(label="🌿 Naturalness Filters", value=False)
     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():
         swapped_image_output = gr.Image(label="✨ Swapped Result", interactive=False, height=450)
         download_output_file = gr.File(label="⬇️ Download Swapped Image")
 
     # Event Handlers
-    def on_target_image_change_or_clear(target_img_pil):
-        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)
-        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)
+    target_image_input.change(fn=lambda x: (Image.new('RGB', DETECTION_SIZE, color='lightgray') if x is None else target_faces_preview_output.value, 
+                                         gr.Slider(interactive=False) if x is None else face_index_slider.value), 
+                              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, 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')
-        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)
-
+    swap_button.click(fn=process_face_swap,
+                      inputs=[source_image_input, target_image_input, face_index_slider,
+                              enhance_checkbox, color_correction_checkbox, naturalness_checkbox],
+                      outputs=[swapped_image_output, download_output_file])
+    def clear_all():
+        return None, None, Image.new('RGB', DETECTION_SIZE, color='lightgray'), gr.Slider(interactive=False), None, None
+    clear_button.click(fn=clear_all, 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, True],
-        ],
-        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, cache_examples=False,
-        label="Example Face Swaps (Click to run)"
+        [["examples/source_face.jpg", "examples/target_group.jpg", 0, True, True, True]],
+        [source_image_input, target_image_input, face_index_slider, enhance_checkbox, color_correction_checkbox, naturalness_checkbox],
+        [swapped_image_output, download_output_file],
+        process_face_swap, cache_examples=False, label="Example"
     )
 
 # --- Main Execution Block ---
 if __name__ == "__main__":
     os.makedirs("models", exist_ok=True)
     os.makedirs("examples", exist_ok=True)
-
-    print("\n" + "="*60)
-    print("🚀 ReSwapper AI FACE SWAP - STARTUP STATUS 🚀")
-    print("="*60)
+    print("\n" + "="*60 + "\n🚀 ReSwapper AI FACE SWAP - STARTUP STATUS 🚀\n" + "="*60)
     if not core_models_loaded_successfully:
-        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"   - ReSwapper Model: '{SWAPPER_MODEL_PATH}' (Status: {'Loaded' if reswapper_session else 'Failed'})") # MODIFIED
-        print("   The application UI will load but will be NON-FUNCTIONAL.")
-    else:
-        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. Enhancement feature disabled.")
-    else:
-        print("🟢 Face Restoration model loaded successfully.")
-    print("="*60 + "\n")
-    
-    print("Launching Gradio Interface...")
+        print(f"🔴 CRITICAL: Core models FAILED. Analyzer: {'OK' if face_analyzer else 'FAIL'}, ReSwapper: {'OK' if reswapper_session else 'FAIL'}")
+    else: print("🟢 Core models (Analyzer & ReSwapper) loaded.")
+    if not restoration_model_loaded_successfully: print(f"🟡 INFO: Restoration model ('{RESTORATION_MODEL_PATH}') N/A.")
+    else: print("🟢 Restoration model loaded.")
+    print("="*60 + "\nLaunching Gradio Interface...")
     demo.launch()