Update managers/seedvr_manager.py

managers/seedvr_manager.py

@@ -1,22 +1,12 @@
 # managers/seedvr_manager.py
-# AducSdr: Uma implementação aberta e funcional da arquitetura ADUC-SDR
-# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
 #
-# Contato:
-# Carlos Rodrigues dos Santos
-# [email protected]
-# Rua Eduardo Carlos Pereira, 4125, B1 Ap32, Curitiba, PR, Brazil, CEP 8102025
+# Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Repositórios e Projetos Relacionados:
-# GitHub: https://github.com/carlex22/Aduc-sdr
+# Version: 2.3.1
 #
-# PENDING PATENT NOTICE: Please see NOTICE.md.
-#
-# Version: 2.3.0
-#
-# This file implements the SeedVrManager, which uses the SeedVR model for
-# video super-resolution. It is self-contained, automatically cloning its own
-# dependencies from the official SeedVR repository.
+# 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.
 
 import torch
 import os
@@ -31,41 +21,46 @@ 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__)
 
-# --- Dependency Management ---
+# --- Gerenciamento de Dependências ---
 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():
     """
-    Ensures the SeedVR repository is cloned and available in the sys.path.
+    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.
     """
     if not SEEDVR_REPO_DIR.exists():
-        logger.info(f"SeedVR repository not found at '{SEEDVR_REPO_DIR}'. Cloning from GitHub...")
+        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
             )
-            logger.info("SeedVR repository cloned successfully.")
+            logger.info("Repositório SeedVR clonado com sucesso.")
         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 SeedVR. Git stderr: {e.stderr}")
+            raise RuntimeError("Não foi possível clonar a dependência necessária do SeedVR do GitHub.")
     else:
-        logger.info("Found local SeedVR repository.")
+        logger.info("Repositório SeedVR local encontrado.")
     
     if str(SEEDVR_REPO_DIR.resolve()) not in sys.path:
         sys.path.insert(0, str(SEEDVR_REPO_DIR.resolve()))
-        logger.info(f"Added '{SEEDVR_REPO_DIR.resolve()}' to sys.path.")
+        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
@@ -78,33 +73,37 @@ 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))
     if not os.path.exists(cached_file):
-        logger.info(f'Downloading: "{url}" to {cached_file}')
+        logger.info(f'Baixando: "{url}" para {cached_file}')
         download_url_to_file(url, cached_file, hash_prefix=None, progress=True)
     return cached_file
 
 class SeedVrManager:
     """
-    Manages the SeedVR model for HD Mastering tasks.
+    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
         self.workspace_dir = workspace_dir
         self.is_initialized = False
-        logger.info("SeedVrManager initialized. Model will be loaded on demand.")
+        logger.info("SeedVrManager inicializado. O modelo será carregado sob demanda.")
 
     def _download_models_and_configs(self):
-        """Downloads the necessary checkpoints AND the missing VAE config file."""
-        logger.info("Verifying and downloading SeedVR2 models and configs...")
+        """
+        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 = {
@@ -118,13 +117,15 @@ class SeedVrManager:
         for key, url in pretrain_model_urls.items():
             _load_file_from_url(url=url, model_dir=str(ckpt_dir))
             
-        logger.info("SeedVR2 models and configs downloaded successfully.")
+        logger.info("Modelos e configurações do SeedVR2 baixados com sucesso.")
 
     def _initialize_runner(self, model_version: str):
-        """Loads and configures the SeedVR model on demand based on the selected version."""
+        """Carrega e configura o modelo SeedVR sob demanda com base na versão selecionada."""
         if self.runner is not None: return
+
         self._download_models_and_configs()
-        logger.info(f"Initializing SeedVR2 {model_version} runner...")
+
+        logger.info(f"Inicializando o executor do SeedVR2 {model_version}...")
         if model_version == '3B':
             config_path = SEEDVR_REPO_DIR / 'configs_3b' / 'main.yaml'
             checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
@@ -132,16 +133,18 @@ class SeedVrManager:
             config_path = SEEDVR_REPO_DIR / 'configs_7b' / 'main.yaml'
             checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_7b.pth'
         else:
-            raise ValueError(f"Unsupported SeedVR model version: {model_version}")
+            raise ValueError(f"Versão do modelo SeedVR não suportada: {model_version}")
 
         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=str(checkpoint_path))
         
-        # --- PATH CORRECTION ---
+        # --- CORREÇÃO APLICADA AQUI ---
         correct_vae_config_path = SEEDVR_REPO_DIR / 'configs' / 'vae' / 's8_c16_t4_inflation_sd3.yaml'
-        logger.info(f"Correcting VAE config path to: {correct_vae_config_path}")
+        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()
@@ -150,18 +153,20 @@ class SeedVrManager:
             self.runner.vae.set_memory_limit(**self.runner.config.vae.memory_limit)
         
         self.is_initialized = True
-        logger.info(f"Runner for SeedVR2 {model_version} initialized and ready.")
-
+        logger.info(f"Executor para SeedVR2 {model_version} inicializado e pronto.")
+    
     def _unload_runner(self):
+        """Remove o executor da VRAM para liberar recursos."""
         if self.runner is not None:
             del self.runner; self.runner = None
             gc.collect(); torch.cuda.empty_cache()
             self.is_initialized = False
-            logger.info("SeedVR2 runner unloaded from VRAM.")
+            logger.info("Executor do SeedVR2 descarregado da VRAM.")
 
     def process_video(self, input_video_path: str, output_video_path: str, prompt: str,
                       model_version: str = '3B', steps: int = 50, seed: int = 666, 
                       progress: gr.Progress = None) -> str:
+        """Aplica o aprimoramento HD a um vídeo usando a lógica do SeedVR."""
         try:
             self._initialize_runner(model_version)
             set_seed(seed, same_across_ranks=True)
@@ -204,10 +209,10 @@ class SeedVrManager:
             final_sample = final_sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
             final_sample_np = final_sample.to(torch.uint8).cpu().numpy()
             mediapy.write_video(output_video_path, final_sample_np, fps=24)
-            logger.info(f"HD Mastered video saved to: {output_video_path}")
+            logger.info(f"Vídeo Masterizado em HD salvo em: {output_video_path}")
             return output_video_path
         finally:
             self._unload_runner()
 
-# --- Singleton Instance ---
+# --- Instância Singleton ---
 seedvr_manager_singleton = SeedVrManager()