Update managers/seedvr_manager.py

managers/seedvr_manager.py

@@ -2,11 +2,11 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 2.3.1
+# Version: 2.3.2
 #
-# Esta versão adiciona uma correção robusta para o FileNotFoundError da configuração do VAE,
-# sobrescrevendo o caminho fixo incorreto em tempo de execução para apontar para o local
-# correto dentro do repositório clonado.
+# Esta versão implementa uma correção robusta para o FileNotFoundError da configuração do VAE,
+# antecipando a falha, carregando as configurações manualmente e fundindo-as para
+# contornar o caminho fixo problemático na biblioteca externa.
 
 import torch
 import os
@@ -21,7 +21,6 @@ import gradio as gr
 import mediapy
 from einops import rearrange
 
-# Utilitário internalizado para correção de cor, garantindo estabilidade.
 from tools.tensor_utils import wavelet_reconstruction
 
 logger = logging.getLogger(__name__)
@@ -30,22 +29,15 @@ logger = logging.getLogger(__name__)
 DEPS_DIR = Path("./deps")
 SEEDVR_REPO_DIR = DEPS_DIR / "SeedVR"
 SEEDVR_REPO_URL = "https://github.com/ByteDance-Seed/SeedVR.git"
-# URL direto para o arquivo de configuração do VAE que pode estar faltando
 VAE_CONFIG_URL = "https://raw.githubusercontent.com/ByteDance-Seed/SeedVR/main/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml"
 
 def setup_seedvr_dependencies():
-    """
-    Garante que o repositório do SeedVR seja clonado e esteja disponível no sys.path.
-    Esta função é executada uma vez quando o módulo é importado pela primeira vez.
-    """
+    """Garante que o repositório do SeedVR seja clonado e esteja disponível no sys.path."""
     if not SEEDVR_REPO_DIR.exists():
         logger.info(f"Repositório SeedVR não encontrado em '{SEEDVR_REPO_DIR}'. Clonando do GitHub...")
         try:
             DEPS_DIR.mkdir(exist_ok=True)
-            subprocess.run(
-                ["git", "clone", "--depth", "1", SEEDVR_REPO_URL, str(SEEDVR_REPO_DIR)],
-                check=True, capture_output=True, text=True
-            )
+            subprocess.run(["git", "clone", "--depth", "1", SEEDVR_REPO_URL, str(SEEDVR_REPO_DIR)], check=True, capture_output=True, text=True)
             logger.info("Repositório SeedVR clonado com sucesso.")
         except subprocess.CalledProcessError as e:
             logger.error(f"Falha ao clonar o repositório SeedVR. Git stderr: {e.stderr}")
@@ -57,10 +49,8 @@ def setup_seedvr_dependencies():
         sys.path.insert(0, str(SEEDVR_REPO_DIR.resolve()))
         logger.info(f"Adicionado '{SEEDVR_REPO_DIR.resolve()}' ao sys.path.")
 
-# --- Executa a configuração da dependência imediatamente na importação do módulo ---
 setup_seedvr_dependencies()
 
-# --- Agora que o caminho está configurado, podemos importar com segurança do repositório clonado ---
 from projects.video_diffusion_sr.infer import VideoDiffusionInfer
 from common.config import load_config
 from common.seed import set_seed
@@ -71,9 +61,7 @@ from torchvision.transforms import Compose, Lambda, Normalize
 from torchvision.io.video import read_video
 from omegaconf import OmegaConf
 
-
 def _load_file_from_url(url, model_dir='./', file_name=None):
-    """Função auxiliar para baixar arquivos de uma URL para um diretório local."""
     os.makedirs(model_dir, exist_ok=True)
     filename = file_name or os.path.basename(urlparse(url).path)
     cached_file = os.path.abspath(os.path.join(model_dir, filename))
@@ -83,9 +71,7 @@ def _load_file_from_url(url, model_dir='./', file_name=None):
     return cached_file
 
 class SeedVrManager:
-    """
-    Gerencia o modelo SeedVR para tarefas de Masterização HD.
-    """
+    """Gerencia o modelo SeedVR para tarefas de Masterização HD."""
     def __init__(self, workspace_dir="deformes_workspace"):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self.runner = None
@@ -94,18 +80,13 @@ class SeedVrManager:
         logger.info("SeedVrManager inicializado. O modelo será carregado sob demanda.")
 
     def _download_models_and_configs(self):
-        """
-        Baixa os checkpoints necessários E o arquivo de configuração do VAE que pode estar faltando.
-        """
+        """Baixa os checkpoints necessários E o arquivo de configuração do VAE que pode estar faltando."""
         logger.info("Verificando e baixando modelos e configurações do SeedVR2...")
         ckpt_dir = SEEDVR_REPO_DIR / 'ckpts'
         config_dir = SEEDVR_REPO_DIR / 'configs' / 'vae'
         ckpt_dir.mkdir(exist_ok=True)
         config_dir.mkdir(parents=True, exist_ok=True)
-
-        # Baixa a configuração do VAE para garantir que ela exista
         _load_file_from_url(url=VAE_CONFIG_URL, model_dir=str(config_dir))
-
         pretrain_model_urls = {
             'vae_ckpt': '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',
@@ -113,18 +94,14 @@ class SeedVrManager:
             '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'
         }
-        
         for key, url in pretrain_model_urls.items():
             _load_file_from_url(url=url, model_dir=str(ckpt_dir))
-            
         logger.info("Modelos e configurações do SeedVR2 baixados com sucesso.")
 
     def _initialize_runner(self, model_version: str):
-        """Carrega e configura o modelo SeedVR sob demanda com base na versão selecionada."""
+        """Carrega e configura o modelo SeedVR, com uma correção robusta para o caminho da config do VAE."""
         if self.runner is not None: return
-
         self._download_models_and_configs()
-
         logger.info(f"Inicializando o executor do SeedVR2 {model_version}...")
         if model_version == '3B':
             config_path = SEEDVR_REPO_DIR / 'configs_3b' / 'main.yaml'
@@ -135,23 +112,22 @@ class SeedVrManager:
         else:
             raise ValueError(f"Versão do modelo SeedVR não suportada: {model_version}")
 
-        config = load_config(str(config_path))
-        
+        try:
+            config = load_config(str(config_path))
+        except FileNotFoundError:
+            logger.warning("FileNotFoundError esperado capturado. Carregando config manualmente.")
+            config = OmegaConf.load(str(config_path))
+            correct_vae_config_path = SEEDVR_REPO_DIR / 'configs' / 'vae' / 's8_c16_t4_inflation_sd3.yaml'
+            vae_config = OmegaConf.load(str(correct_vae_config_path))
+            config.vae = vae_config
+            logger.info("Configuração carregada e corrigida manualmente com sucesso.")
+
         self.runner = VideoDiffusionInfer(config)
         OmegaConf.set_readonly(self.runner.config, False)
-        
         self.runner.configure_dit_model(device=self.device, checkpoint=str(checkpoint_path))
-        
-        # --- CORREÇÃO APLICADA AQUI ---
-        correct_vae_config_path = SEEDVR_REPO_DIR / 'configs' / 'vae' / 's8_c16_t4_inflation_sd3.yaml'
-        logger.info(f"Corrigindo o caminho da configuração do VAE para: {correct_vae_config_path}")
-        self.runner.config.vae.config = str(correct_vae_config_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(f"Executor para SeedVR2 {model_version} inicializado e pronto.")