Update aduc_framework/managers/seedvr_manager.py

aduc_framework/managers/seedvr_manager.py

@@ -1,5 +1,5 @@
-# hd_specialist.py (Versão Final - Estrutura de Arquivos Corrigida)
-#https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B
+# hd_specialist.py (Versão Final - Estrutura de Arquivos Corrigida e Autossuficiente)
+# https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B
 
 import torch
 import imageio
@@ -18,6 +18,7 @@ from omegaconf import OmegaConf
 import sys
 
 
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
 APP_ROOT = Path("/home/user/app")
@@ -109,19 +110,16 @@ class SeedVrManager:
         ckpt_dir.mkdir(exist_ok=True)
 
         pretrain_model_url = {
-            'configs_3b': 'https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B/tree/main/configs_3b/main.yaml',
             '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',
             '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'
         }
-
-        _load_file_from_url(url=pretrain_model_url['configs_3b'], model_dir='./ckpts/')
+        
         _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.")
 
     def _initialize_runner(self):
@@ -133,8 +131,27 @@ class SeedVrManager:
         self._download_models()
 
         logger.info("Inicializando o runner do SeedVR2...")
-        config_path = os.path.join('./configs_3b', 'main.yaml')
-        config = load_config(config_path)
+        
+        # --- 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"
+        
+        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 ---
+
+        config = load_config(str(config_path))
         
         self.runner = VideoDiffusionInfer(config)
         OmegaConf.set_readonly(self.runner.config, False)
@@ -161,7 +178,7 @@ class SeedVrManager:
     
 
     def process_video(self, input_video_path: str, output_video_path: str, prompt: str, 
-                      model_version: str = '3B') -> str:
+                      seed: int = 42, fps_out: int = 30) -> str:
         """
         Aplica o aprimoramento HD a um vídeo usando a lógica oficial do SeedVR.
         """
@@ -232,5 +249,5 @@ class SeedVrManager:
         finally:
             self._unload_runner()        
 
-# Instância Singleton
-seedvr_manager_singleton = SeedVrManager()
+# Instância Singleton para facilitar o uso (opcional)
+# seedvr_manager_singleton = SeedVrManager()