app_seedvr.py

#!/usr/bin/env python3
"""
SeedVR UI (Gradio) - Interface de Usuário para Restauração de Mídia

- Permite o upload de um único arquivo de vídeo (.mp4) ou imagem.
- Oferece controle sobre parâmetros de geração como seed, resolução,
  paralelismo e FPS de saída para vídeos.
- Delega a execução da inferência para a classe `SeedVRServer`, que gerencia
  o ambiente e a chamada ao `torchrun`.
- Possui uma lógica de retorno robusta para exibir o resultado correto (imagem ou vídeo)
  ou notificar o usuário se nenhum resultado for encontrado.
"""

import os
import mimetypes
from pathlib import Path
from typing import Optional

import gradio as gr

# Importa a classe do servidor que gerencia a lógica de backend.
# A inicialização do servidor (download de modelos, etc.) acontece aqui.
try:
    from api.seedvr_server import SeedVRServer
except ImportError:
    print("ERRO FATAL: Não foi possível importar o SeedVRServer. Verifique o caminho em services/seed_server.py")
    raise

# Cria uma instância única e persistente do servidor.
server = SeedVRServer()

# Define o caminho de saída para referência, caso seja necessário.
OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/outputs"))


def _is_video(path: str) -> bool:
    """Verifica se um caminho de arquivo corresponde a um tipo de vídeo."""
    if not path:
        return False
    mime, _ = mimetypes.guess_type(path)
    # Verifica tanto o MIME type quanto a extensão do arquivo para maior robustez.
    return (mime or "").startswith("video") or path.lower().endswith((".mp4", ".mov", ".avi", ".mkv"))


def _is_image(path: str) -> bool:
    """Verifica se um caminho de arquivo corresponde a um tipo de imagem."""
    if not path:
        return False
    mime, _ = mimetypes.guess_type(path)
    return (mime or "").startswith("image") or path.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))


def ui_infer(
    input_path: Optional[str],
    seed: int,
    res_h: int,
    res_w: int,
    sp_size: int,
    fps: float,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Função principal de callback do Gradio. Acionada pelo botão "Restaurar".
    """
    progress(0.1, desc="Validando entradas...")
    if not input_path or not Path(input_path).exists():
        gr.Warning("Arquivo de entrada ausente ou inválido. Por favor, faça o upload de um vídeo ou imagem.")
        return None, None, None # Retorna None para todas as saídas

    try:
        progress(0.4, desc="Enviando tarefa para o backend. A inferência pode levar vários minutos...")
        
        # Chama o método do servidor, passando todos os parâmetros da UI.
        video_out, image_out, out_dir = server.run_inference(
            filepath=input_path,
            seed=int(seed),
            resh=int(res_h),
            resw=int(res_w),
            spsize=int(sp_size),
            fps=float(fps) if fps and fps > 0 else None, # Passa None se o campo for 0 ou vazio
        )
        
        progress(0.9, desc="Inferência concluída. Processando resultados...")

        # Lógica de retorno robusta
        final_image = None
        final_video = None
        result_file_for_download = None

        is_video_input = _is_video(input_path)

        if is_video_input:
            # Se a entrada foi um vídeo, a saída principal esperada é um vídeo.
            if video_out:
                final_video = video_out
                result_file_for_download = video_out
            elif image_out: # Fallback
                gr.Info("A inferência de vídeo resultou em uma única imagem.")
                final_image = image_out
                result_file_for_download = image_out
        else: # A entrada foi uma imagem
            # Se a entrada foi uma imagem, a saída principal esperada é uma imagem.
            if image_out:
                final_image = image_out
                result_file_for_download = image_out
            elif video_out: # Fallback
                gr.Info("A inferência de imagem resultou em um vídeo.")
                final_video = video_out
                result_file_for_download = video_out

        if not result_file_for_download:
            gr.Warning("A inferência foi concluída, mas nenhum arquivo de mídia de saída foi encontrado no diretório de resultados.")

        return final_image, final_video, result_file_for_download

    except Exception as e:
        print(f"[UI ERROR] A inferência falhou: {e}")
        # Exibe uma notificação de erro clara na interface do Gradio.
        gr.Error(f"Erro na Inferência: {e}")
        return None, None, None


# --- Construção da Interface Gráfica com Gradio ---
with gr.Blocks(title="SeedVR (Aduc-SDR)", theme=gr.themes.Soft()) as demo:
    gr.HTML(
        """
        <div style='text-align:center; margin-bottom: 20px;'>
            <h1>SeedVR - Restauração de Imagem e Vídeo</h1>
            <p>Implementação com backend Aduc-SDR</p>
        </div>
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            inp = gr.File(label="Arquivo de Entrada (Vídeo ou Imagem)", type="filepath")
            
            with gr.Accordion("Parâmetros de Geração", open=True):
                with gr.Row():
                    seed = gr.Number(label="Seed", value=42, precision=0)
                    fps_out = gr.Number(label="FPS de Saída", value=24, precision=0, info="Apenas para vídeos.")

                with gr.Row():
                    res_h = gr.Number(label="Altura (Height)", value=720, precision=0)
                    res_w = gr.Number(label="Largura (Width)", value=1280, precision=0)

                sp_size = gr.Slider(label="Paralelismo de Sequência (sp_size)", minimum=1, maximum=160, step=4, value=4, info="Para vídeos em multi-GPU. Use 1 para imagens.")

            run_button = gr.Button("Restaurar Mídia", variant="primary")

        with gr.Column(scale=2):
            gr.Markdown("### Resultado")
            out_image = gr.Image(label="Resultado (Imagem)", show_download_button=True)
            out_video = gr.Video(label="Resultado (Vídeo)")
            # Componente de arquivo para download, pode ser visível ou não.
            # Deixá-lo visível pode ser uma boa alternativa se a pré-visualização falhar.
            out_download = gr.File(label="Baixar Resultado")

    # Define a ação do botão
    run_button.click(
        fn=ui_infer,
        inputs=[inp, seed, res_h, res_w, sp_size, fps_out],
        outputs=[out_image, out_video, out_download],
    )
    
    gr.Markdown("---")
    gr.Markdown("### Exemplos")
    # Nota: Exemplos precisam de arquivos presentes no contêiner para funcionar.
    # Adicione os arquivos de exemplo ao seu Dockerfile se desejar usá-los.
    # gr.Examples(...)

# --- Ponto de Entrada da Aplicação ---
if __name__ == "__main__":
    demo.launch(
        server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
        allowed_paths=[str(OUTPUT_ROOT), str(os.getenv("INPUT_ROOT", "/app/inputs"))],
        show_error=True,
    )