Spaces:
Runtime error
Runtime error
File size: 3,520 Bytes
bb9106a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 |
import cv2
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
# -------------------------------------------------
# InSwapper model
inswapper_path = hf_hub_download(
repo_id=REPO_ID,
filename="models/inswapper_128.onnx",
repo_type="model",
local_dir=MODELS_DIR
)
# Buffalo_l ONNX files
buffalo_files = [
"1k3d68.onnx",
"2d106det.onnx",
"genderage.onnx",
"det_10g.onnx",
"w600k_r50.onnx"
]
for f in buffalo_files:
hf_hub_download(
repo_id=REPO_ID,
filename=f"models/buffalo_l/{f}",
repo_type="model",
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_DIR, providers=['CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640))
swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])
# -------------------------------------------------
# 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)
tgt_img = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)
src_faces = app.get(src_img)
tgt_faces = app.get(tgt_img)
if len(src_faces) == 0 or len(tgt_faces) == 0:
return None # No face detected
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 + 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 & Upscaled Output")
btn1 = gr.Button("Swap + Upscale")
btn1.click(swap_face_image, inputs=[src_img, tgt_img], outputs=output_img)
# -------------------------------------------------
# Launch
# -------------------------------------------------
demo.launch()
|