Update aduc_framework/managers/seedvr_manager.py

aduc_framework/managers/seedvr_manager.py

@@ -1,4 +1,4 @@
-# hd_specialist.py (Versão Final - Estrutura de Arquivos Corrigida e Autossuficiente)
+# hd_specialist.py (Versão Final - Gerenciando o Repositório Completo)
 # https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B
 
 import torch
@@ -13,49 +13,53 @@ import shlex
 import subprocess
 from pathlib import Path
 from urllib.parse import urlparse
-from torch.hub import download_url_to_file, get_dir
+from torch.hub import download_url_to_file
 from omegaconf import OmegaConf
 import sys
 
-
+# --- Configuração do Logging ---
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-APP_ROOT = Path("/home/user/app")
-DEPS_DIR = APP_ROOT / "deps"
+# --- 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
+DEPS_DIR = PROJECT_ROOT / "deps"
 SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
 SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"
 
-
-def setup_dependencies():
+def setup_environment():
     """
-    Ensures the SEEDVR repository is cloned and available in the sys.path.
-    This function is run once when the module is first imported.
+    Clona o repositório SeedVR se não existir e o adiciona ao sys.path
+    para que seus módulos (common, projects, etc.) possam ser importados.
     """
-    if not SEEDVR_SPACE_DIR.exists():
-        logger.info(f"SEEDVR repository not found at '{SEEDVR_SPACE_DIR}'. Cloning from GitHub...")
+    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
             )
-            logger.info("SEEDVR repository cloned successfully.")
+            logger.info(f"✅ Repositório clonado com sucesso em '{SEEDVR_SPACE_DIR}'")
         except subprocess.CalledProcessError as e:
-            logger.error(f"Failed to clone SEEDVR repository. Git stderr: {e.stderr}")
-            raise RuntimeError("Could not clone the required SEEDVR dependency from GitHub.")
+            logger.error(f"❌ Falha ao clonar o repositório. Erro do Git: {e.stderr}")
+            raise RuntimeError("Não foi possível clonar a dependência SeedVR do Hugging Face.")
     else:
-        logger.info("Found local SEEDVR repository.")
-    
-    if str(SEEDVR_SPACE_DIR.resolve()) not in sys.path:
-        sys.path.insert(0, str(SEEDVR_SPACE_DIR.resolve()))
-        logger.info(f"Added '{SEEDVR_SPACE_DIR.resolve()}' to sys.path.")
-
-setup_dependencies()
+        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)
+        logger.info(f"Adicionado '{resolved_path}' ao sys.path.")
 
+# Executa a configuração do ambiente assim que o módulo é carregado
+setup_environment()
 
-# Função auxiliar para download
+# Função auxiliar de download (permanece a mesma)
 def _load_file_from_url(url, model_dir='./', file_name=None):
     os.makedirs(model_dir, exist_ok=True)
     filename = file_name or os.path.basename(urlparse(url).path)
@@ -65,9 +69,8 @@ def _load_file_from_url(url, model_dir='./', file_name=None):
         download_url_to_file(url, cached_file, hash_prefix=None, progress=True)
     return cached_file
 
-
-
-# --- Importações diretas, assumindo que as pastas estão na raiz ---
+# --- 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
@@ -79,8 +82,6 @@ from torchvision.transforms import Compose, Lambda, Normalize
 from torchvision.io.video import read_video
 from einops import rearrange
 
-    
-
 class SeedVrManager:
     """
     Implementa o Especialista HD (Δ+) usando a infraestrutura oficial do SeedVR.
@@ -88,27 +89,25 @@ class SeedVrManager:
     def __init__(self, workspace_dir="deformes_workspace"):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self.runner = None
-        self.workspace_dir = workspace_dir
+        self.workspace_dir = Path(workspace_dir)
         self.is_initialized = False
         logger.info("Especialista HD (SeedVR) inicializado. Modelo será carregado sob demanda.")
 
     def _setup_dependencies(self):
-    
         """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'
-        apex_wheel_path = _load_file_from_url(url=apex_url)
-        
-        # Instala a roda do Apex baixada
+        # 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.")
 
     def _download_models(self):
-        """Baixa os checkpoints necessários para o SeedVR2."""
+        """Baixa os checkpoints necessários para o SeedVR2 DENTRO do repositório clonado."""
         logger.info("Verificando e baixando modelos do SeedVR2...")
-        ckpt_dir = Path('./ckpts')
-        ckpt_dir.mkdir(exist_ok=True)
-
+        # O diretório de checkpoints agora é DENTRO do repositório clonado
+        ckpt_dir = SEEDVR_SPACE_DIR / 'ckpts'
+        
         pretrain_model_url = {
             'vae': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/ema_vae.pth',
             'dit': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/seedvr2_ema_3b.pth',
@@ -116,11 +115,12 @@ class SeedVrManager:
             'neg_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/neg_emb.pt'
         }
         
-        _load_file_from_url(url=pretrain_model_url['dit'], model_dir='./ckpts/')
-        _load_file_from_url(url=pretrain_model_url['vae'], model_dir='./ckpts/')
-        _load_file_from_url(url=pretrain_model_url['pos_emb'])
-        _load_file_from_url(url=pretrain_model_url['neg_emb'])
-        logger.info("Modelos do SeedVR2 baixados com sucesso.")
+        # 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))
+        _load_file_from_url(url=pretrain_model_url['neg_emb'], model_dir=str(SEEDVR_SPACE_DIR))
+        logger.info("✅ Modelos do SeedVR2 baixados com sucesso.")
 
     def _initialize_runner(self):
         """Carrega e configura o modelo SeedVR sob demanda."""
@@ -132,66 +132,49 @@ class SeedVrManager:
 
         logger.info("Inicializando o runner do SeedVR2...")
         
-        # --- NOVO: Verificação e download automático do main.yaml ---
-        config_dir = Path('./configs_3b')
-        config_path = config_dir / 'main.yaml'
-        config_url = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B/resolve/main/configs_3b/main.yaml"
+        # 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():
-            logger.warning(f"Arquivo de configuração '{config_path}' não encontrado.")
-            logger.info("Tentando baixar automaticamente...")
-            try:
-                # Reutiliza a função de download existente
-                _load_file_from_url(url=config_url, model_dir=str(config_dir), file_name='main.yaml')
-                logger.info(f"✅ Arquivo de configuração baixado com sucesso para '{config_path}'")
-            except Exception as e:
-                logger.error(f"Falha ao baixar o arquivo de configuração. Por favor, baixe-o manualmente de {config_url} e coloque-o em {config_dir}. Erro: {e}")
-                raise
-        else:
-            logger.info(f"Arquivo de configuração encontrado em '{config_path}'.")
-        # --- FIM DA MODIFICAÇÃO ---
+            # 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.")
 
+        # `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)
         
-        self.runner.configure_dit_model(device=self.device, checkpoint='./ckpts/seedvr2_ema_3b.pth')
+        # 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.")
+        logger.info("✅ Runner do SeedVR2 inicializado e pronto.")
 
     def _unload_runner(self):
         """Remove o runner da VRAM para liberar recursos."""
         if self.runner is not None:
-            del self.runner
-            self.runner = None
-            gc.collect()
-            torch.cuda.empty_cache()
+            del self.runner; self.runner = None
+            gc.collect(); torch.cuda.empty_cache()
             self.is_initialized = False
             logger.info("Runner do SeedVR2 descarregado da VRAM.")
 
-    
-
     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.
-        """
+        """Aplica o aprimoramento HD a um vídeo usando a lógica oficial do SeedVR."""
         try:
             self._initialize_runner()
             set_seed(seed, same_across_ranks=True)
 
-            # --- Configuração do Pipeline (adaptado de app.py) ---
-            self.runner.config.diffusion.cfg.scale = 1.0 # cfg_scale
-            self.runner.config.diffusion.cfg.rescale = 0.0 # cfg_rescale
-            self.runner.config.diffusion.timesteps.sampling.steps = 1 # sample_steps (one-step model)
+            self.runner.config.diffusion.cfg.scale = 1.0
+            self.runner.config.diffusion.cfg.rescale = 0.0
+            self.runner.config.diffusion.timesteps.sampling.steps = 1
             self.runner.configure_diffusion()
 
-            # --- Preparação do Vídeo de Entrada ---
             logger.info(f"Processando vídeo de entrada: {input_video_path}")
             video_tensor = read_video(input_video_path, output_format="TCHW")[0] / 255.0
             if video_tensor.size(0) > 121:
@@ -200,22 +183,22 @@ class SeedVrManager:
             
             video_transform = Compose([
                 NaResize(resolution=(1280 * 720)**0.5, mode="area", downsample_only=False),
-                Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
-                DivisibleCrop((16, 16)),
-                Normalize(0.5, 0.5),
-                Rearrange("t c h w -> c t h w"),
+                Lambda(lambda x: torch.clamp(x, 0.0, 1.0)), DivisibleCrop((16, 16)),
+                Normalize(0.5, 0.5), Rearrange("t c h w -> c t h w"),
             ])
             
             cond_latent = video_transform(video_tensor.to(self.device))
-            input_video_for_colorfix = cond_latent.clone() # Salva para o color fix
+            input_video_for_colorfix = cond_latent.clone()
             ori_length = cond_latent.size(1)
 
-            # --- Codificação VAE e Geração ---
             logger.info("Codificando vídeo para o espaço latente...")
             cond_latent = self.runner.vae_encode([cond_latent])[0]
 
-            text_pos_embeds = torch.load('pos_emb.pt').to(self.device)
-            text_neg_embeds = torch.load('neg_emb.pt').to(self.device)
+            # 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)
             text_embeds_dict = {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}
 
             noise = torch.randn_like(cond_latent)
@@ -225,28 +208,25 @@ class SeedVrManager:
                 video_tensor_out = self.runner.inference(
                     noises=[noise],
                     conditions=[self.runner.get_condition(noise, task="sr", latent_blur=cond_latent)],
-                    dit_offload=False,
-                    **text_embeds_dict,
+                    dit_offload=False, **text_embeds_dict,
                 )[0]
 
             sample = rearrange(video_tensor_out, "c t h w -> t c h w")
-            
-            # --- Pós-processamento e Salvamento ---
             if ori_length < sample.shape[0]:
                 sample = sample[:ori_length]
             
             input_video_for_colorfix = rearrange(input_video_for_colorfix, "c t h w -> t c h w")
             sample = wavelet_reconstruction(sample.cpu(), input_video_for_colorfix[:sample.size(0)].cpu())
-            
             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}")
+            self.workspace_dir.mkdir(parents=True, exist_ok=True)
             imageio.get_writer(output_video_path, fps=fps_out, codec='libx264', quality=9).extend(sample)
 
             return output_video_path
         finally:
-            self._unload_runner()        
+            self._unload_runner()
 
-# Instância Singleton para facilitar o uso (opcional)
+# Instância Singleton
 seedvr_manager_singleton = SeedVrManager()