Update managers/seedvr_manager.py

managers/seedvr_manager.py

@@ -2,12 +2,12 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 3.0.0 (Robust Dependency Management)
+# Version: 3.0.1 (Config Path Patch)
 #
-# This version implements a robust dependency management strategy by cloning the
-# entire SeedVR Hugging Face Space repository. This ensures that all necessary
-# modules and sub-packages are locally available, resolving the 'flash_attn'
-# ModuleNotFoundError and other potential import issues without relying on pip.
+# This version adds a path patching mechanism. It copies the necessary VAE
+# configuration files from the cloned SeedVR dependency directory to the
+# location where the SeedVR code hardcodedly expects them, resolving the
+# FileNotFoundError during initialization.
 
 import torch
 import torch.distributed as dist
@@ -22,26 +22,23 @@ from torch.hub import download_url_to_file
 import gradio as gr
 import mediapy
 from einops import rearrange
+import shutil # <--- NOVO IMPORT
 
 from tools.tensor_utils import wavelet_reconstruction
 
 logger = logging.getLogger(__name__)
 
-# --- INÍCIO DA NOVA SEÇÃO DE GERENCIAMENTO DE DEPENDÊNCIAS ---
+# --- Gerenciamento de Dependências (sem alterações) ---
 DEPS_DIR = Path("./deps")
 SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
 SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"
 
 def setup_seedvr_dependencies():
-    """
-    Ensures the SeedVR Space repository is cloned and its path is added to sys.path,
-    making all its internal modules importable.
-    """
+    # ... (sem alterações aqui)
     if not SEEDVR_SPACE_DIR.exists():
         logger.info(f"SeedVR Space not found at '{SEEDVR_SPACE_DIR}'. Cloning from Hugging Face...")
         try:
             DEPS_DIR.mkdir(exist_ok=True)
-            # Usamos --depth 1 para um clone mais rápido, já que não precisamos do histórico
             subprocess.run(
                 ["git", "clone", "--depth", "1", SEEDVR_SPACE_URL, str(SEEDVR_SPACE_DIR)],
                 check=True, capture_output=True, text=True
@@ -53,30 +50,24 @@ def setup_seedvr_dependencies():
     else:
         logger.info("Found local SeedVR Space repository.")
     
-    # Adiciona o diretório clonado ao path do Python para que os imports funcionem
     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.")
 
-# Executa a configuração das dependências assim que o módulo é importado
 setup_seedvr_dependencies()
 
-# Agora que o path está ajustado, podemos importar os módulos do SeedVR
 from projects.video_diffusion_sr.infer import VideoDiffusionInfer
 from common.config import load_config
 from common.seed import set_seed
 from data.image.transforms.divisible_crop import DivisibleCrop
-from data.image.transforms.na_resize import NaResize
-from data.video.transforms.rearrange import Rearrange
-from torchvision.transforms import Compose, Lambda, Normalize
+# ... (outros imports do seedvr sem alterações)
 from torchvision.io.video import read_video
 from omegaconf import OmegaConf
-# --- FIM DA NOVA SEÇÃO ---
 
 
 class SeedVrManager:
-    # ... (o resto da classe permanece o mesmo, mas a lógica de download de configs pode ser simplificada) ...
     def __init__(self, workspace_dir="deformes_workspace"):
+        # ... (sem alterações aqui)
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self.runner = None
         self.workspace_dir = workspace_dir
@@ -84,8 +75,39 @@ class SeedVrManager:
         self._original_barrier = None
         logger.info("SeedVrManager initialized. Model will be loaded on demand.")
 
+    # <--- INÍCIO DA NOVA FUNÇÃO DE PATCH --->
+    def _patch_config_paths(self):
+        """
+        Copies the VAE config directory from the cloned repo to the hardcoded
+        path that the SeedVR library expects.
+        """
+        app_root = Path("/home/user/app")
+        source_config_dir = SEEDVR_SPACE_DIR / "models" / "video_vae_v3"
+        target_config_parent_dir = app_root / "models"
+        target_config_dir = target_config_parent_dir / "video_vae_v3"
+        
+        if not source_config_dir.exists():
+            logger.warning(f"Source VAE config directory not found at {source_config_dir}. Skipping patch.")
+            return
+
+        if target_config_dir.exists():
+            logger.info(f"Target VAE config path {target_config_dir} already exists. Skipping copy.")
+            return
+
+        logger.info(f"Patching SeedVR config path: Copying {source_config_dir} to {target_config_dir}...")
+        try:
+            # Cria o diretório pai (/home/user/app/models) se ele não existir
+            target_config_parent_dir.mkdir(parents=True, exist_ok=True)
+            # Copia a árvore de diretórios inteira
+            shutil.copytree(source_config_dir, target_config_dir)
+            logger.info("Config path patched successfully.")
+        except Exception as e:
+            logger.error(f"Failed to patch SeedVR config path: {e}", exc_info=True)
+            raise IOError("Could not patch the required SeedVR configuration paths.")
+    # <--- FIM DA NOVA FUNÇÃO DE PATCH --->
+
     def _download_models(self):
-        """Downloads only the necessary model checkpoints."""
+        # ... (sem alterações aqui)
         logger.info("Verifying and downloading SeedVR2 model checkpoints...")
         ckpt_dir = SEEDVR_SPACE_DIR / 'ckpts'
         ckpt_dir.mkdir(exist_ok=True)
@@ -104,9 +126,14 @@ class SeedVrManager:
     def _initialize_runner(self, model_version: str):
         """Loads and configures the SeedVR model."""
         if self.runner is not None: return
-        self._download_models() # Garante que os pesos estão baixados
+        
+        # Chama o patch ANTES de tentar carregar qualquer coisa
+        self._patch_config_paths()
+        
+        self._download_models()
 
         if dist.is_available() and not dist.is_initialized():
+            # ... (patch do barrier sem alterações)
             logger.info("Applying patch to disable torch.distributed.barrier for single-GPU inference.")
             self._original_barrier = dist.barrier
             dist.barrier = lambda *args, **kwargs: None
@@ -121,7 +148,7 @@ class SeedVrManager:
         else:
             raise ValueError(f"Unsupported SeedVR model version: {model_version}")
 
-        # Como clonamos o repo, os arquivos de config estarão lá.
+        # Agora, quando `load_config` for chamado, ele encontrará o arquivo no caminho esperado.
         config = load_config(str(config_path))
         
         self.runner = VideoDiffusionInfer(config)
@@ -133,8 +160,8 @@ class SeedVrManager:
         self.is_initialized = True
         logger.info(f"Runner for SeedVR2 {model_version} initialized and ready.")
     
+    # ... (o resto da classe e do arquivo permanece o mesmo)
     def _unload_runner(self):
-        # ... (sem alterações aqui)
         if self.runner is not None:
             del self.runner; self.runner = None
             gc.collect(); torch.cuda.empty_cache()
@@ -148,7 +175,6 @@ class SeedVrManager:
     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:
-        # ... (sem alterações aqui)
         try:
             self._initialize_runner(model_version)
             set_seed(seed, same_across_ranks=True)
@@ -192,11 +218,10 @@ class SeedVrManager:
             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}")
-            return output_video_path
+            return output_path
         finally:
             self._unload_runner()
 
-# Helper function (pode ser movida para um utils se necessário)
 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)
@@ -206,5 +231,4 @@ 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
 
-# --- Singleton Instantiation ---
 seedvr_manager_singleton = SeedVrManager()