Update services/vince_pool_manager.py

services/vince_pool_manager.py

@@ -11,11 +11,12 @@ from huggingface_hub import snapshot_download
 from omegaconf import OmegaConf, open_dict
 
 # --- Configurações Globais ---
+# Os diretórios são definidos para usar /data para persistência, espelhando o SeedVR.
 VINCIE_DIR = Path(os.getenv("VINCIE_DIR", "/data/VINCIE"))
 VINCE_GIT_URL = os.getenv("VINCE_GIT_URL", "https://github.com/ByteDance-Seed/VINCIE")
 VINCE_REPO_ID = os.getenv("VINCE_REPO_ID", "ByteDance-Seed/VINCIE-3B")
 VINCE_CKPT = Path(os.getenv("VINCE_CKPT", "/data/ckpt/VINCIE-3B"))
-HF_HOME_CACHE = os.getenv("HF_HOME")
+HF_HOME_CACHE = os.getenv("HF_HOME", "/data/.cache/huggingface")
 
 # --- Classe Worker (Gerencia uma única GPU de forma isolada) ---
 class VinceWorker:
@@ -30,44 +31,26 @@ class VinceWorker:
         print(f"[VinceWorker-{self.device_id_str}] Inicializado. Mapeado para o índice de GPU {self.gpu_index}.")
 
     def _execute_in_isolated_env(self, function_to_run, *args, **kwargs):
-        """
-        Wrapper que define CUDA_VISIBLE_DEVICES antes de executar uma função
-        e o limpa depois. Garante que o código VINCIE veja apenas uma GPU.
-        """
-        # Guarda a variável de ambiente original, se existir
+        """Wrapper que define CUDA_VISIBLE_DEVICES para isolar a visibilidade da GPU."""
         original_cuda_visible = os.environ.get('CUDA_VISIBLE_DEVICES')
-        
         try:
-            # Isola a visibilidade da GPU para este worker
             os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_index
-            
-            # Limpa qualquer cache de dispositivo que o PyTorch possa ter
             if torch.cuda.is_available():
-                torch.cuda.set_device(0) # Agora 'cuda:0' é a nossa GPU alvo
-
-            # Executa a função alvo (carregamento ou inferência)
+                torch.cuda.set_device(0) # 'cuda:0' agora é a nossa GPU alvo
             return function_to_run(*args, **kwargs)
-            
         finally:
-            # Restaura a variável de ambiente original
             if original_cuda_visible is not None:
                 os.environ['CUDA_VISIBLE_DEVICES'] = original_cuda_visible
-            else:
-                # Remove a variável se ela não existia antes
-                if 'CUDA_VISIBLE_DEVICES' in os.environ:
-                    del os.environ['CUDA_VISIBLE_DEVICES']
+            elif 'CUDA_VISIBLE_DEVICES' in os.environ:
+                del os.environ['CUDA_VISIBLE_DEVICES']
 
     def _load_model_task(self):
         """Tarefa de carregamento do modelo, a ser executada no ambiente isolado."""
         print(f"[VinceWorker-{self.device_id_str}] Carregando modelo para a VRAM (GPU física visível: {self.gpu_index})...")
-        
-        # Para o código VINCIE, a única GPU que ele vê é 'cuda:0'
         device_for_vincie = 'cuda:0'
-
         original_cwd = Path.cwd()
         try:
             os.chdir(str(VINCIE_DIR))
-            
             if str(VINCIE_DIR) not in sys.path: sys.path.insert(0, str(VINCIE_DIR))
 
             from common.config import load_config, create_object
@@ -81,12 +64,10 @@ class VinceWorker:
             
             self.gen.to(torch.device(device_for_vincie))
             print(f"[VinceWorker-{self.device_id_str}] ✅ Modelo VINCIE está 'quente' e pronto na GPU física {self.gpu_index}.")
-
         finally:
             os.chdir(original_cwd)
     
     def load_model_to_gpu(self):
-        """Wrapper público para carregar o modelo no ambiente isolado."""
         if self.gen is None:
             self._execute_in_isolated_env(self._load_model_task)
 
@@ -107,18 +88,17 @@ class VinceWorker:
             
             self.gen.inference_loop()
             return Path(kwargs["output_dir"])
-            
         finally:
             os.chdir(original_cwd)
             gc.collect()
             torch.cuda.empty_cache()
 
     def infer(self, **kwargs) -> Path:
-        """Wrapper público para executar a inferência no ambiente isolado."""
         if self.gen is None:
             raise RuntimeError(f"Modelo no worker {self.device_id_str} não foi carregado.")
         return self._execute_in_isolated_env(self._infer_task, **kwargs)
 
+
 # --- Classe Pool Manager (A Orquestradora Singleton) ---
 class VincePoolManager:
     _instance = None
@@ -132,54 +112,61 @@ class VincePoolManager:
                 cls._instance._initialized = False
             return cls._instance
 
-    def __init__(self, num_gpus: int = 4, config_path: str = "/data/VINCIE/configs/generate.yaml", output_root: str = "/app/outputs"):
+    def __init__(self, num_gpus: int = 4, output_root: str = "/app/outputs"):
         if self._initialized: return
         with self._lock:
             if self._initialized: return
 
-            print("Inicializando o VincePoolManager (Multi-GPU, Cache-Enabled, OOM-Fixed)...")
+            print("Inicializando o VincePoolManager (Padrão SeedVR)...")
             self.output_root = Path(output_root)
             self.output_root.mkdir(parents=True, exist_ok=True)
             self.worker_lock = threading.Lock()
             self.next_worker_idx = 0
 
-            self._setup_dependencies()
+            self.setup_dependencies()
 
             if not torch.cuda.is_available() or torch.cuda.device_count() < num_gpus:
                 raise RuntimeError(f"Erro: {num_gpus} GPUs são necessárias, mas {torch.cuda.device_count()} foram encontradas.")
             
             devices = [f'cuda:{i}' for i in range(num_gpus)]
-            self.workers = [VinceWorker(dev_id, config_path) for dev_id in devices]
+            vincie_config_path = VINCIE_DIR / "configs/generate.yaml"
+            self.workers = [VinceWorker(dev_id, str(vincie_config_path)) for dev_id in devices]
 
             print("Iniciando carregamento dos modelos em paralelo para todas as GPUs...")
             threads = [threading.Thread(target=worker.load_model_to_gpu) for worker in self.workers]
-            for t in threads:
-                t.start()
-            for t in threads:
-                t.join()
+            for t in threads: t.start()
+            for t in threads: t.join()
             
             self._initialized = True
             print(f"✅ VincePoolManager pronto com {len(self.workers)} workers 'quentes'.")
 
-    def _setup_dependencies(self):
+    def setup_dependencies(self):
         """Garante que o código e os modelos do VINCIE estejam disponíveis, usando cache."""
-        if not VINCIE_DIR.exists():
+        # 1. Código do Repositório (clona em /data/VINCIE)
+        if not (VINCIE_DIR / ".git").exists():
             print(f"Clonando repositório VINCIE para {VINCIE_DIR}...")
+            VINCIE_DIR.parent.mkdir(parents=True, exist_ok=True)
             subprocess.run(["git", "clone", "--depth", "1", VINCE_GIT_URL, str(VINCIE_DIR)], check=True)
         else:
-            print("Repositório VINCIE já existe.")
+            print("Repositório VINCIE já existe em /data/VINCIE.")
 
+        # 2. Modelos (com cache, baixando para /data/ckpt/VINCIE-3B)
         print(f"Verificando checkpoints VINCIE em {VINCE_CKPT} (usando cache em {HF_HOME_CACHE})...")
-        snapshot_download(
-            repo_id=VINCE_REPO_ID,
-            local_dir=VINCE_CKPT,
-            local_dir_use_symlinks=False,
-            cache_dir=HF_HOME_CACHE,
-            resume_download=True,
-            token=os.getenv("HF_TOKEN")
-        )
-        print("Checkpoints VINCIE prontos.")
-
+        try:
+            snapshot_download(
+                repo_id=VINCE_REPO_ID,
+                local_dir=VINCE_CKPT,
+                local_dir_use_symlinks=False,
+                cache_dir=HF_HOME_CACHE,
+                resume_download=True,
+                token=os.getenv("HF_TOKEN")
+            )
+            print("Checkpoints VINCIE prontos.")
+        except Exception as e:
+            print(f"ERRO durante o snapshot_download para VINCIE: {e}")
+            raise
+
+        # 3. Symlink para compatibilidade
         repo_ckpt_dir = VINCIE_DIR / "ckpt"
         repo_ckpt_dir.mkdir(parents=True, exist_ok=True)
         link = repo_ckpt_dir / "VINCIE-3B"
@@ -190,7 +177,6 @@ class VincePoolManager:
             print("Symlink de checkpoint já existe.")
 
     def _get_next_worker(self) -> VinceWorker:
-        """Obtém o próximo worker disponível em um sistema de rodízio (thread-safe)."""
         with self.worker_lock:
             worker = self.workers[self.next_worker_idx]
             self.next_worker_idx = (self.next_worker_idx + 1) % len(self.workers)
@@ -217,7 +203,9 @@ class VincePoolManager:
 # --- Instância Singleton Global ---
 try:
     NUM_GPUS_FOR_VINCE = int(os.getenv("VINCE_GPUS", "4"))
-    vince_pool_manager_singleton = VincePoolManager(num_gpus=NUM_GPUS_FOR_VINCE)
+    # Passamos `output_root` lido da env var para o construtor.
+    output_root_path = os.getenv("OUTPUT_ROOT", "/app/outputs")
+    vince_pool_manager_singleton = VincePoolManager(num_gpus=NUM_GPUS_FOR_VINCE, output_root=output_root_path)
 except Exception as e:
     print(f"ERRO CRÍTICO ao inicializar o VincePoolManager: {e}", file=sys.stderr)
     vince_pool_manager_singleton = None