esrgan added

app.py

@@ -3,24 +3,28 @@ import insightface
 from insightface.app import FaceAnalysis
 import gradio as gr
 from huggingface_hub import hf_hub_download
+from basicsr.archs.rrdbnet_arch import RRDBNet
+from realesrgan import RealESRGANer
+import os
 
 # -------------------------------------------------
 # Hugging Face repo ID where models are stored
 # -------------------------------------------------
 REPO_ID = "HariLogicgo/face_swap_models"
+MODELS_DIR = "./models"
 
 # -------------------------------------------------
 # Download models from HF Hub
 # -------------------------------------------------
-# Download inswapper model
+# InSwapper model
 inswapper_path = hf_hub_download(
     repo_id=REPO_ID,
     filename="models/inswapper_128.onnx",
     repo_type="model",
-    local_dir="./models"
+    local_dir=MODELS_DIR
 )
 
-# Download buffalo_l ONNX files one by one
+# Buffalo_l ONNX files
 buffalo_files = [
     "1k3d68.onnx",
     "2d106det.onnx",
@@ -34,19 +38,51 @@ for f in buffalo_files:
         repo_id=REPO_ID,
         filename=f"models/buffalo_l/{f}",
         repo_type="model",
-        local_dir="./models"
+        local_dir=MODELS_DIR
     )
 
+# RealESRGAN model
+realesrgan_path = hf_hub_download(
+    repo_id=REPO_ID,
+    filename="models/RealESRGAN_x4plus.pth",
+    repo_type="model",
+    local_dir=MODELS_DIR
+)
+
 # -------------------------------------------------
 # Initialize Face Analysis + Swapper
 # -------------------------------------------------
-app = FaceAnalysis(name="buffalo_l", root="./models", providers=['CPUExecutionProvider'])
+app = FaceAnalysis(name="buffalo_l", root=MODELS_DIR, providers=['CPUExecutionProvider'])
 app.prepare(ctx_id=0, det_size=(640, 640))
 
 swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])
 
 # -------------------------------------------------
-# Face swap function (images only)
+# Initialize RealESRGAN Upscaler (CPU mode)
+# -------------------------------------------------
+model = RRDBNet(
+    num_in_ch=3,
+    num_out_ch=3,
+    num_feat=64,
+    num_block=23,
+    num_grow_ch=32,
+    scale=4
+)
+
+upscaler = RealESRGANer(
+    scale=4,
+    model_path=realesrgan_path,
+    dni_weight=None,
+    model=model,
+    tile=0,
+    tile_pad=10,
+    pre_pad=0,
+    half=False,    # must be False for CPU
+    gpu_id=None    # ensures CPU mode
+)
+
+# -------------------------------------------------
+# Face swap + upscale function
 # -------------------------------------------------
 def swap_face_image(source_img, target_img):
     src_img = cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR)
@@ -58,22 +94,27 @@ def swap_face_image(source_img, target_img):
     if len(src_faces) == 0 or len(tgt_faces) == 0:
         return None  # No face detected
 
-    result = swapper.get(tgt_img, tgt_faces[0], src_faces[0])
-    result = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
+    swapped = swapper.get(tgt_img, tgt_faces[0], src_faces[0])
+
+    # Upscale with RealESRGAN
+    upscaled, _ = upscaler.enhance(swapped, outscale=4)
+
+    # Convert back to RGB for Gradio
+    result = cv2.cvtColor(upscaled, cv2.COLOR_BGR2RGB)
     return result
 
 # -------------------------------------------------
 # Gradio Interface
 # -------------------------------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("Face Swap API (Image Only, CPU Deployment)")
+    gr.Markdown("### Face Swap + Upscale (CPU, HuggingFace Models)")
 
     with gr.Row():
         src_img = gr.Image(type="numpy", label="Source Image")
         tgt_img = gr.Image(type="numpy", label="Target Image")
 
-    output_img = gr.Image(type="numpy", label="Swapped Output")
-    btn1 = gr.Button("Swap Faces")
+    output_img = gr.Image(type="numpy", label="Swapped & Upscaled Output")
+    btn1 = gr.Button("Swap + Upscale")
 
     btn1.click(swap_face_image, inputs=[src_img, tgt_img], outputs=output_img)
 

requirements.txt

@@ -3,4 +3,6 @@ onnxruntime
 opencv-python
 numpy
 huggingface_hub
-gradio
+gradio
+basicsr 
+realesrgan