mainmainmina / app.py
Uniaff's picture
Update app.py
95e31cf verified
raw
history blame
3.98 kB
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:
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:
# Вызов функции инференса через subprocess
cmd = ["python", "inference.py"] + args
subprocess.run(cmd, check=True)
except subprocess.CalledProcessError as e:
return f"Произошла ошибка при обработке: {e}"
# Проверка наличия выходного файла
if not os.path.exists(output_path):
return "Не удалось создать выходное видео."
# Возвращаем путь к выходному файлу
return output_path
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)