app.py

import gradio as gr
import torch
import os
import subprocess
import tempfile
from inference import main as inference_main  # Предполагается, что вы интегрируете функцию main из inference.py
import shutil

def generate(video, audio, checkpoint, no_smooth, resize_factor, pad_top, pad_bottom, pad_left, pad_right):
    if video is None or audio is None or checkpoint is None:
        return "Пожалуйста, загрузите видео/изображение и аудио файл, а также выберите чекпойнт."

    # Поскольку Gradio возвращает пути к файлам, используем их напрямую
    video_path = video  # Уже строка с путем к видео
    audio_path = audio  # Уже строка с путем к аудио

    # Создание временной директории для сохранения выходного видео
    with tempfile.TemporaryDirectory() as temp_dir:
        # Создание поддиректории 'temp' внутри temp_dir для inference.py
        temp_subdir = os.path.join(temp_dir, 'temp')
        os.makedirs(temp_subdir, exist_ok=True)

        # Определение выходного файла
        output_path = os.path.join(temp_dir, "output.mp4")

        # Подготовка аргументов для инференса
        args = [
            "--checkpoint_path", f"checkpoints/{checkpoint}.pth",
            "--segmentation_path", "checkpoints/face_segmentation.pth",
            "--enhance_face", "gfpgan",
            "--face", video_path,
            "--audio", audio_path,
            "--outfile", output_path,
            "--resize_factor", str(resize_factor),
            "--pads", str(pad_top), str(pad_bottom), str(pad_left), str(pad_right)
        ]

        if no_smooth:
            args.append("--nosmooth")

        try:
            # Вызов команды с установкой рабочей директории в temp_dir
            cmd = ["python", "inference.py"] + args
            subprocess.run(cmd, check=True, cwd=temp_dir)
        except subprocess.CalledProcessError as e:
            return f"Произошла ошибка при обработке: {e}"

        # Проверка наличия выходного файла
        if not os.path.exists(output_path):
            return "Не удалось создать выходное видео."

        # Возвращаем путь к выходному файлу
        return output_path  # Gradio автоматически обработает путь и выведет видео

with gr.Blocks() as ui:
    gr.Markdown("## Wav2Lip - Синхронизация губ в видео")
    with gr.Row():
        video = gr.File(label="Видео или Изображение", type="filepath")
        audio = gr.File(label="Аудио", type="filepath")
        with gr.Column():
            checkpoint = gr.Radio(["wav2lip", "wav2lip_gan"], label="Чекпойнт", value="wav2lip_gan")
            no_smooth = gr.Checkbox(label="Без сглаживания", value=False)
            resize_factor = gr.Slider(minimum=1, maximum=4, step=1, label="Фактор изменения размера", value=1)
    with gr.Row():
        with gr.Column():
            pad_top = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Отступ сверху")
            pad_bottom = gr.Slider(minimum=0, maximum=50, step=1, value=10, label="Отступ снизу (рекомендуется 20 для включения подбородка)")
            pad_left = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Отступ слева")
            pad_right = gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Отступ справа")
            generate_btn = gr.Button("Сгенерировать")
        with gr.Column():
            result = gr.Video(label="Результат")

    generate_btn.click(
        generate, 
        inputs=[video, audio, checkpoint, no_smooth, resize_factor, pad_top, pad_bottom, pad_left, pad_right], 
        outputs=result,
        concurrency_limit=1  # Устанавливаем лимит на количество одновременно выполняемых операций
    )

ui.launch(debug=True)