app.py

import os
import cv2
import numpy as np
from PIL import Image
import urllib.request
import gradio as gr
import insightface
from insightface.app import FaceAnalysis
from gfpgan import GFPGANer

# --- Загрузка моделей при первом старте ---
def download_file(url: str, filename: str):
    if not os.path.exists(filename):
        print(f"Скачиваю {filename} из {url}...")
        urllib.request.urlretrieve(url, filename)
        print(f"{filename} скачан.")

# Пути к моделям
INSWAPPER_URL = "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx"
INSWAPPER_PATH = "inswapper_128.onnx"
GFPGAN_URL = "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth"
GFPGAN_PATH = "GFPGANv1.3.pth"

download_file(INSWAPPER_URL, INSWAPPER_PATH)
download_file(GFPGAN_URL, GFPGAN_PATH)

assert os.path.exists(INSWAPPER_PATH), "inswapper_128.onnx не найден."
assert os.path.exists(GFPGAN_PATH), "GFPGANv1.3.pth не найден."

# --- Инициализация моделей ---
try:
    import torch
    has_gpu = torch.cuda.is_available()
except Exception:
    has_gpu = False

ctx_id = 0 if has_gpu else -1  # 0 = GPU, -1 = CPU
print("GPU доступен" if has_gpu else "Работаем на CPU")

providers = ['CPUExecutionProvider']  # Только CPU для бесплатного тарифа

app_face = FaceAnalysis(name="buffalo_l", providers=providers)
app_face.prepare(ctx_id=ctx_id, det_size=(640, 640))

swapper = insightface.model_zoo.get_model(INSWAPPER_PATH, providers=providers)
face_enhancer = GFPGANer(
    model_path=GFPGAN_PATH,
    upscale=1,
    arch="clean",
    channel_multiplier=2,
    bg_upsampler=None,
)

MAX_PREVIEWS = 8  # Максимум маленьких превью

# --- Тексты для локализации ---
TEXTS = {
    "ru": {
        "lang_radio_label": "Язык / Language",
        "title_md": "# FaceSwap Pro\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": "Запустить замену",
        "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\nLeft — **donor photo**, right — photo where faces will be replaced.",
        "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": "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)

# --- Основная логика обработки ---
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 estimate_time(lang: str, num_faces: int, use_enhancer: bool) -> str:
    if num_faces <= 0:
        return TEXTS[lang]["eta_initial"]
    base_per_face = 4 if has_gpu else 10
    extra_per_face = 6 if use_enhancer else 0
    total = num_faces * (base_per_face + extra_per_face)
    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)

# --- Интерфейс Gradio ---
custom_css = """
.gradio-container { max-width: 1200px; margin: 0 auto; }
.step-card { background: radial-gradient(circle at top left, #0b1220, #020617); border-radius: 18px; padding: 18px 20px; }
.face-thumb img { width: 100%; height: 100%; object-fit: cover; border-radius: 8px; }
.progress-ring { width: 120px; height: 120px; border-radius: 50%; display: flex; align-items: center; justify-content: center; }
.progress-ring-text { font-size: 14px; color: #e5e7eb; text-align: center; }
"""

premium_theme = gr.themes.Default(
    primary_hue="violet",
    neutral_hue="slate",
).set(
    body_text_size="16px",
    button_large_padding="12px 22px",
)

with gr.Blocks(title="FaceSwap Pro", css=custom_css) as demo:
    lang_state = gr.State("ru")
    lang_radio = gr.Radio(choices=["RU", "EN"], value="RU", label="Язык / Language")
    title_md = gr.Markdown(TEXTS["ru"]["title_md"])

    with gr.Row():
        with gr.Column():
            donor_img = gr.Image(label="Фото донора", type="numpy")
            donor_previews = [gr.Image(visible=False, interactive=False, width=64, height=64) for _ in range(MAX_PREVIEWS)]
            donor_choice = gr.Radio(label="Выберите лицо-донора", choices=[])
        with gr.Column():
            target_img = gr.Image(label="Целевое фото", type="numpy")
            target_previews = [gr.Image(visible=False, interactive=False, width=64, height=64) for _ in range(MAX_PREVIEWS)]
            target_choices = gr.CheckboxGroup(label="Выберите лица для замены", choices=[])

    with gr.Row():
        use_enh = gr.Checkbox(label="Улучшить качество результата (GFPGAN)", value=True)
        fmt = gr.Dropdown(label="Формат файла для скачивания", choices=["png", "jpeg", "webp"], value="png")
        run_btn = gr.Button("Запустить замену", variant="primary")
        download_btn = gr.DownloadButton("Скачать результат")

    result_img = gr.Image(label="Результат", interactive=False)
    status_md = gr.Markdown("")
    before_after = gr.Gallery(label="До / После", columns=2)

    def switch_language(choice):
        lang = "ru" if choice == "RU" else "en"
        t = TEXTS[lang]
        return (
            lang,
            gr.update(value=t["title_md"]),
            gr.update(label=t["step1_input_label"]),
            gr.update(label=t["step1_donor_choice_label"]),
            gr.update(label=t["step2_input_label"]),
            gr.update(label=t["step2_target_choices_label"]),
            gr.update(label=t["use_enh_label"]),
            gr.update(label=t["fmt_label"]),
            gr.update(value=t["run_btn"]),
            gr.update(value=t["download_btn"]),
        )

    lang_radio.change(
        fn=switch_language,
        inputs=lang_radio,
        outputs=[
            lang_state,
            title_md,
            donor_img,
            donor_choice,
            target_img,
            target_choices,
            use_enh,
            fmt,
            run_btn,
            download_btn,
        ],
    )

    donor_img.change(
        fn=update_donor_faces,
        inputs=[donor_img, lang_state],
        outputs=donor_previews + [donor_choice],
    )

    target_img.change(
        fn=update_target_faces,
        inputs=[target_img, use_enh, lang_state],
        outputs=target_previews + [target_choices],
    )

demo.launch(
    share=False,
    ssr_mode=False,
    theme=premium_theme,
    css=custom_css
)