Update app.py

app.py

@@ -45,16 +45,31 @@ def upscale_bg_tiled(frame, tile_size=128, overlap=16):
             upscaled_img[py1:py2, px1:px2] = tile_out[oy : oy+(py2-py1), ox : ox+(px2-px1)]
     return upscaled_img
 
-def restore_face_core(img):
-    """CodeFormer Face Restoration"""
-    # Pre-process (CodeFormer expects 512x512)
+def restore_face_core(img, fidelity=0.5):
+    """CodeFormer Face Restoration with Dynamic Input Support"""
+    # 1. Pre-process (CodeFormer expects 512x512)
     img_512 = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
     img_in = np.transpose(cv2.cvtColor(img_512, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0, (2, 0, 1))[np.newaxis, :]
     
-    # AI
-    out = face_session.run(None, {face_session.get_inputs()[0].name: img_in})[0]
+    # 2. Map Inputs Dynamically
+    input_names = [i.name for i in face_session.get_inputs()]
+    input_feed = {}
     
-    # Post-process
+    # Map the main image input (usually 'input' or 'x')
+    input_feed[input_names[0]] = img_in
+    
+    # If the model requires the 'weight' or 'fidelity' input, provide it
+    if 'weight' in input_names:
+        # CodeFormer weights are usually 1D arrays
+        input_feed['weight'] = np.array([fidelity], dtype=np.float32)
+    elif len(input_names) > 1:
+        # Fallback for other potential names for the second input
+        input_feed[input_names[1]] = np.array([fidelity], dtype=np.float32)
+    
+    # 3. AI Inference
+    out = face_session.run(None, input_feed)[0]
+    
+    # 4. Post-process
     res = cv2.cvtColor((np.clip(np.squeeze(out), 0, 1).transpose(1, 2, 0) * 255.0).astype(np.uint8), cv2.COLOR_RGB2BGR)
     return cv2.resize(res, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_LANCZOS4)
 
@@ -67,13 +82,12 @@ def hybrid_enhancer(img_data, mode, face_strength, progress=gr.Progress()):
         progress(0, desc="Stage 1: Upscaling Background...")
         bg_upscaled = upscale_bg_tiled(img)
         progress(0.5, desc="Stage 2: Restoring Face Details...")
-        # In a CPU-only "Noob" version, we treat the whole crop as a face for simplicity
-        face_restored = restore_face_core(bg_upscaled)
-        # Blend based on strength
+        # Use the strength slider as the internal model weight too
+        face_restored = restore_face_core(bg_upscaled, fidelity=face_strength)
         return cv2.addWeighted(face_restored, face_strength, bg_upscaled, 1 - face_strength, 0)
     else:
-        # Just Face Restoration (No background upscale)
-        return restore_face_core(img)
+        # Just Face Restoration
+        return restore_face_core(img, fidelity=face_strength)
 
 # --- UI ---
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="indigo")) as demo: