Update api/seedvr_server.py

api/seedvr_server.py

@@ -1,147 +1,128 @@
+# api/seedvr_server.py
+
 import os
-import shutil
-import subprocess
 import sys
-import time
+import shutil
 import mimetypes
+import time
 from pathlib import Path
 from typing import Optional, Tuple
+from types import SimpleNamespace
+
 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.
+SEEDVR_REPO_PATH = Path(os.getenv("SEEDVR_ROOT", "/data/SeedVR"))
+if str(SEEDVR_REPO_PATH) not in sys.path:
+    sys.path.insert(0, str(SEEDVR_REPO_PATH))
+
+# Tenta importar as funções necessárias APÓS a modificação do path.
+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'. Verifique se o repositório em {SEEDVR_REPO_PATH} está correto.")
+    raise e
+
 class SeedVRServer:
     def __init__(self, **kwargs):
-        self.SEEDVR_ROOT = Path(os.getenv("SEEDVR_ROOT", "/data/SeedVR"))
-        self.CKPTS_ROOT = Path("/data/seedvr_models_fp16")
+        """
+        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"))
-        # Use 8 por padrão, mas nunca maior que o visível no container/host
-        self.NUM_GPUS_TOTAL = min(int(os.getenv("NUM_GPUS", "8")),
-                                  int(os.getenv("MAX_VISIBLE_GPUS", "8")))
-        print("🚀 SeedVRServer (FP16) inicializando e preparando o ambiente...")
-
-        for p in [self.SEEDVR_ROOT.parent, self.CKPTS_ROOT, self.OUTPUT_ROOT, self.INPUT_ROOT, self.HF_HOME_CACHE]:
+        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...")
+        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 (FP16) pronto.")
+        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}...")
-            subprocess.run(["git", "clone", "--depth", "1", os.getenv("SEEDVR_GIT_URL",
-                            "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler"), str(self.SEEDVR_ROOT)], check=True)
+            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",
+            "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"),
+                    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.")
-
-    def _prepare_job(self, input_file: str) -> Tuple[Path, Path]:
-        ts = f"{int(time.time())}_{os.urandom(4).hex()}"
-        job_input_dir = self.INPUT_ROOT / f"job_{ts}"
-        out_dir = self.OUTPUT_ROOT / f"run_{ts}"
-        job_input_dir.mkdir(parents=True, exist_ok=True)
-        out_dir.mkdir(parents=True, exist_ok=True)
-        shutil.copy2(input_file, job_input_dir / Path(input_file).name)
-        return job_input_dir, out_dir
-
-    def _visible_devices_for(self, nproc: int) -> str:
-        # Mapeia 0..nproc-1 (lógico) para o espaço visível do container
-        return ",".join(str(i) for i in range(nproc))
-
-    def run_inference(
+        
+    def run_inference_direct(
         self,
-        file_path: str,
-        *,
-        seed: int,
-        res_h: int,
-        res_w: int,
-        sp_size: int,
-        fps: Optional[float] = None,
-    ) -> Tuple[Optional[str], Optional[str], Path]:
-
-        script = self.SEEDVR_ROOT / "inference_cli.py"
-        job_input_dir, out_dir = self._prepare_job(file_path)
-
-        media_type, _ = mimetypes.guess_type(file_path)
-        is_image = bool(media_type and media_type.startswith("image"))
-
-        # Política: 1 GPU para imagem, 8 GPUs (ou NUM_GPUS_TOTAL) para vídeo
-        effective_nproc = 1 if is_image else self.NUM_GPUS_TOTAL
-        effective_sp_size = 1 if is_image else sp_size
-
-        output_filename = f"result_{Path(file_path).stem}.mp4"
-        output_filepath = out_dir / output_filename
-
-        cmd = [
-            "torchrun",
-            "--standalone",
-            "--nnodes=1",
-            f"--nproc-per-node={effective_nproc}",
-            str(script),
-            "--video_path", str(file_path),
-            "--output", str(output_filepath),
-            "--model_dir", str(self.CKPTS_ROOT),
-            "--seed", str(seed),
-            "--resolution", str(res_h),
-            "--batch_size", str(effective_sp_size),
-            "--model", "seedvr2_ema_3b_fp16.safetensors",
-            "--preserve_vram",
-            "--debug",
-            "--output_format", "video",
-        ]
-        # Removido: --cuda_device ... (torchrun + LOCAL_RANK fará o binding correto)
-
-        env = os.environ.copy()
-        # Alinhar espaço lógico de devices com nproc
-        env["CUDA_VISIBLE_DEVICES"] = self._visible_devices_for(effective_nproc)
-        # Dicas úteis de debug (opcional):
-        # env["NCCL_DEBUG"] = "WARN"
-        # env["CUDA_LAUNCH_BLOCKING"] = "1"
-
-        print("[SeedVRServer] Comando:", " ".join(cmd))
-        print("[SeedVRServer] CUDA_VISIBLE_DEVICES:", env.get("CUDA_VISIBLE_DEVICES", ""))
+        file_path: str, *,
+        seed: int, res_h: int, res_w: int, sp_size: int,
+        fps: Optional[float] = None, progress=None
+    ) -> str:
+        """
+        Executa a inferência diretamente no mesmo processo e retorna o caminho do arquivo de saída.
+        """
+        out_dir = self.OUTPUT_ROOT / f"run_{int(time.time())}"
+        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 espera
+        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,
+            cuda_device=",".join(map(str, range(self.NUM_GPUS_TOTAL))),
+            skip_first_frames=0,
+            load_cap=0
+        )
 
         try:
-            subprocess.run(
-                cmd,
-                cwd=str(self.SEEDVR_ROOT),
-                check=True,
-                env=env,
-                stdout=sys.stdout,
-                stderr=sys.stderr,
-            )
-
-            if is_image:
-                # Se output_format=png no CLI, aqui poderia ser diretório; com "video" mantemos mp4, mas
-                # preservamos compatibilidade caso o CLI mude:
-                image_dir = output_filepath if output_filepath.suffix == "" else output_filepath.with_suffix("")
-                return str(image_dir), None, out_dir
-            else:
-                return None, str(output_filepath), out_dir
-
+            if progress:
+                progress(0.1, desc="Iniciando a lógica de inferência...")
+            
+            # Chama a função importada do script original
+            result_tensor, original_fps, _, _ = run_inference_logic(args, progress_callback=progress)
+            
+            if progress:
+                progress(0.9, desc="Salvando o vídeo resultante...")
+
+            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"✅ Vídeo salvo com sucesso em: {output_filepath}")
+            
+            return str(output_filepath)
+            
         except Exception as e:
-            print(f"[UI ERROR] A inferência falhou: {e}")
-            return None, None, out_dir
+            print(f"❌ Erro durante a execução direta da inferência: {e}")
+            import traceback
+            traceback.print_exc()
+            raise # Propaga o erro para a UI do Gradio, que o exibirá.