import gradio as gr
import subprocess
import os

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

    print(f"Текущая рабочая директория: {os.getcwd()}")
    print(f"Содержимое текущей директории: {os.listdir('.')}")
    print(f"Проверка наличия 'inference.py': {os.path.exists('inference.py')}")

    video_path = video  
    audio_path = audio  

    print(f"Путь к видео: {video_path}")
    print(f"Путь к аудио: {audio_path}")

    output_dir = "outputs"
    os.makedirs(output_dir, exist_ok=True)

    output_path = os.path.join(output_dir, "output.mp4")
    print(f"Путь к выходному файлу: {output_path}")

    args = [
    "--checkpoint_path", f"checkpoints/{checkpoint}.pth",
    "--segmentation_path", "checkpoints/face_segmentation.pth",    
    "--no_seg",  
    "--no_sr",      
    "--face", video_path,
    "--audio", audio_path,
    "--outfile", output_path,
    "--resize_factor", "2", 
    "--face_det_batch_size", "4",    
    "--wav2lip_batch_size", "64",    
    "--fps", "30",                    
    
    "--pads", str(pad_top), str(pad_bottom), str(pad_left), str(pad_right)
    ]

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

    if save_as_video:
        args.append("--save_as_video")

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

    if not os.path.exists(output_path):
        print("Выходной файл не существует.")
        return "Не удалось создать выходное видео."

    print(f"Выходной файл создан по пути: {output_path}")
    return output_path  

with gr.Blocks() as ui:
    gr.Markdown("## Lypsinc")
    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", visible=False)
            no_smooth = gr.Checkbox(label="Без сглаживания", value=False)
            resize_factor = gr.Slider(minimum=1, maximum=4, step=1, label="Фактор изменения размера", value=2)
    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="Отступ снизу")
            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="Отступ справа")
            save_as_video = gr.Checkbox(label="Сохранять как видео", value=True)  
            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, save_as_video], 
        outputs=result,
        concurrency_limit=30  
    )

ui.launch(debug=True)