Update api/seedvr_server.py

api/seedvr_server.py

@@ -1,248 +1,142 @@
-# app_seedvr.py
+# api/seedvr_server.py
 
 import os
 import sys
+import shutil
+import mimetypes
+import time
+import subprocess # Necessário para clonar o repositório na configuração inicial
 from pathlib import Path
-from typing import Optional
-import gradio as gr
-import cv2
+from typing import Optional, Callable
+from types import SimpleNamespace
 
-# --- INTEGRAÇÃO COM A LÓGICA DO SERVIDOR ---
-try:
-    # Importa a classe SeedVRServer que agora atua como nossa biblioteca de inferência.
-    from api.seedvr_server import SeedVRServer
-except ImportError as e:
-    print(f"ERRO FATAL: Não foi possível importar o SeedVRServer. Detalhes: {e}")
-    # A aplicação não pode rodar sem a lógica do servidor.
-    raise
-
-# --- INICIALIZAÇÃO ---
-# Cria uma instância única e persistente do servidor.
-# A inicialização (clonar repo, baixar modelos) acontece apenas uma vez, no início.
-server = SeedVRServer()
-
-# --- FUNÇÕES AUXILIARES ---
-
-def _is_video(path: str) -> bool:
-    """Verifica se um caminho de arquivo corresponde a um tipo de vídeo."""
-    if not path: return False
-    import mimetypes
-    mime, _ = mimetypes.guess_type(path)
-    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():
-            print(f"Erro: Não foi possível abrir o vídeo em {video_path}")
-            return None
-        success, image = vid_cap.read()
-        vid_cap.release()
-        if not success:
-            print(f"Erro: Não foi possível ler o primeiro frame de {video_path}")
-            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 on_file_upload(file_obj):
-    """
-    Callback acionado quando o usuário faz o upload de um arquivo.
-    Verifica se o arquivo é um vídeo e sugere um `sp_size` apropriado.
-    """
-    if file_obj is None:
-        # Limpa os resultados e o log se o arquivo for removido
-        return gr.update(value=1), None, None, None, gr.update(value=None, visible=False)
-    
-    if _is_video(file_obj.name):
-        # Para vídeos, sugere um valor padrão para multi-GPU e torna o slider interativo
-        return gr.update(value=4, interactive=True), None, None, None, gr.update(value=None, visible=False)
-    else:
-        # Para imagens, trava o valor em 1
-        return gr.update(value=1, interactive=False), None, None, None, gr.update(value=None, visible=False)
-
-# --- FUNÇÃO PRINCIPAL DE INFERÊNCIA DA UI ---
+from huggingface_hub import hf_hub_download
 
-def run_inference_ui(
-    input_file_path: Optional[str],
-    resolution: str,
-    sp_size: int,
-    fps: float,
-    progress=gr.Progress(track_tqdm=True)
-):
-    """
-    A função de callback principal do Gradio. Usa geradores (`yield`)
-    para permitir atualizações da UI em tempo real durante a tarefa de longa duração.
-    """
-    # 1. Estado Inicial e Validação
-    # No início, desabilita o botão, limpa resultados anteriores e mostra a janela de log.
-    yield (
-        gr.update(interactive=False, value="Processing... 🚀"),
-        gr.update(value=None, visible=False),
-        gr.update(value=None, visible=False),
-        gr.update(value=None, visible=False),
-        gr.update(value="▶ Starting inference process...\n", visible=True)
-    )
+# Adiciona dinamicamente o caminho do repositório clonado ao sys.path.
+# Isso é crucial para que a importação do 'inference_cli' funcione.
+SEEDVR_REPO_PATH = Path(os.getenv("SEEDVR_ROOT", "/data/SeedVR"))
+if str(SEEDVR_REPO_PATH) not in sys.path:
+    # Insere no início da lista para garantir prioridade de importação.
+    sys.path.insert(0, str(SEEDVR_REPO_PATH))
 
-    if not input_file_path:
-        gr.Warning("Please upload a media file first.")
-        # Reabilita o botão e esconde os componentes de saída
-        yield (gr.update(interactive=True, value="Restore Media"), None, None, None, gr.update(visible=False))
-        return
-
-    log_buffer = ["▶ Starting inference process...\n"]
-    last_log_message = ""
-    was_input_video = _is_video(input_file_path)
-
-    try:
-        # Define um callback que será chamado pelo backend para atualizar o progresso e o log
-        def progress_callback_wrapper(step: float, desc: str):
-            """ Wrapper para formatar logs e atualizar o progresso. """
-            nonlocal last_log_message
-            # Só adiciona ao log se a mensagem for nova, para evitar poluição visual
-            if desc != last_log_message:
-                log_buffer.append(f"⏳ {desc}\n")
-                last_log_message = desc
-            # Atualiza o objeto de progresso do Gradio
-            progress(step, desc=desc)
+# Tenta importar as funções necessárias APÓS a modificação do path.
+# Se falhar, a aplicação não pode continuar.
+try:
+    from inference_cli import run_inference_logic, save_frames_to_video
+except ImportError as e:
+    print(f"ERRO FATAL: Não foi possível importar de 'inference_cli.py'.")
+    print(f"Verifique se o repositório em '{SEEDVR_REPO_PATH}' está correto e completo.")
+    raise e
 
-        # 2. Executa a Inferência
-        # Chama o método direto do servidor, passando o nosso callback.
-        video_result_path = server.run_inference_direct(
-            file_path=input_file_path,
-            seed=42, # Semente fixa conforme solicitado
-            res_h=int(resolution),
-            res_w=int(resolution), # Largura igual à altura
-            sp_size=int(sp_size),
-            fps=float(fps) if fps and fps > 0 else None,
-            progress=progress_callback_wrapper, # Passa nossa função de callback
-        )
-        
-        progress(1.0, desc="Complete!")
-        log_buffer.append("✅ Inference complete! Processing final output...\n")
+class SeedVRServer:
+    def __init__(self, **kwargs):
+        """
+        Inicializa o servidor, define os caminhos e prepara o ambiente.
+        """
+        self.SEEDVR_ROOT = SEEDVR_REPO_PATH
+        self.CKPTS_ROOT = Path("/data/seedvr_models_fp16") 
+        self.OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/outputs"))
+        self.INPUT_ROOT = Path(os.getenv("INPUT_ROOT", "/app/inputs"))
+        self.HF_HOME_CACHE = Path(os.getenv("HF_HOME", "/data/.cache/huggingface"))
+        self.REPO_URL = os.getenv("SEEDVR_GIT_URL", "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler")
+        self.NUM_GPUS_TOTAL = int(os.getenv("NUM_GPUS", "4"))
         
-        # 3. Processa e Exibe os Resultados
-        final_image, final_video = None, None
-        if was_input_video:
-            final_video = video_result_path
-            log_buffer.append("✅ Video result is ready.\n")
-        else: # Se a entrada foi uma imagem
-            final_image = _extract_first_frame(video_result_path)
-            final_video = video_result_path # Também disponibiliza o vídeo de 1 frame
-            log_buffer.append("✅ Image result extracted from video.\n")
+        print("🚀 SeedVRServer (Modo de Chamada Direta) inicializando...")
+        for p in [self.CKPTS_ROOT, self.OUTPUT_ROOT, self.INPUT_ROOT, self.HF_HOME_CACHE]:
+            p.mkdir(parents=True, exist_ok=True)
         
-        # Yield final para mostrar os resultados e reabilitar o botão
-        yield (
-            gr.update(interactive=True, value="Restore Media"),
-            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),
-            ''.join(log_buffer)
-        )
-
-    except Exception as e:
-        error_message = f"❌ Inference failed: {e}"
-        gr.Error(error_message)
-        log_buffer.append(f"\n{error_message}")
-        import traceback
-        traceback.print_exc()
+        self.setup_dependencies()
+        print("✅ SeedVRServer (Modo de Chamada Direta) pronto.")
+
+    def setup_dependencies(self):
+        """ Garante que o repositório e os modelos estão presentes. """
+        self._ensure_repo()
+        self._ensure_model()
+
+    def _ensure_repo(self) -> None:
+        """ Clona o repositório do SeedVR se ele não existir. """
+        if not (self.SEEDVR_ROOT / ".git").exists():
+            print(f"[SeedVRServer] Clonando repositório para {self.SEEDVR_ROOT}...")
+            # Usamos subprocess.run aqui porque é uma tarefa de inicialização única.
+            subprocess.run(["git", "clone", "--depth", "1", self.REPO_URL, str(self.SEEDVR_ROOT)], check=True)
+        else:
+            print("[SeedVRServer] Repositório SeedVR já existe.")
+
+    def _ensure_model(self) -> None:
+        """ Baixa os checkpoints do Hugging Face se não existirem localmente. """
+        print(f"[SeedVRServer] Verificando checkpoints (FP16) em {self.CKPTS_ROOT}...")
+        model_files = {
+            "seedvr2_ema_3b_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+            "ema_vae_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
+            "pos_emb.pt": "ByteDance-Seed/SeedVR2-3B",
+            "neg_emb.pt": "ByteDance-Seed/SeedVR2-3B"
+        }
+        for filename, repo_id in model_files.items():
+            if not (self.CKPTS_ROOT / filename).exists():
+                print(f"Baixando {filename} de {repo_id}...")
+                hf_hub_download(
+                    repo_id=repo_id, filename=filename, local_dir=str(self.CKPTS_ROOT),
+                    cache_dir=str(self.HF_HOME_CACHE), token=os.getenv("HF_TOKEN")
+                )
+        print("[SeedVRServer] Checkpoints (FP16) estão no local correto.")
         
-        # Yield para estado de erro: reabilita o botão e mostra o log com o erro
-        yield (
-            gr.update(interactive=True, value="Restore Media"),
-            None, None, None,
-            gr.update(value=''.join(log_buffer), visible=True)
-        )
-
-# --- LAYOUT DA INTERFACE GRÁFICA (GRADIO) ---
-
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="SeedVR Media Restoration") as demo:
-    # Cabeçalho
-    gr.Markdown(
+    def run_inference_direct(
+        self,
+        file_path: str, *,
+        seed: int, res_h: int, res_w: int, sp_size: int,
+        fps: Optional[float] = None, progress: Optional[Callable] = None
+    ) -> str:
         """
-        <div style='text-align: center; margin-bottom: 20px;'>
-            <h1>📸 SeedVR - Image & Video Restoration 🚀</h1>
-            <p>High-quality media upscaling powered by SeedVR-3B. Upload your file and see the magic.</p>
-        </div>
+        Executa a inferência diretamente no mesmo processo e retorna o caminho do arquivo de saída.
         """
-    )
+        # Cria um diretório de saída único para salvar o resultado.
+        out_dir = self.OUTPUT_ROOT / f"run_{int(time.time())}_{Path(file_path).stem}"
+        out_dir.mkdir(parents=True, exist_ok=True)
+        output_filepath = out_dir / f"result_{Path(file_path).stem}.mp4"
+
+        # Simula o objeto 'args' que a função de lógica do inference_cli espera.
+        # Usamos SimpleNamespace para criar um objeto simples com atributos.
+        args = SimpleNamespace(
+            video_path=file_path,
+            output=str(output_filepath),
+            model_dir=str(self.CKPTS_ROOT),
+            seed=seed,
+            resolution=res_h, # O script do SeedVR usa a altura (lado menor) como referência.
+            batch_size=sp_size,
+            model="seedvr2_ema_3b_fp16.safetensors",
+            preserve_vram=True,
+            debug=True, # Mantém o debug ativo para logs detalhados.
+            cuda_device=",".join(map(str, range(self.NUM_GPUS_TOTAL))),
+            skip_first_frames=0,
+            load_cap=0,
+            output_format='video' # Garante que sempre gere vídeo
+        )
 
-    with gr.Row():
-        # --- Coluna da Esquerda: Entradas e Controles ---
-        with gr.Column(scale=1):
-            gr.Markdown("### 1. Upload Media")
-            input_media = gr.File(label="Input File (Video or Image)", type="filepath", interactive=True)
+        try:
+            # Informa a UI que o processo começou.
+            if progress:
+                progress(0.01, "Initializing...")
             
-            gr.Markdown("### 2. Configure Settings")
-            with gr.Accordion("Generation Parameters", open=True):
-                resolution_select = gr.Dropdown(
-                    label="Resolution",
-                    choices=["480", "560", "720", "960", "1024"],
-                    value="480",
-                    info="Sets the output height and width to this value."
-                )
-                
-                sp_size_slider = gr.Slider(
-                    label="Frames per Batch (sp_size)",
-                    minimum=1, maximum=16, step=1, value=4,
-                    info="For multi-GPU videos. Automatically set to 1 for images."
-                )
-
-                fps_out = gr.Number(label="Output FPS (for Videos)", value=24, precision=0, info="Set to 0 to use the original FPS.")
-
-            run_button = gr.Button("Restore Media", variant="primary", icon="✨")
-
-        # --- Coluna da Direita: Resultados ---
-        with gr.Column(scale=2):
-            gr.Markdown("### 3. Results")
+            # Chama a função importada do script original, passando o callback de progresso.
+            # Este callback será chamado de dentro da lógica de multi-processamento.
+            result_tensor, original_fps, _, _ = run_inference_logic(args, progress_callback=progress)
             
-            # Janela de Log
-            log_window = gr.Textbox(
-                label="Inference Log 📝",
-                lines=8, max_lines=15,
-                interactive=False, visible=False, autoscroll=True
-            )
-
-            # Componentes de saída (começam invisíveis)
-            output_image = gr.Image(label="Image Result", show_download_button=True, type="filepath", visible=False)
-            output_video = gr.Video(label="Video Result", visible=False)
-            output_download = gr.File(label="Download Full Result (Video)", visible=False)
+            # Informa a UI que a inferência terminou e o salvamento vai começar.
+            if progress:
+                progress(0.95, "Saving the final video...")
 
-    # --- Rodapé ---
-    gr.Markdown(
-        """
-        ---
-        *Space and Docker were developed by Carlex.*
-        *Contact: Email: Carlex22@gmail.com | GitHub: [carlex22](https://github.com/carlex22)*
-        """
-    )
-    
-    # --- Lógica de Eventos da UI ---
-    
-    # Ao fazer upload de um arquivo, ajusta o slider `sp_size` e limpa saídas antigas.
-    input_media.upload(
-        fn=on_file_upload,
-        inputs=[input_media],
-        outputs=[sp_size_slider, output_image, output_video, output_download, log_window]
-    )
-    
-    # Ao clicar no botão, executa a função de inferência principal.
-    run_button.click(
-        fn=run_inference_ui,
-        inputs=[input_media, resolution_select, sp_size_slider, fps_out],
-        outputs=[run_button, output_image, output_video, output_download, log_window],
-    )
-
-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")),
-        show_error=True
-    )
+            # Define o FPS final: usa o valor da UI ou o original do vídeo de entrada.
+            final_fps = fps if fps and fps > 0 else original_fps
+            save_frames_to_video(result_tensor, str(output_filepath), final_fps, args.debug)
+            
+            print(f"✅ Video saved successfully to: {output_filepath}")
+            
+            # Retorna o caminho do arquivo gerado para a UI.
+            return str(output_filepath)
+            
+        except Exception as e:
+            print(f"❌ Error during direct inference execution: {e}")
+            import traceback
+            traceback.print_exc()
+            # Propaga o erro para a UI do Gradio, que o exibirá de forma amigável.
+            raise