Update api/seedvr_server.py

api/seedvr_server.py

@@ -3,31 +3,74 @@
 import os
 import sys
 import shutil
-import mimetypes
 import time
-import subprocess # Necessário para clonar o repositório na configuração inicial
+import subprocess
+import multiprocessing as mp
+import queue
 from pathlib import Path
 from typing import Optional, Callable
 from types import SimpleNamespace
 
+# --- Importações diretas de bibliotecas e do repositório SeedVR ---
+import torch
+import cv2
+import numpy as np
 from huggingface_hub import hf_hub_download
 
-# Adiciona dinamicamente o caminho do repositório clonado ao sys.path.
-# Isso é crucial para que a importação do 'inference_cli' funcione.
+# Adiciona o caminho do repositório para importar os módulos do 'src'
 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))
 
-# 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
+    from src.core.generation import generation_loop
+    from src.core.model_manager import configure_runner
+    from src.utils.downloads import download_weight
 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.")
+    print(f"ERRO FATAL: Não foi possível importar os módulos do SeedVR a partir de '{SEEDVR_REPO_PATH}'.")
     raise e
 
+# --- Função do Worker (definida fora da classe para ser 'picklable' pelo multiprocessing) ---
+
+def _worker_entry_point(proc_idx, device_id, frames_np, shared_args, return_queue, progress_queue=None):
+    """
+    Ponto de entrada para cada processo filho (worker).
+    Esta função executa em um processo separado e em uma GPU designada.
+    """
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
+    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
+
+    # Imports são feitos aqui para garantir um ambiente limpo para cada processo
+    import torch
+    
+    frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
+    
+    local_progress_callback = None
+    if progress_queue:
+        def callback_wrapper(batch_idx, total_batches, current_frames, message):
+            progress_queue.put((proc_idx, batch_idx, total_batches, message))
+        local_progress_callback = callback_wrapper
+
+    try:
+        # Cada worker configura seu próprio 'runner' para sua GPU
+        runner = configure_runner(shared_args["model"], shared_args["model_dir"], shared_args["preserve_vram"], shared_args["debug"])
+        
+        result_tensor = generation_loop(
+            runner=runner, images=frames_tensor, cfg_scale=shared_args["cfg_scale"],
+            seed=shared_args["seed"], res_w=shared_args["res_w"], batch_size=shared_args["batch_size"],
+            preserve_vram=shared_args["preserve_vram"], temporal_overlap=shared_args["temporal_overlap"],
+            debug=shared_args["debug"],
+            progress_callback=local_progress_callback 
+        )
+        return_queue.put((proc_idx, result_tensor.cpu().numpy()))
+    except Exception as e:
+        import traceback
+        error_msg = f"ERROR in worker {proc_idx}: {e}\n{traceback.format_exc()}"
+        print(error_msg)
+        if progress_queue: progress_queue.put((proc_idx, -1, -1, error_msg))
+        return_queue.put((proc_idx, error_msg))
+
+
 class SeedVRServer:
     def __init__(self, **kwargs):
         """
@@ -41,12 +84,12 @@ class SeedVRServer:
         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"))
         
-        print("🚀 SeedVRServer (Modo de Chamada Direta) inicializando...")
+        print("🚀 SeedVRServer (Arquitetura Integrada) inicializando...")
         for p in [self.CKPTS_ROOT, self.OUTPUT_ROOT, self.INPUT_ROOT, self.HF_HOME_CACHE]:
             p.mkdir(parents=True, exist_ok=True)
         
         self.setup_dependencies()
-        print("✅ SeedVRServer (Modo de Chamada Direta) pronto.")
+        print("✅ SeedVRServer (Arquitetura Integrada) pronto.")
 
     def setup_dependencies(self):
         """ Garante que o repositório e os modelos estão presentes. """
@@ -57,7 +100,6 @@ class SeedVRServer:
         """ 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.")
@@ -79,65 +121,135 @@ class SeedVRServer:
                     cache_dir=str(self.HF_HOME_CACHE), token=os.getenv("HF_TOKEN")
                 )
         print("[SeedVRServer] Checkpoints (FP16) estão no local correto.")
-        
-    def run_inference_direct(
+
+    def run_inference(
         self,
         file_path: str, *,
         seed: int, res_h: int, res_w: int, sp_size: int,
         fps: Optional[float] = None, progress: Optional[Callable] = None
     ) -> str:
         """
-        Executa a inferência diretamente no mesmo processo e retorna o caminho do arquivo de saída.
+        Método público principal para executar a inferência.
         """
-        # 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
-        )
-
         try:
-            # Informa a UI que o processo começou.
-            if progress:
-                progress(0.01, "Initializing...")
+            if progress: progress(0.05, "Extracting frames...")
+            frames_tensor, original_fps = self._extract_frames(file_path, debug=True)
             
-            # 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)
+            if progress: progress(0.1, "Starting parallel processing...")
             
-            # Informa a UI que a inferência terminou e o salvamento vai começar.
-            if progress:
-                progress(0.95, "Saving the final video...")
-
-            # Define o FPS final: usa o valor da UI ou o original do vídeo de entrada.
+            # Chama o método de processamento paralelo
+            result_tensor = self._process_in_parallel(
+                frames_tensor=frames_tensor,
+                seed=seed, resolution=res_h, batch_size=sp_size,
+                progress_callback=progress
+            )
+            
+            if progress: progress(0.95, "Saving the final video...")
             final_fps = fps if fps and fps > 0 else original_fps
-            save_frames_to_video(result_tensor, str(output_filepath), final_fps, args.debug)
+            self._save_video(result_tensor, str(output_filepath), final_fps, debug=True)
             
             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}")
+            print(f"❌ Error during inference execution: {e}")
             import traceback
             traceback.print_exc()
-            # Propaga o erro para a UI do Gradio, que o exibirá de forma amigável.
             raise
-            
+
+    def _extract_frames(self, video_path: str, debug: bool = False):
+        """Método privado para extrair quadros de um vídeo."""
+        # (Este é o código da função extract_frames_from_video, agora como um método de classe)
+        if debug: print(f"🎬 Extracting frames from video: {video_path}")
+        if not os.path.exists(video_path): raise FileNotFoundError(f"Video file not found: {video_path}")
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened(): raise ValueError(f"Cannot open video file: {video_path}")
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        frames = []
+        while True:
+            ret, frame = cap.read()
+            if not ret: break
+            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frame = frame.astype(np.float32) / 255.0
+            frames.append(frame)
+        cap.release()
+        if not frames: raise ValueError(f"No frames extracted from video: {video_path}")
+        return torch.from_numpy(np.stack(frames)).to(torch.float16), fps
+
+    def _save_video(self, frames_tensor: torch.Tensor, output_path: str, fps: float, debug: bool = False):
+        """Método privado para salvar um tensor de quadros em um arquivo de vídeo."""
+        # (Este é o código da função save_frames_to_video, agora como um método de classe)
+        if debug: print(f"🎬 Saving {frames_tensor.shape[0]} frames to video: {output_path}")
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8)
+        _, H, W, _ = frames_np.shape
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
+        for frame in frames_np:
+            out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
+        out.release()
+        if debug: print(f"✅ Video saved successfully: {output_path}")
+
+    def _process_in_parallel(self, frames_tensor: torch.Tensor, seed: int, resolution: int, batch_size: int, progress_callback: Optional[Callable] = None):
+        """
+        Método privado que gerencia a lógica de multiprocessamento.
+        """
+        device_list = list(range(self.NUM_GPUS_TOTAL))
+        num_devices = len(device_list)
+        chunks = torch.chunk(frames_tensor, num_devices, dim=0)
+
+        manager = mp.Manager()
+        return_queue = manager.Queue()
+        progress_queue = manager.Queue() if progress_callback else None
+        workers = []
+
+        shared_args = {
+            "model": "seedvr2_ema_3b_fp16.safetensors", "model_dir": str(self.CKPTS_ROOT),
+            "preserve_vram": True, "debug": True, "cfg_scale": 1.0,
+            "seed": seed, "res_w": resolution, "batch_size": batch_size, "temporal_overlap": 0,
+        }
+
+        for idx, device_id in enumerate(device_list):
+            p = mp.Process(target=_worker_entry_point, args=(idx, device_id, chunks[idx].cpu().numpy(), shared_args, return_queue, progress_queue))
+            p.start()
+            workers.append(p)
+
+        results_np = [None] * num_devices
+        finished_workers_count = 0
+        worker_progress = [0.0] * num_devices
+        
+        while finished_workers_count < num_devices:
+            if progress_queue:
+                while not progress_queue.empty():
+                    try:
+                        proc_idx, batch_idx, total_batches, message = progress_queue.get_nowait()
+                        if batch_idx == -1: raise RuntimeError(f"Worker {proc_idx} error: {message}")
+                        if total_batches > 0: worker_progress[proc_idx] = batch_idx / total_batches
+                        total_progress = sum(worker_progress) / num_devices
+                        progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: {message}")
+                    except queue.Empty:
+                        break
+
+            try:
+                proc_idx, result = return_queue.get(timeout=0.2)
+                if isinstance(result, str) and result.startswith("ERROR"):
+                     raise RuntimeError(f"Worker {proc_idx} failed: {result}")
+                results_np[proc_idx] = result
+                worker_progress[proc_idx] = 1.0
+                finished_workers_count += 1
+                if progress_callback:
+                    total_progress = sum(worker_progress) / num_devices
+                    progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: Completed!")
+            except queue.Empty:
+                pass
+
+        for p in workers: p.join()
+
+        if any(r is None for r in results_np):
+            raise RuntimeError("One or more workers failed to return a result.")
+
+        return torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)