Update app_seedvr.py

app_seedvr.py

@@ -1,11 +1,14 @@
 #!/usr/bin/env python3
 """
-SeedVR UI (Gradio) — CLI torchrun
-
-- Upload único: vídeo (.mp4) ou imagem (.png/.jpg/.jpeg/.webp).
-- Parâmetros: seed, res_h, res_w, sp_size.
-- Executa via torchrun com NUM_GPUS (do ambiente).
-- Exibe vídeo se a entrada for vídeo; imagem se for imagem.
+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
@@ -15,24 +18,37 @@ from typing import Optional
 
 import gradio as gr
 
-from services.seed_server import SeedVRServer
+# 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 services.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
 
-# Instância única do servidor (clona repo, baixa modelo, cria symlink)
+# Cria uma instância única e persistente do servidor.
 server = SeedVRServer()
 
-# Paths padrão (para allowed_paths e debug)
+# Define o caminho de saída para referência, caso seja necessário.
 OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/outputs"))
-CKPTS_ROOT = Path(os.getenv("CKPTS_ROOT", "/app/ckpts/SeedVR2-3B"))
+
 
 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)
-    return (mime or "").startswith("video") or str(path).lower().endswith(".mp4")
+    # 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)
-    if mime and mime.startswith("image"):
-        return True
-    return str(path).lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
+    return (mime or "").startswith("image") or path.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
+
 
 def ui_infer(
     input_path: Optional[str],
@@ -40,78 +56,124 @@ def ui_infer(
     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.")
-        return None, None
-
-    is_vid = _is_video(input_path)
-    is_img = _is_image(input_path)
-    if not (is_vid or is_img):
-        gr.Warning("Tipo de arquivo não suportado. Envie .mp4, .png, .jpg, .jpeg ou .webp.")
-        return None, None
+        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:
-        video_out, image_out, _ = server.run_inference(
-            file_path=input_path,
+        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),
-            res_h=int(res_h),
-            res_w=int(res_w),
-            sp_size=int(sp_size),
+            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:
-        gr.Warning(f"Erro na inferência: {e}")
-        return None, None
-
-    if is_vid:
-        if video_out and Path(video_out).exists():
-            return None, video_out
-        if image_out and Path(image_out).exists():
-            return image_out, None
-        gr.Warning("Nenhum resultado encontrado.")
-        return None, None
-    else:
-        if image_out and Path(image_out).exists():
-            return image_out, None
-        if video_out and Path(video_out).exists():
-            return None, video_out
-        gr.Warning("Nenhum resultado encontrado.")
-        return None, None
-
-with gr.Blocks(title="SeedVR (CLI torchrun)") as demo:
-    gr.Markdown(
-        "\n".join([
-            "# SeedVR — Restauração (CLI torchrun)",
-            "- Envie um vídeo (.mp4) ou uma imagem (.png/.jpg/.jpeg/.webp).",
-            "- A execução utiliza torchrun com múltiplas GPUs.",
-        ])
+        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():
-        inp = gr.File(label="Entrada (vídeo .mp4 ou imagem)", type="filepath")
-
+    
     with gr.Row():
-        seed = gr.Number(label="Seed", value=int(os.getenv("SEED", "42")), precision=0)
-        res_h = gr.Number(label="Altura (res_h)", value=int(os.getenv("RES_H", "720")), precision=0)
-        res_w = gr.Number(label="Largura (res_w)", value=int(os.getenv("RES_W", "1280")), precision=0)
-        sp_size = gr.Number(label="sp_size", value=int(os.getenv("SP_SIZE", "4")), precision=0)
-
-    run = gr.Button("Restaurar", variant="primary")
-
-    out_image = gr.Image(label="Resultado (imagem)")
-    out_video = gr.Video(label="Resultado (vídeo)")
-
-    run.click(
-        ui_infer,
-        inputs=[inp, seed, res_h, res_w, sp_size],
-        outputs=[out_image, out_video],
+        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", os.getenv("PORT", "7860"))),
-        allowed_paths=[str(OUTPUT_ROOT), str(CKPTS_ROOT)],
+        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
+        allowed_paths=[str(OUTPUT_ROOT), str(os.getenv("INPUT_ROOT", "/app/inputs"))],
         show_error=True,
-    )
-    
+    )