Update aduc_framework/managers/seedvr_manager.py

aduc_framework/managers/seedvr_manager.py

@@ -2,13 +2,10 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 6.0.0 (GPU Isolation Fix)
+# Version: 6.1.0 (Force 3B Model)
 #
-# Esta versão implementa a solução definitiva para os erros de dispositivo
-# usando a variável de ambiente CUDA_VISIBLE_DEVICES. Cada worker agora opera
-# em um ambiente completamente isolado, vendo apenas sua própria GPU.
-# Isso força todo o código de terceiros a usar o dispositivo correto e
-# elimina a necessidade de gerenciar manualmente o torch.distributed.
+# Esta versão remove a seleção de modelo e força o uso exclusivo do SeedVR 3B.
+# Também otimiza o setup para não baixar o checkpoint do modelo 7B.
 
 import torch
 import os
@@ -26,12 +23,10 @@ import shutil
 from omegaconf import OmegaConf
 import yaml
 
-# Imports relativos para o hardware_manager
 from ..tools.hardware_manager import hardware_manager
 
 logger = logging.getLogger(__name__)
 
-# --- Caminhos Globais ---
 APP_ROOT = Path("/home/user/app")
 DEPS_DIR = APP_ROOT / "deps"
 SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
@@ -43,7 +38,6 @@ class SeedVrWorker:
         self.global_device_id = device_id
         self.local_device_name = 'cuda:0'
         self.gpu_index = self.global_device_id.split(':')[-1]
-        
         self.runner = None
         self.is_initialized = False
         self.setup_complete = self._check_and_run_global_setup()
@@ -55,22 +49,17 @@ class SeedVrWorker:
         setup_flag = DEPS_DIR / "seedvr.setup.complete"
         if str(APP_ROOT) not in sys.path:
             sys.path.insert(0, str(APP_ROOT))
-            
-        if setup_flag.exists():
-            return True
+        if setup_flag.exists(): return True
         
         logger.info("--- Iniciando Setup Global do SeedVR (primeira execução) ---")
-        
         if not SEEDVR_SPACE_DIR.exists():
             DEPS_DIR.mkdir(exist_ok=True, parents=True)
             subprocess.run(["git", "clone", "--depth", "1", SEEDVR_SPACE_URL, str(SEEDVR_SPACE_DIR)], check=True)
         
         required_dirs = ["projects", "common", "models", "configs_3b", "configs_7b", "data"]
         for dirname in required_dirs:
-            source = SEEDVR_SPACE_DIR / dirname
-            target = APP_ROOT / dirname
-            if not target.exists():
-                shutil.copytree(source, target)
+            source, target = SEEDVR_SPACE_DIR / dirname, APP_ROOT / dirname
+            if not target.exists(): shutil.copytree(source, target)
         
         try:
             import apex
@@ -81,10 +70,11 @@ class SeedVrWorker:
 
         ckpt_dir = APP_ROOT / 'ckpts'
         ckpt_dir.mkdir(exist_ok=True)
+        
+        # <<< MODIFICAÇÃO: Removido o download do modelo 7B >>>
         model_urls = {
             'vae': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/ema_vae.pth',
             'dit_3b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/seedvr2_ema_3b.pth',
-            #'dit_7b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-7B/resolve/main/seedvr2_ema_7b.pth',
             'pos_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/pos_emb.pt',
             'neg_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/neg_emb.pt'
         }
@@ -95,8 +85,8 @@ class SeedVrWorker:
         logger.info("--- Setup Global do SeedVR Concluído ---")
         return True
 
-    def initialize_runner(self, model_version: str):
-        """Carrega os modelos para a VRAM do dispositivo, usando um ambiente de GPU isolado."""
+    def initialize_runner(self):
+        """Carrega o modelo 3B para a VRAM do dispositivo."""
         if self.runner is not None: return
         
         os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_index
@@ -104,13 +94,11 @@ class SeedVrWorker:
         from projects.video_diffusion_sr.infer import VideoDiffusionInfer
         from common.config import load_config
         
-        logger.info(f"Worker {self.global_device_id}: Inicializando runner... (Processo vê apenas {self.local_device_name})")
+        logger.info(f"Worker {self.global_device_id}: Inicializando runner SeedVR 3B... (Processo vê apenas {self.local_device_name})")
         
-        config_path_str = f'configs_{model_version.lower()}'
-        checkpoint_path_str = f'seedvr2_ema_{model_version.lower()}.pth'
-
-        config_path = APP_ROOT / config_path_str / 'main.yaml'
-        checkpoint_path = APP_ROOT / 'ckpts' / checkpoint_path_str
+        # <<< MODIFICAÇÃO: Caminhos fixados para o modelo 3B >>>
+        config_path = APP_ROOT / 'configs_3b' / 'main.yaml'
+        checkpoint_path = APP_ROOT / 'ckpts' / 'seedvr2_ema_3b.pth'
 
         config = load_config(str(config_path))
         self.runner = VideoDiffusionInfer(config)
@@ -120,7 +108,7 @@ class SeedVrWorker:
         self.runner.configure_vae_model()
         
         self.is_initialized = True
-        logger.info(f"Worker {self.global_device_id}: Runner pronto na VRAM.")
+        logger.info(f"Worker {self.global_device_id}: Runner 3B pronto na VRAM.")
         
     def unload_runner(self):
         """Descarrega os modelos da VRAM e limpa o ambiente."""
@@ -135,7 +123,7 @@ class SeedVrWorker:
         if 'CUDA_VISIBLE_DEVICES' in os.environ:
             del os.environ['CUDA_VISIBLE_DEVICES']
 
-    def process_video_internal(self, input_video_path, output_video_path, prompt, model_version, steps, seed):
+    def process_video_internal(self, input_video_path, output_video_path, prompt, steps, seed):
         """Executa a inferência em um ambiente de GPU isolado."""
         os.environ['CUDA_VISIBLE_DEVICES'] = self.gpu_index
         device = torch.device(self.local_device_name)
@@ -182,7 +170,7 @@ class SeedVrWorker:
         self.runner.vae.to(device)
         samples = self.runner.vae_decode(video_tensors)
         final_sample = samples[0]
-        input_video_sample = cond_latents[0] # Usar o latente de condição como base
+        input_video_sample = cond_latents[0]
         if final_sample.shape[1] < input_video_sample.shape[1]:
             input_video_sample = input_video_sample[:, :final_sample.shape[1]]
             
@@ -214,8 +202,9 @@ class SeedVrPoolManager:
         logger.info(f"SEEDVR CLEANUP THREAD: Iniciando limpeza de {worker.global_device_id} em background...")
         worker.unload_runner()
 
+    # <<< MODIFICAÇÃO: Removido o argumento 'model_version' da assinatura pública >>>
     def process_video(self, input_video_path: str, output_video_path: str, prompt: str,
-                      model_version: str = '3B', steps: int = 100, seed: int = 666) -> str:
+                      steps: int = 100, seed: int = 666) -> str:
         worker_to_use = None
         try:
             with self.lock:
@@ -230,13 +219,15 @@ class SeedVrPoolManager:
                 cleanup_thread.start()
                 self.last_cleanup_thread = cleanup_thread
                 
-                worker_to_use.initialize_runner(model_version)
+                # Chama initialize_runner sem argumentos, pois ele agora sabe que deve usar o 3B
+                worker_to_use.initialize_runner()
                 
                 self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
             
             logger.info(f"SEEDVR POOL MANAGER: Processando vídeo na GPU {worker_to_use.global_device_id}...")
+            # Passa os argumentos para a função interna, sem 'model_version'
             return worker_to_use.process_video_internal(
-                input_video_path, output_video_path, prompt, model_version, steps, seed
+                input_video_path, output_video_path, prompt, steps, seed
             )
         except Exception as e:
             logger.error(f"SEEDVR POOL MANAGER: Erro durante o processamento de vídeo: {e}", exc_info=True)