Update app_seedvr.py

app_seedvr.py

@@ -1,145 +1,120 @@
-#!/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.
-"""
+# app_seedvr.py
 
 import os
-import mimetypes
 from pathlib import Path
 from typing import Optional
-
 import gradio as gr
+import cv2
 
-# Importa a classe do servidor que gerencia a lógica de backend.
-# A inicialização do servidor (download de modelos, etc.) acontece aqui.
 try:
+    # Importa a classe de servidor que agora é uma biblioteca local
     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")
+except ImportError as e:
+    print(f"ERRO FATAL: Não foi possível importar o SeedVRServer. Detalhes: {e}")
+    # Se a importação falhar, a aplicação não pode continuar.
     raise
 
-# Cria uma instância única e persistente do servidor.
+# Cria uma instância única do servidor. A inicialização (clonar repo, baixar modelos) acontece aqui.
 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
+    if not path: return False
+    import mimetypes
     mime, _ = mimetypes.guess_type(path)
-    return (mime or "").startswith("image") or path.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
+    return (mime or "").startswith("video")
 
+def _extract_first_frame(video_path: str) -> Optional[str]:
+    """Extrai o primeiro frame de um vídeo e o salva como uma imagem JPG."""
+    if not video_path or not os.path.exists(video_path): return None
+    try:
+        vid_cap = cv2.VideoCapture(video_path)
+        if not vid_cap.isOpened(): return None
+        success, image = vid_cap.read()
+        vid_cap.release()
+        if not success: return None
+        
+        # Salva o frame no mesmo diretório do vídeo, com extensão .jpg
+        image_path = Path(video_path).with_suffix(".jpg")
+        cv2.imwrite(str(image_path), image)
+        return str(image_path)
+    except Exception as e:
+        print(f"Erro ao extrair o primeiro frame: {e}")
+        return None
 
 def ui_infer(
     input_path: Optional[str],
-    seed: int,
-    res_h: int,
-    res_w: int,
-    sp_size: int,
-    fps: float,
+    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".
+    Função de callback principal do Gradio. Agora chama a lógica de inferência diretamente.
     """
-    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
+    if not input_path:
+        gr.Warning("Por favor, faça o upload de um arquivo.")
+        return None, None, None
 
+    was_input_video = _is_video(input_path)
+    
     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(
+        # Desabilita o botão enquanto processa
+        yield gr.update(interactive=False, value="Processando..."), None, None, None
+
+        # Chama o método direto do servidor, passando o objeto de progresso do Gradio
+        video_result_path = server.run_inference_direct(
             file_path=input_path,
             seed=int(seed),
             res_h=int(res_h),
             res_w=int(res_w),
             sp_size=int(sp_size),
-            fps=float(fps) if fps and fps > 0 else None, # Passa None se o campo for 0 ou vazio
+            fps=float(fps) if fps and fps > 0 else None,
+            progress=progress,
         )
         
-        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
+        progress(1.0, desc="Concluído!")
+
+        final_image, final_video = None, None
+        if was_input_video:
+            final_video = video_result_path
+        else: # Se a entrada foi uma imagem
+            final_image = _extract_first_frame(video_result_path)
+            final_video = video_result_path
+            
+        # Retorna o resultado e reabilita o botão
+        yield (
+            gr.update(interactive=True, value="Restaurar Mídia"),
+            gr.update(value=final_image, visible=final_image is not None),
+            gr.update(value=final_video, visible=final_video is not None),
+            gr.update(value=video_result_path, visible=video_result_path is not None)
+        )
 
     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 ---
+        error_message = f"A inferência falhou: {e}"
+        gr.Error(error_message)
+        print(error_message)
+        import traceback
+        traceback.print_exc()
+        # Limpa os resultados e reabilita o botão em caso de erro
+        yield gr.update(interactive=True, value="Restaurar Mídia"), None, None, None
+
+# --- Construção da Interface Gráfica ---
 with gr.Blocks(title="SeedVR (Aduc-SDR)", theme=gr.themes.Soft()) as demo:
-    gr.HTML(
-        """
+    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")
+            inp = gr.File(label="Arquivo de Entrada (Vídeo .mp4 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.")
-
+                    fps_out = gr.Number(label="FPS de Saída (para Vídeos)", value=24, precision=0, info="0 para usar o FPS original.")
                 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)
@@ -150,30 +125,20 @@ with gr.Blocks(title="SeedVR (Aduc-SDR)", theme=gr.themes.Soft()) as demo:
 
         with gr.Column(scale=2):
             gr.Markdown("### Resultado")
-            out_image = gr.Image(label="Resultado (Imagem)", show_download_button=True)
+            out_image = gr.Image(label="Resultado (Imagem)", show_download_button=True, type="filepath", visible=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")
+            out_download = gr.File(label="Baixar Resultado (Vídeo)")
 
-    # Define a ação do botão
+    # A função click agora é um gerador.
     run_button.click(
         fn=ui_infer,
         inputs=[inp, seed, res_h, res_w, sp_size, fps_out],
-        outputs=[out_image, out_video, out_download],
+        outputs=[run_button, 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,
+        show_error=True
     )