vadim71's picture
Create app.py
9923ac4 verified
import os
import sys
import cv2
import numpy as np
from PIL import Image
import gradio as gr
import gradio_client.utils as gc_utils
import insightface
from insightface.app import FaceAnalysis
# =========================================================
# Patch for gradio_client schema bugs (schema can be bool)
# =========================================================
_orig_get_type = gc_utils.get_type
_orig_json_schema_to_python_type = gc_utils.json_schema_to_python_type
def _safe_get_type(schema):
if not isinstance(schema, dict):
return "Any"
return _orig_get_type(schema)
def _safe_json_schema_to_python_type(schema):
if not isinstance(schema, dict):
return "Any"
try:
return _orig_json_schema_to_python_type(schema)
except Exception:
return "Any"
gc_utils.get_type = _safe_get_type
gc_utils.json_schema_to_python_type = _safe_json_schema_to_python_type
# =========================================================
# Paths (OFFLINE)
# =========================================================
APP_DIR = os.path.dirname(os.path.abspath(__file__))
MODELS_DIR = os.path.join(APP_DIR, "models")
INSWAPPER_PATH = os.path.join(MODELS_DIR, "inswapper_128.onnx")
GFPGAN_PATH = os.path.join(MODELS_DIR, "GFPGANv1.3.pth")
def require_file(path: str, hint: str = ""):
if os.path.exists(path):
print(f"[OK] Found: {path}")
return
msg = f"Required file not found: {path}\n"
if hint:
msg += hint + "\n"
raise RuntimeError(msg)
require_file(
INSWAPPER_PATH,
hint="Put inswapper_128.onnx into: models/inswapper_128.onnx",
)
require_file(
GFPGAN_PATH,
hint="Put GFPGANv1.3.pth into: models/GFPGANv1.3.pth",
)
from gfpgan import GFPGANer # after files check
# =========================================================
# Localization texts
# =========================================================
TEXTS = {
"ru": {
"lang_radio_label": "Язык / Language",
"title_md": (
"# FaceSwap Pro (Docker)\n\n"
"Слева — фото **донора** (может быть несколько лиц), справа — фото, где нужно заменить лица. \n"
"Внизу — результат, круговой индикатор прогресса, оценка времени и блок «До / После».\n"
),
"step1_title": "### 1. Фото донора",
"step1_input_label": "Загрузите фото донора (может быть несколько лиц)",
"step1_donor_choice_label": "Выберите лицо‑донора",
"step2_title": "### 2. Фото, которое меняем",
"step2_input_label": "Загрузите изображение для замены лиц",
"step2_target_choices_label": "Выберите лица для замены (если ничего не выбрано — меняем все)",
"step3_title": "### 3. Настройки и экспорт",
"use_enh_label": "Улучшить качество результата (GFPGAN)",
"eta_initial": "Оценка времени появится после обнаружения лиц.",
"fmt_label": "Формат файла для скачивания",
"run_btn": "3. Запустить замену",
"download_btn": "Скачать результат",
"before_after_label": "До / После",
"eta_fmt_sec": "Оценочное время обработки: ~{sec} сек.",
"eta_fmt_min": "Оценочное время обработки: ~{min} мин {sec} сек.",
"msg_need_images": "Загрузите оба изображения, чтобы начать обработку.",
"msg_no_donor_faces": "На фото донора не найдено ни одного лица.",
"msg_no_target_faces": "На целевом фото не найдено ни одного лица.",
"msg_prep": "Подготовка к обработке...",
"msg_swap_step": "Замена лица {i} из {n}",
"msg_enh": "Улучшение качества (GFPGAN)...",
"msg_done": "Готово. Обработано лиц: {n}.",
"progress_done": "Готово!",
"donor_option": "Донор {i}",
"target_option": "Лицо {i}",
},
"en": {
"lang_radio_label": "Language / Язык",
"title_md": (
"# FaceSwap Pro (Docker)\n\n"
"Left — **donor photo** (can contain multiple faces), right — photo where faces will be replaced. \n"
"Below — result, circular progress indicator, time estimate and **Before / After** gallery.\n"
),
"step1_title": "### 1. Donor photo",
"step1_input_label": "Upload donor photo (can contain multiple faces)",
"step1_donor_choice_label": "Choose donor face",
"step2_title": "### 2. Target photo",
"step2_input_label": "Upload image where faces will be replaced",
"step2_target_choices_label": "Choose faces to replace (if none selected — replace all)",
"step3_title": "### 3. Settings & Export",
"use_enh_label": "Enhance result quality (GFPGAN)",
"eta_initial": "Time estimate will appear after detecting faces.",
"fmt_label": "Download file format",
"run_btn": "3. Start swapping",
"download_btn": "Download result",
"before_after_label": "Before / After",
"eta_fmt_sec": "Estimated processing time: ~{sec} sec.",
"eta_fmt_min": "Estimated processing time: ~{min} min {sec} sec.",
"msg_need_images": "Please upload both images to start.",
"msg_no_donor_faces": "No faces detected on donor image.",
"msg_no_target_faces": "No faces detected on target image.",
"msg_prep": "Preparing for processing...",
"msg_swap_step": "Replacing face {i} of {n}",
"msg_enh": "Enhancing quality (GFPGAN)...",
"msg_done": "Done. Faces processed: {n}.",
"progress_done": "Done!",
"donor_option": "Donor {i}",
"target_option": "Face {i}",
},
}
def T(lang: str, key: str, **kwargs) -> str:
s = TEXTS[lang][key]
return s.format(**kwargs)
# =========================================================
# Model init
# =========================================================
try:
import torch
has_gpu = torch.cuda.is_available()
except Exception:
has_gpu = False
ctx_id = 0 if has_gpu else -1
print("[INFO] GPU available" if has_gpu else "[INFO] Running on CPU")
# buffalo_l must already exist in /root/.insightface/models/buffalo_l (Dockerfile copies it)
app_face = FaceAnalysis(name="buffalo_l")
app_face.prepare(ctx_id=ctx_id, det_size=(640, 640))
# load inswapper from local path; do not download
swapper = insightface.model_zoo.get_model(INSWAPPER_PATH, download=False)
# GFPGAN from local weights
face_enhancer = GFPGANer(
model_path=GFPGAN_PATH,
upscale=1,
arch="clean",
channel_multiplier=2,
bg_upsampler=None,
)
MAX_PREVIEWS = 8
# =========================================================
# ETA
# =========================================================
def estimate_time(lang: str, num_faces: int, use_enhancer: bool) -> str:
if num_faces <= 0:
return TEXTS[lang]["eta_initial"]
if has_gpu:
base_per_face = 4
extra = 6 if use_enhancer else 0
else:
base_per_face = 10
extra = 25 if use_enhancer else 0
total = num_faces * (base_per_face + extra)
if total < 60:
return T(lang, "eta_fmt_sec", sec=int(total))
m = int(total // 60)
s = int(total % 60)
return T(lang, "eta_fmt_min", min=m, sec=s)
# =========================================================
# Face detection
# =========================================================
def detect_faces_generic(img):
if img is None:
return [], []
bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
faces = app_face.get(bgr)
faces = sorted(faces, key=lambda f: f.bbox[0])
previews = []
for f in faces:
x1, y1, x2, y2 = map(int, f.bbox)
h, w = bgr.shape[:2]
x1 = max(0, x1)
y1 = max(0, y1)
x2 = min(w, x2)
y2 = min(h, y2)
crop = bgr[y1:y2, x1:x2]
previews.append(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
return previews, faces
def update_donor_faces(img, lang: str):
previews, faces = detect_faces_generic(img)
updates = []
for i in range(MAX_PREVIEWS):
if i < len(previews):
updates.append(gr.update(value=previews[i], visible=True))
else:
updates.append(gr.update(value=None, visible=False))
labels = [T(lang, "donor_option", i=i + 1) for i in range(len(previews))]
default = labels[0] if labels else None
radio_update = gr.update(choices=labels, value=default)
return updates + [radio_update, faces]
def update_target_faces(img, use_enhancer: bool, lang: str):
previews, faces = detect_faces_generic(img)
updates = []
for i in range(MAX_PREVIEWS):
if i < len(previews):
updates.append(gr.update(value=previews[i], visible=True))
else:
updates.append(gr.update(value=None, visible=False))
labels = [T(lang, "target_option", i=i + 1) for i in range(len(previews))]
default = labels if labels else []
checkbox_update = gr.update(choices=labels, value=default)
eta_text = estimate_time(lang, len(faces), use_enhancer)
return updates + [checkbox_update, faces, eta_text]
def update_eta_only(faces, use_enhancer: bool, lang: str):
num = len(faces) if faces else 0
return estimate_time(lang, num, use_enhancer)
# =========================================================
# Progress HTML
# =========================================================
def make_progress_html(percent: int, text: str) -> str:
percent = max(0, min(100, int(percent)))
return f"""
<div class="progress-container">
<div class="progress-ring" style="--p:{percent};">
<div class="progress-ring-inner">
<span class="progress-ring-percent">{percent}%</span>
</div>
</div>
<div class="progress-ring-text">{text}</div>
</div>
"""
# =========================================================
# Main swap generator
# =========================================================
def swap_from_ui(
donor_img,
target_img,
donor_choice,
target_choices,
donor_faces,
target_faces,
use_enhancer: bool,
lang: str,
):
if donor_img is None or target_img is None:
msg = T(lang, "msg_need_images")
html = make_progress_html(0, msg)
yield None, msg, html, []
return
if not donor_faces:
msg = T(lang, "msg_no_donor_faces")
html = make_progress_html(0, msg)
yield target_img, msg, html, [target_img]
return
if not target_faces:
msg = T(lang, "msg_no_target_faces")
html = make_progress_html(0, msg)
yield target_img, msg, html, [target_img]
return
# donor face index
if donor_choice:
try:
donor_idx = int(donor_choice.split()[-1]) - 1
except Exception:
donor_idx = 0
else:
donor_idx = 0
donor_idx = max(0, min(donor_idx, len(donor_faces) - 1))
donor_face = donor_faces[donor_idx]
# target indices
if target_choices:
target_indices = []
for lbl in target_choices:
try:
idx = int(lbl.split()[-1]) - 1
if 0 <= idx < len(target_faces):
target_indices.append(idx)
except Exception:
continue
if not target_indices:
target_indices = list(range(len(target_faces)))
else:
target_indices = list(range(len(target_faces)))
tar_bgr = cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR)
result = tar_bgr.copy()
total_steps = len(target_indices) + (1 if use_enhancer else 0)
total_steps = max(total_steps, 1)
msg = T(lang, "msg_prep")
yield cv2.cvtColor(result, cv2.COLOR_BGR2RGB), msg, make_progress_html(0, msg), []
# swap loop
for i, idx in enumerate(target_indices):
step = i + 1
msg = T(lang, "msg_swap_step", i=step, n=len(target_indices))
pct = int(step / total_steps * 100)
yield cv2.cvtColor(result, cv2.COLOR_BGR2RGB), msg, make_progress_html(pct, msg), []
face_obj = target_faces[idx]
result = swapper.get(result, face_obj, donor_face, paste_back=True)
# GFPGAN
if use_enhancer:
msg = T(lang, "msg_enh")
pct = int((total_steps - 1) / total_steps * 100)
yield cv2.cvtColor(result, cv2.COLOR_BGR2RGB), msg, make_progress_html(pct, msg), []
try:
_, _, enhanced = face_enhancer.enhance(
result,
has_aligned=False,
only_center_face=False,
paste_back=True,
)
result = enhanced
except Exception as e:
print("[WARN] GFPGAN failed:", e)
final_img = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
msg = T(lang, "msg_done", n=len(target_indices))
html = make_progress_html(100, T(lang, "progress_done"))
before_after = [target_img, final_img]
yield final_img, msg, html, before_after
# =========================================================
# Save
# =========================================================
def save_result(image, fmt: str):
if image is None:
return None
img = Image.fromarray(image)
fmt = fmt.lower()
path = f"/tmp/faceswap.{fmt}"
img.save(path, format=fmt.upper())
return path
# =========================================================
# CSS + Theme
# =========================================================
custom_css = """
.gradio-container { max-width: 1240px !important; margin: 0 auto !important; }
.step-card {
background: radial-gradient(circle at top left, #0b1220, #020617 55%);
border-radius: 18px;
padding: 18px 20px;
border: 1px solid #111827;
box-shadow: 0 22px 60px rgba(15,23,42,0.9);
}
.face-thumb img {
width: 100% !important;
height: 100% !important;
object-fit: cover;
border-radius: 8px;
}
.progress-container { display:flex; flex-direction:column; align-items:center; justify-content:center; margin-top:14px; }
.progress-ring {
width:120px; height:120px; border-radius:50%;
background:conic-gradient(#8b5cf6 calc(var(--p) * 1%), #1f2937 0);
display:flex; align-items:center; justify-content:center;
box-shadow: 0 0 25px rgba(139,92,246,0.6);
transition: background 0.6s ease-out, box-shadow 0.6s ease-out;
}
.progress-ring-inner {
width:85px; height:85px; border-radius:50%;
background:#020617; display:flex; align-items:center; justify-content:center;
}
.progress-ring-percent { font-size:22px; font-weight:700; color:#e5e7eb; }
.progress-ring-text { margin-top:8px; font-size:14px; color:#e5e7eb; text-align:center; max-width:260px; }
button#run-btn {
background-image: linear-gradient(135deg, #a855f7, #ec4899);
border-radius: 999px; border: none; font-weight: 600;
letter-spacing: 0.03em;
box-shadow: 0 12px 30px rgba(236,72,153,0.45);
}
button#run-btn:hover { filter: brightness(1.08); box-shadow: 0 16px 44px rgba(236,72,153,0.6); }
button#download-btn {
background-image: linear-gradient(135deg, #22c55e, #16a34a);
border-radius: 999px; border: none; color: #0f172a; font-weight: 600;
letter-spacing: 0.03em;
box-shadow: 0 10px 24px rgba(34,197,94,0.45);
}
button#download-btn:hover { filter: brightness(1.05); box-shadow: 0 14px 32px rgba(34,197,94,0.6); }
"""
premium_theme = gr.themes.Default(primary_hue="violet", neutral_hue="slate").set(
body_text_size="16px",
body_text_weight="500",
button_large_text_size="16px",
button_large_padding="12px 22px",
)
# =========================================================
# UI
# =========================================================
with gr.Blocks(
title="FaceSwap Pro (Docker)",
css=custom_css,
theme=premium_theme,
) as demo:
lang_state = gr.State("ru")
lang_radio = gr.Radio(
choices=["RU", "EN"],
value="RU",
label=TEXTS["ru"]["lang_radio_label"],
)
title_md = gr.Markdown(TEXTS["ru"]["title_md"])
donor_faces_state = gr.State([])
target_faces_state = gr.State([])
with gr.Row(elem_classes=["step-card"]):
with gr.Column():
step1_title_md = gr.Markdown(TEXTS["ru"]["step1_title"])
donor_img = gr.Image(
label=TEXTS["ru"]["step1_input_label"],
type="numpy",
height=420,
)
gr.Markdown("**Найденные лица (донор):**")
with gr.Row():
donor_previews = [
gr.Image(
value=None,
visible=False,
show_label=False,
interactive=False,
width=32,
height=32,
type="numpy",
elem_classes=["face-thumb"],
)
for _ in range(MAX_PREVIEWS)
]
donor_choice = gr.Radio(
label=TEXTS["ru"]["step1_donor_choice_label"],
choices=[],
)
with gr.Column():
step2_title_md = gr.Markdown(TEXTS["ru"]["step2_title"])
target_img = gr.Image(
label=TEXTS["ru"]["step2_input_label"],
type="numpy",
height=420,
)
gr.Markdown("**Найденные лица (цель):**")
with gr.Row():
target_previews = [
gr.Image(
value=None,
visible=False,
show_label=False,
interactive=False,
width=32,
height=32,
type="numpy",
elem_classes=["face-thumb"],
)
for _ in range(MAX_PREVIEWS)
]
target_choices = gr.CheckboxGroup(
label=TEXTS["ru"]["step2_target_choices_label"],
choices=[],
)
with gr.Row(elem_classes=["step-card"]):
with gr.Column(scale=1):
step3_title_md = gr.Markdown(TEXTS["ru"]["step3_title"])
use_enh = gr.Checkbox(label=TEXTS["ru"]["use_enh_label"], value=True)
eta_text = gr.Markdown(TEXTS["ru"]["eta_initial"])
fmt = gr.Dropdown(
label=TEXTS["ru"]["fmt_label"],
choices=["png", "jpeg", "webp"],
value="png",
)
run_btn = gr.Button(TEXTS["ru"]["run_btn"], variant="primary", elem_id="run-btn")
progress_html = gr.HTML("")
download_btn = gr.DownloadButton(TEXTS["ru"]["download_btn"], elem_id="download-btn")
with gr.Column(scale=2):
result_img = gr.Image(label="", show_label=False, interactive=False, type="numpy", height=520)
status_md = gr.Markdown("")
before_after = gr.Gallery(label=TEXTS["ru"]["before_after_label"], columns=2, height=350)
def switch_language(choice):
lang = "ru" if choice == "RU" else "en"
t = TEXTS[lang]
return (
lang,
gr.update(label=t["lang_radio_label"]),
gr.update(value=t["title_md"]),
gr.update(value=t["step1_title"]),
gr.update(label=t["step1_input_label"]),
gr.update(label=t["step1_donor_choice_label"]),
gr.update(value=t["step2_title"]),
gr.update(label=t["step2_input_label"]),
gr.update(label=t["step2_target_choices_label"]),
gr.update(value=t["step3_title"]),
gr.update(label=t["use_enh_label"]),
gr.update(value=t["eta_initial"]),
gr.update(label=t["fmt_label"]),
gr.update(value=t["run_btn"]),
gr.update(value=t["download_btn"]),
gr.update(label=t["before_after_label"]),
)
lang_radio.change(
fn=switch_language,
inputs=lang_radio,
outputs=[
lang_state,
lang_radio,
title_md,
step1_title_md,
donor_img,
donor_choice,
step2_title_md,
target_img,
target_choices,
step3_title_md,
use_enh,
eta_text,
fmt,
run_btn,
download_btn,
before_after,
],
)
donor_img.change(
fn=update_donor_faces,
inputs=[donor_img, lang_state],
outputs=donor_previews + [donor_choice, donor_faces_state],
)
target_img.change(
fn=update_target_faces,
inputs=[target_img, use_enh, lang_state],
outputs=target_previews + [target_choices, target_faces_state, eta_text],
)
use_enh.change(
fn=update_eta_only,
inputs=[target_faces_state, use_enh, lang_state],
outputs=eta_text,
)
run_btn.click(
fn=swap_from_ui,
inputs=[
donor_img,
target_img,
donor_choice,
target_choices,
donor_faces_state,
target_faces_state,
use_enh,
lang_state,
],
outputs=[result_img, status_md, progress_html, before_after],
)
download_btn.click(
fn=save_result,
inputs=[result_img, fmt],
outputs=download_btn,
)
demo.launch(
server_name="0.0.0.0",
server_port=int(os.getenv("PORT", "7860")),
)