Update aduc_framework/managers/seedvr_manager.py

aduc_framework/managers/seedvr_manager.py

@@ -1,4 +1,4 @@
-# hd_specialist.py (Versão Final - Gerenciando o Repositório Completo)
+# hd_specialist.py (Versão Final - Corrigindo o Contexto de Execução de Caminhos)
 # https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B
 
 import torch
@@ -7,8 +7,6 @@ import os
 import gc
 import logging
 import numpy as np
-from PIL import Image
-from tqdm import tqdm
 import shlex
 import subprocess
 from pathlib import Path
@@ -23,7 +21,8 @@ logger = logging.getLogger(__name__)
 
 # --- Constantes de Caminho ---
 # Define a raiz do projeto (onde este script está) e cria um diretório para dependências
-PROJECT_ROOT = Path(__file__).resolve().parent
+# Usamos 'Path.cwd()' para ter certeza que é relativo ao diretório de execução do projeto principal.
+PROJECT_ROOT = Path.cwd() 
 DEPS_DIR = PROJECT_ROOT / "deps"
 SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
 SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"
@@ -36,9 +35,7 @@ def setup_environment():
     if not SEEDVR_SPACE_DIR.is_dir():
         logger.info(f"Repositório SeedVR não encontrado. Clonando de '{SEEDVR_SPACE_URL}'...")
         try:
-            # Garante que o diretório de dependências exista
             DEPS_DIR.mkdir(exist_ok=True)
-            # Clona o repositório. '--depth 1' baixa apenas a versão mais recente.
             subprocess.run(
                 ["git", "clone", "--depth", "1", SEEDVR_SPACE_URL, str(SEEDVR_SPACE_DIR)],
                 check=True, capture_output=True, text=True
@@ -50,7 +47,6 @@ def setup_environment():
     else:
         logger.info(f"Repositório SeedVR já existe em '{SEEDVR_SPACE_DIR}'.")
 
-    # Adiciona a raiz do repositório clonado ao path do Python
     resolved_path = str(SEEDVR_SPACE_DIR.resolve())
     if resolved_path not in sys.path:
         sys.path.insert(0, resolved_path)
@@ -70,7 +66,6 @@ def _load_file_from_url(url, model_dir='./', file_name=None):
     return cached_file
 
 # --- Importações do Repositório Clonado ---
-# Agora estas importações funcionam porque `setup_environment` adicionou o repositório ao sys.path
 from projects.video_diffusion_sr.infer import VideoDiffusionInfer
 from common.config import load_config
 from common.seed import set_seed
@@ -97,7 +92,6 @@ class SeedVrManager:
         """Instala dependências complexas como Apex."""
         logger.info("Configurando dependências do SeedVR (Apex)...")
         apex_url = 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x86_64.whl'
-        # Baixa para um diretório temporário ou de dependências
         apex_wheel_path = _load_file_from_url(url=apex_url, model_dir=str(DEPS_DIR))
         subprocess.run(shlex.split(f"pip install {apex_wheel_path}"), check=True)
         logger.info("✅ Dependência Apex instalada com sucesso.")
@@ -105,7 +99,6 @@ class SeedVrManager:
     def _download_models(self):
         """Baixa os checkpoints necessários para o SeedVR2 DENTRO do repositório clonado."""
         logger.info("Verificando e baixando modelos do SeedVR2...")
-        # O diretório de checkpoints agora é DENTRO do repositório clonado
         ckpt_dir = SEEDVR_SPACE_DIR / 'ckpts'
         
         pretrain_model_url = {
@@ -115,7 +108,6 @@ class SeedVrManager:
             'neg_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/neg_emb.pt'
         }
         
-        # Salva os modelos nos locais corretos que o código espera
         _load_file_from_url(url=pretrain_model_url['dit'], model_dir=str(ckpt_dir))
         _load_file_from_url(url=pretrain_model_url['vae'], model_dir=str(ckpt_dir))
         _load_file_from_url(url=pretrain_model_url['pos_emb'], model_dir=str(SEEDVR_SPACE_DIR))
@@ -132,29 +124,39 @@ class SeedVrManager:
 
         logger.info("Inicializando o runner do SeedVR2...")
         
-        # CORREÇÃO: O caminho para o config agora é relativo ao repositório clonado.
-        config_path = SEEDVR_SPACE_DIR / 'configs_3b' / 'main.yaml'
-        
-        if not config_path.is_file():
-            # Esta verificação agora é uma salvaguarda, pois o git clone já deve ter baixado o arquivo.
-            raise FileNotFoundError(f"Arquivo de configuração principal não encontrado em {config_path}. O clone do repositório pode ter falhado.")
+        # --- CORREÇÃO CRÍTICA: MUDANÇA DE DIRETÓRIO DE TRABALHO ---
+        original_cwd = Path.cwd()
+        try:
+            # Muda para o diretório do repositório clonado. Isso é essencial para que
+            # o `load_config` encontre os arquivos .yaml herdados (como os da pasta `models`).
+            os.chdir(SEEDVR_SPACE_DIR)
+            
+            logger.info(f"Diretório de trabalho alterado para: {SEEDVR_SPACE_DIR}")
 
-        # `load_config` agora encontrará os arquivos herdados pois estamos no contexto do repo clonado.
-        config = load_config(str(config_path))
-        
-        self.runner = VideoDiffusionInfer(config)
-        OmegaConf.set_readonly(self.runner.config, False)
-        
-        # CORREÇÃO: Os caminhos para os checkpoints também são relativos ao repositório clonado.
-        dit_checkpoint_path = SEEDVR_SPACE_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
-        self.runner.configure_dit_model(device=self.device, checkpoint=str(dit_checkpoint_path))
-        self.runner.configure_vae_model()
-        
-        if hasattr(self.runner.vae, "set_memory_limit"):
-            self.runner.vae.set_memory_limit(**self.runner.config.vae.memory_limit)
-        
-        self.is_initialized = True
-        logger.info("✅ Runner do SeedVR2 inicializado e pronto.")
+            # Agora todos os caminhos são relativos à raiz do repositório
+            config_path = './configs_3b/main.yaml'
+            dit_checkpoint_path = './ckpts/seedvr2_ema_3b.pth'
+            
+            config = load_config(config_path)
+            
+            self.runner = VideoDiffusionInfer(config)
+            OmegaConf.set_readonly(self.runner.config, False)
+            
+            self.runner.configure_dit_model(device=self.device, checkpoint=dit_checkpoint_path)
+            self.runner.configure_vae_model()
+            
+            if hasattr(self.runner.vae, "set_memory_limit"):
+                self.runner.vae.set_memory_limit(**self.runner.config.vae.memory_limit)
+            
+            self.is_initialized = True
+            logger.info("✅ Runner do SeedVR2 inicializado e pronto.")
+
+        finally:
+            # Garante que o diretório de trabalho original seja restaurado,
+            # não importa se a inicialização foi bem-sucedida ou falhou.
+            os.chdir(original_cwd)
+            logger.info(f"Diretório de trabalho restaurado para: {original_cwd}")
+        # --- FIM DA CORREÇÃO ---
 
     def _unload_runner(self):
         """Remove o runner da VRAM para liberar recursos."""
@@ -166,8 +168,14 @@ class SeedVrManager:
 
     def process_video(self, input_video_path: str, output_video_path: str, prompt: str) -> str:
         """Aplica o aprimoramento HD a um vídeo usando a lógica oficial do SeedVR."""
+        original_cwd = Path.cwd() # Salva o diretório original
         try:
             self._initialize_runner()
+            
+            # --- CORREÇÃO CRÍTICA 2: MUDAR DIRETÓRIO TAMBÉM DURANTE O PROCESSAMENTO ---
+            # O código pode precisar acessar arquivos de embedding novamente
+            os.chdir(SEEDVR_SPACE_DIR)
+
             set_seed(seed, same_across_ranks=True)
 
             self.runner.config.diffusion.cfg.scale = 1.0
@@ -176,7 +184,11 @@ class SeedVrManager:
             self.runner.configure_diffusion()
 
             logger.info(f"Processando vídeo de entrada: {input_video_path}")
-            video_tensor = read_video(input_video_path, output_format="TCHW")[0] / 255.0
+            # Garante que os caminhos de entrada/saída sejam absolutos para não quebrar com a mudança de CWD
+            abs_input_path = original_cwd / input_video_path
+            abs_output_path = original_cwd / output_video_path
+
+            video_tensor = read_video(str(abs_input_path), output_format="TCHW")[0] / 255.0
             if video_tensor.size(0) > 121:
                 logger.warning(f"Vídeo com {video_tensor.size(0)} frames. Truncando para 121 frames.")
                 video_tensor = video_tensor[:121]
@@ -194,11 +206,9 @@ class SeedVrManager:
             logger.info("Codificando vídeo para o espaço latente...")
             cond_latent = self.runner.vae_encode([cond_latent])[0]
 
-            # CORREÇÃO: Carrega os embeddings a partir do repositório clonado
-            pos_emb_path = SEEDVR_SPACE_DIR / 'pos_emb.pt'
-            neg_emb_path = SEEDVR_SPACE_DIR / 'neg_emb.pt'
-            text_pos_embeds = torch.load(pos_emb_path).to(self.device)
-            text_neg_embeds = torch.load(neg_emb_path).to(self.device)
+            # Carrega embeddings com caminhos relativos ao diretório do repo
+            text_pos_embeds = torch.load('pos_emb.pt').to(self.device)
+            text_neg_embeds = torch.load('neg_emb.pt').to(self.device)
             text_embeds_dict = {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}
 
             noise = torch.randn_like(cond_latent)
@@ -220,12 +230,13 @@ class SeedVrManager:
             sample = rearrange(sample, "t c h w -> t h w c")
             sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round().to(torch.uint8).numpy()
 
-            logger.info(f"Salvando vídeo aprimorado em: {output_video_path}")
+            logger.info(f"Salvando vídeo aprimorado em: {abs_output_path}")
             self.workspace_dir.mkdir(parents=True, exist_ok=True)
-            imageio.get_writer(output_video_path, fps=fps_out, codec='libx264', quality=9).extend(sample)
+            imageio.get_writer(str(abs_output_path), fps=fps_out, codec='libx264', quality=9).extend(sample)
 
-            return output_video_path
+            return str(abs_output_path)
         finally:
+            os.chdir(original_cwd) # Restaura o diretório de trabalho original
             self._unload_runner()
 
 # Instância Singleton