Spaces:
Runtime error
Runtime error
File size: 3,520 Bytes
66c1064 ae1d47c 66c1064 713ca36 66c1064 d7c2a8c ae1d47c 66c1064 713ca36 ae1d47c 66c1064 ae1d47c 66c1064 ae1d47c 713ca36 ae1d47c 713ca36 ae1d47c 66c1064 713ca36 66c1064 ae1d47c 66c1064 ae1d47c 66c1064 ae1d47c 66c1064 ae1d47c 66c1064 ae1d47c 66c1064 |
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()
|