Update api/seedvr_server.py

api/seedvr_server.py

@@ -9,46 +9,50 @@ import multiprocessing as mp
 from pathlib import Path
 from typing import Optional, Callable
 
-from huggingface_hub import hf_hub_download
+# --- 1. Import dos Módulos Compartilhados ---
+# É crucial que estes imports venham antes dos imports pesados (torch, etc.)
+# para que o ambiente de multiprocessing seja configurado corretamente.
 
-# -------------------------------------------------------------
-# 1. CONFIGURAÇÃO DE AMBIENTE E CUDA
-# -------------------------------------------------------------
+try:
+    # Importa o gerenciador de GPUs que centraliza a lógica de alocação
+    from api.gpu_manager import gpu_manager
+    # Importa o serviço do LTX para podermos comandá-lo a liberar a VRAM
+    from api.ltx_server_refactored import video_generation_service
+except ImportError:
+    print("ERRO FATAL: Não foi possível importar `gpu_manager` ou `video_generation_service`.")
+    print("Certifique-se de que os arquivos `gpu_manager.py` e `ltx_server_refactored.py` existem em `api/`.")
+    sys.exit(1)
 
-# Garante o uso seguro de CUDA com multiprocessing para estabilidade.
+
+# --- 2. Configuração de Ambiente e CUDA ---
 if mp.get_start_method(allow_none=True) != 'spawn':
     mp.set_start_method('spawn', force=True)
 
-# Configuração de alocação de memória da VRAM
 os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
 
-# Adiciona dinamicamente o caminho do repositório clonado ao sys.path.
+# Adiciona o caminho do repositório SeedVR
 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))
 
-# Importações pesadas (torch, etc.) são feitas após a configuração do ambiente.
+# Imports pesados
 import torch
 import cv2
 import numpy as np
 from datetime import datetime
 
-# -------------------------------------------------------------
-# 2. FUNÇÕES AUXILIARES DE PROCESSAMENTO (Workers e I/O)
-# -------------------------------------------------------------
+
+# --- 3. Funções Auxiliares de Processamento (Workers e I/O) ---
+# (Estas funções não precisam de alteração)
 
 def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load_cap=None):
-    """Extrai quadros de um vídeo e os converte para o formato de tensor."""
-    if debug: print(f"🎬 Extraindo frames de: {video_path}")
+    if debug: print(f"🎬 [SeedVR] Extraindo frames de: {video_path}")
     if not os.path.exists(video_path): raise FileNotFoundError(f"Arquivo de vídeo não encontrado: {video_path}")
-
     cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened(): raise ValueError(f"Não foi possível abrir o arquivo de vídeo: {video_path}")
-
+    if not cap.isOpened(): raise ValueError(f"Não foi possível abrir o vídeo: {video_path}")
+    
     fps = cap.get(cv2.CAP_PROP_FPS)
     frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
-    if debug: print(f"📊 Info do vídeo: {frame_count} frames, {fps:.2f} FPS")
-
     frames = []
     frames_loaded = 0
     for i in range(frame_count):
@@ -56,32 +60,26 @@ def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load
         if not ret: break
         if i < skip_first_frames: continue
         if load_cap and frames_loaded >= load_cap: break
-        
         frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         frames.append(frame.astype(np.float32) / 255.0)
         frames_loaded += 1
     cap.release()
-
-    if not frames: raise ValueError(f"Nenhum frame foi extraído do vídeo: {video_path}")
-    if debug: print(f"✅ {len(frames)} frames extraídos com sucesso.")
+    if not frames: raise ValueError(f"Nenhum frame extraído de: {video_path}")
+    if debug: print(f"✅ [SeedVR] {len(frames)} frames extraídos com sucesso.")
     return torch.from_numpy(np.stack(frames)).to(torch.float16), fps
 
 def save_frames_to_video(frames_tensor, output_path, fps=30.0, debug=False):
-    """Salva um tensor de quadros em um arquivo de vídeo."""
-    if debug: print(f"🎬 Salvando {frames_tensor.shape[0]} frames em: {output_path}")
+    if debug: print(f"💾 [SeedVR] Salvando {frames_tensor.shape[0]} frames em: {output_path}")
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    
     frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8)
     T, H, W, _ = frames_np.shape
-    
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
-    if not out.isOpened(): raise ValueError(f"Não foi possível criar o arquivo de vídeo: {output_path}")
-    
+    if not out.isOpened(): raise ValueError(f"Não foi possível criar o vídeo: {output_path}")
     for frame in frames_np:
         out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
     out.release()
-    if debug: print(f"✅ Vídeo salvo com sucesso: {output_path}")
+    if debug: print(f"✅ [SeedVR] Vídeo salvo com sucesso: {output_path}")
 
 def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, progress_queue=None):
     """Processo filho (worker) que executa o upscaling em uma GPU dedicada."""
@@ -94,7 +92,6 @@ def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, p
     
     try:
         frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
-        
         callback = (lambda b, t, _, m: progress_queue.put((proc_idx, b, t, m))) if progress_queue else None
 
         runner = configure_runner(shared_args["model"], shared_args["model_dir"], shared_args["preserve_vram"], shared_args["debug"])
@@ -107,14 +104,12 @@ def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, p
         return_queue.put((proc_idx, result_tensor.cpu().numpy()))
     except Exception as e:
         import traceback
-        error_msg = f"ERRO no worker {proc_idx}: {e}\n{traceback.format_exc()}"
+        error_msg = f"ERRO no worker {proc_idx} (GPU {device_id}): {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))
 
-# -------------------------------------------------------------
-# 3. CLASSE DO SERVIDOR PRINCIPAL
-# -------------------------------------------------------------
+# --- 4. CLASSE DO SERVIDOR PRINCIPAL ---
 
 class SeedVRServer:
     def __init__(self, **kwargs):
@@ -122,13 +117,16 @@ class SeedVRServer:
         print("⚙️ SeedVRServer inicializando...")
         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.OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/output"))
         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 = torch.cuda.device_count()
+        
+        # OBTÉM AS GPUS ALOCADAS PELO GERENCIADOR CENTRAL
+        self.device_list = gpu_manager.get_seedvr_devices()
+        self.num_gpus = len(self.device_list)
+        print(f"[SeedVR] Alocado para usar {self.num_gpus} GPU(s): {self.device_list}")
 
-        for p in [self.CKPTS_ROOT, self.OUTPUT_ROOT, self.INPUT_ROOT, self.HF_HOME_CACHE]:
+        for p in [self.CKPTS_ROOT, self.OUTPUT_ROOT, self.HF_HOME_CACHE]:
             p.mkdir(parents=True, exist_ok=True)
 
         self.setup_dependencies()
@@ -136,15 +134,10 @@ class SeedVRServer:
 
     def setup_dependencies(self):
         """Garante que o repositório e os modelos estão presentes."""
-        # Clona o repositório do SeedVR se não existir
         if not (self.SEEDVR_ROOT / ".git").exists():
-            print(f"[SeedVRServer] Clonando repositório para {self.SEEDVR_ROOT}...")
+            print(f"[SeedVR] Clonando repositório para {self.SEEDVR_ROOT}...")
             subprocess.run(["git", "clone", "--depth", "1", self.REPO_URL, str(self.SEEDVR_ROOT)], check=True)
-        else:
-            print("[SeedVRServer] Repositório SeedVR já existe.")
         
-        # Baixa os checkpoints do Hugging Face se não existirem
-        print(f"[SeedVRServer] Verificando checkpoints em {self.CKPTS_ROOT}...")
         model_files = {
             "seedvr2_ema_7b_sharp_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses",
             "ema_vae_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses"
@@ -157,7 +150,7 @@ class SeedVRServer:
                     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 estão no local correto.")
+        print("[SeedVR] Checkpoints verificados.")
 
     def run_inference(
         self,
@@ -172,106 +165,104 @@ class SeedVRServer:
         progress: Optional[Callable] = None
     ) -> str:
         """
-        Executa o pipeline completo de upscaling de vídeo e retorna o caminho do arquivo de saída.
+        Executa o pipeline completo de upscaling de vídeo, gerenciando a memória da GPU.
         """
-        if progress: progress(0.01, "⌛ Inicializando...")
-        
-        # --- 1. Extração de Frames ---
-        if progress: progress(0.05, "🎬 Extraindo frames do vídeo...")
-        frames_tensor, original_fps = extract_frames_from_video(file_path, debug)
+        if progress: progress(0.01, "⌛ Inicializando inferência SeedVR...")
 
-        # --- 2. Preparação do Processamento Multi-GPU ---
-        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 else None
-        
-        shared_args = {
-            "model": model, "model_dir": str(self.CKPTS_ROOT), "preserve_vram": preserve_vram,
-            "debug": debug, "seed": seed, "resolution": resolution, "batch_size": batch_size
-        }
+        # --- NÓ 1: GERENCIAMENTO DE MEMÓRIA (SWAP) ---
+        if gpu_manager.requires_memory_swap():
+            print("[SWAP] SeedVR precisa da GPU. Movendo LTX para a CPU...")
+            if progress: progress(0.02, "🔄 Liberando VRAM para o SeedVR...")
+            video_generation_service.move_to_cpu()
+            print("[SWAP] LTX movido para a CPU. VRAM liberada.")
 
-        # --- 3. Inicia os Workers ---
-        if progress: progress(0.1, f"🚀 Iniciando geração em {num_devices} GPUs...")
-        workers = []
-        for idx, device_id in enumerate(device_list):
-            p = mp.Process(target=_worker_process, args=(idx, device_id, chunks[idx].cpu().numpy(), shared_args, return_queue, progress_queue))
-            p.start()
-            workers.append(p)
-        
-        # --- 4. Coleta de Resultados e Monitoramento de Progresso ---
-        results_np = [None] * num_devices
-        finished_workers = 0
-        worker_progress = [0.0] * num_devices
-        while finished_workers < num_devices:
-            # Atualiza a barra de progresso com informações da fila
-            if progress_queue:
-                while not progress_queue.empty():
-                    try:
-                        p_idx, b_idx, b_total, msg = progress_queue.get_nowait()
-                        if b_idx == -1: raise RuntimeError(f"Erro no Worker {p_idx}: {msg}")
-                        if b_total > 0: worker_progress[p_idx] = b_idx / b_total
-                        total_progress = sum(worker_progress) / num_devices
-                        progress(0.1 + total_progress * 0.85, desc=f"GPU {p_idx+1}/{num_devices}: {msg}")
-                    except queue.Empty: pass
+        try:
+            # --- NÓ 2: EXTRAÇÃO DE FRAMES ---
+            if progress: progress(0.05, "🎬 Extraindo frames do vídeo...")
+            frames_tensor, original_fps = extract_frames_from_video(file_path, debug)
+
+            # --- NÓ 3: DIVISÃO PARA MULTI-GPU ---
+            if self.num_gpus == 0:
+                raise RuntimeError("SeedVR requer pelo menos 1 GPU alocada, mas não encontrou nenhuma.")
+            
+            print(f"[SeedVR] Dividindo {frames_tensor.shape[0]} frames em {self.num_gpus} chunks para processamento paralelo.")
+            chunks = torch.chunk(frames_tensor, self.num_gpus, dim=0)
+            
+            manager = mp.Manager()
+            return_queue = manager.Queue()
+            progress_queue = manager.Queue() if progress else None
             
-            # Verifica se algum worker terminou
-            try:
-                proc_idx, result = return_queue.get(timeout=0.2)
-                if isinstance(result, str): raise RuntimeError(f"Worker {proc_idx} falhou: {result}")
-                results_np[proc_idx] = result
-                worker_progress[proc_idx] = 1.0
-                finished_workers += 1
-            except queue.Empty: pass
+            shared_args = {
+                "model": model, "model_dir": str(self.CKPTS_ROOT), "preserve_vram": preserve_vram,
+                "debug": debug, "seed": seed, "resolution": resolution, "batch_size": batch_size
+            }
 
-        for p in workers: p.join()
+            # --- NÓ 4: INÍCIO DOS WORKERS ---
+            if progress: progress(0.1, f"🚀 Iniciando geração em {self.num_gpus} GPU(s)...")
+            workers = []
+            for idx, device_id in enumerate(self.device_list):
+                p = mp.Process(target=_worker_process, args=(idx, device_id, chunks[idx].cpu().numpy(), shared_args, return_queue, progress_queue))
+                p.start()
+                workers.append(p)
+            
+            # --- NÓ 5: COLETA DE RESULTADOS E MONITORAMENTO ---
+            results_np = [None] * self.num_gpus
+            finished_workers = 0
+            worker_progress = [0.0] * self.num_gpus
+            while finished_workers < self.num_gpus:
+                if progress_queue:
+                    while not progress_queue.empty():
+                        try:
+                            p_idx, b_idx, b_total, msg = progress_queue.get_nowait()
+                            if b_idx == -1: raise RuntimeError(f"Erro no Worker {p_idx}: {msg}")
+                            if b_total > 0: worker_progress[p_idx] = b_idx / b_total
+                            total_progress = sum(worker_progress) / self.num_gpus
+                            progress(0.1 + total_progress * 0.85, desc=f"GPU {p_idx+1}/{self.num_gpus}: {msg}")
+                        except queue.Empty: pass
+                
+                try:
+                    proc_idx, result = return_queue.get(timeout=0.2)
+                    if isinstance(result, str): raise RuntimeError(f"Worker {proc_idx} falhou: {result}")
+                    results_np[proc_idx] = result
+                    worker_progress[proc_idx] = 1.0
+                    finished_workers += 1
+                except queue.Empty: pass
 
-        if any(r is None for r in results_np):
-            raise RuntimeError("Um ou mais workers falharam ao retornar um resultado.")
+            for p in workers: p.join()
 
-        # --- 5. Combina os resultados e salva o vídeo final ---
-        result_tensor = torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)
+            # --- NÓ 6: FINALIZAÇÃO ---
+            if any(r is None for r in results_np):
+                raise RuntimeError("Um ou mais workers falharam ao retornar um resultado.")
 
-        if progress: progress(0.95, "💾 Salvando o vídeo final...")
-        
-        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"
+            result_tensor = torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)
+            if progress: progress(0.95, "💾 Salvando o vídeo final...")
+            
+            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"
 
-        final_fps = fps if fps and fps > 0 else original_fps
-        save_frames_to_video(result_tensor, str(output_filepath), final_fps, debug)
-        
-        print(f"✅ Vídeo salvo com sucesso em: {output_filepath}")
-        return str(output_filepath)
+            final_fps = fps if fps and fps > 0 else original_fps
+            save_frames_to_video(result_tensor, str(output_filepath), final_fps, debug)
+            
+            print(f"✅ Vídeo salvo com sucesso em: {output_filepath}")
+            return str(output_filepath)
 
-# -------------------------------------------------------------
-# 4. PONTO DE ENTRADA PARA EXECUÇÃO
-# -------------------------------------------------------------
+        finally:
+            # --- NÓ 7: RESTAURAÇÃO DE MEMÓRIA (SWAP BACK) ---
+            if gpu_manager.requires_memory_swap():
+                print("[SWAP] Inferência do SeedVR concluída. Movendo LTX de volta para a GPU...")
+                if progress: progress(0.99, "🔄 Restaurando o ambiente LTX...")
+                ltx_device = gpu_manager.get_ltx_device()
+                video_generation_service.move_to_device(ltx_device)
+                print(f"[SWAP] LTX de volta em {ltx_device}.")
 
+# --- PONTO DE ENTRADA ---
 if __name__ == "__main__":
-    # Bloco para testes ou inicialização autônoma.
-    print("🚀 Executando o servidor SeedVR em modo autônomo...")
+    print("🚀 Executando o servidor SeedVR em modo autônomo para inicialização...")
     try:
         server = SeedVRServer()
         print("✅ Servidor inicializado com sucesso. Pronto para receber chamadas.")
-        # Exemplo de como chamar a inferência (requer um arquivo de vídeo):
-        # input_video = "caminho/para/seu/video.mp4"
-        # if os.path.exists(input_video):
-        #     server.run_inference(
-        #         file_path=input_video,
-        #         seed=42,
-        #         resolution=1072,
-        #         batch_size=4,
-        #         progress=lambda p, desc: print(f"Progresso: {p*100:.1f}% - {desc}")
-        #     )
-        # else:
-        #     print(f"Vídeo de teste não encontrado em '{input_video}'. Pulei a execução da inferência.")
     except Exception as e:
-        print(f"❌ Falha ao inicializar o servidor: {e}")
-        import traceback
+        print(f"❌ Falha ao inicializar o servidor SeedVR: {e}")
         traceback.print_exc()
-        sys.exit(1)
-        
+        sys.exit(1)