Update managers/seedvr_manager.py

managers/seedvr_manager.py

@@ -1,19 +1,21 @@
 # managers/seedvr_manager.py
 #
-# Version: 3.2.0 (3B Model Focus)
+# Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# This version simplifies the manager to exclusively use the SeedVR 3B model.
-# The 7B model download and selection logic have been removed to streamline
-# the code and reduce resource usage.
+# Version: 3.1.0 (Full Environment Setup)
+#
+# This version now fully replicates the environment setup from the original
+# SeedVR Space. It sets the necessary torch.distributed environment variables
+# and forces the installation of flash-attn via subprocess, ensuring complete
+# compatibility and resolving runtime dependency issues.
 
-# ... (imports permanecem os mesmos) ...
 import torch
 import torch.distributed as dist
 import os
 import gc
 import logging
 import sys
-import subprocess
+import subprocess # <--- NOVO IMPORT
 from pathlib import Path
 from urllib.parse import urlparse
 from torch.hub import download_url_to_file
@@ -21,10 +23,11 @@ import gradio as gr
 import mediapy
 from einops import rearrange
 import shutil
+
 from tools.tensor_utils import wavelet_reconstruction
 
 logger = logging.getLogger(__name__)
-# ... (setup_seedvr_environment_and_dependencies e imports do seedvr permanecem os mesmos) ...
+
 # --- INÍCIO DA SEÇÃO DE GERENCIAMENTO DE DEPENDÊNCIAS E AMBIENTE ---
 DEPS_DIR = Path("./deps")
 SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
@@ -64,6 +67,8 @@ def setup_seedvr_environment_and_dependencies():
             logger.info("flash-attn installed successfully.")
         except subprocess.CalledProcessError as e:
             logger.error(f"Failed to install flash-attn. Stderr: {e.stderr}")
+            # Não lançamos um erro aqui, pois pode não ser fatal em todos os sistemas
+            # O import posterior vai falhar se for realmente necessário.
 
     # 3. Clonar o repositório do SeedVR Space
     if not SEEDVR_SPACE_DIR.exists():
@@ -86,19 +91,21 @@ def setup_seedvr_environment_and_dependencies():
         sys.path.insert(0, str(SEEDVR_SPACE_DIR.resolve()))
         logger.info(f"Added '{SEEDVR_SPACE_DIR.resolve()}' to sys.path.")
 
+# Executa o setup completo uma única vez
 setup_seedvr_environment_and_dependencies()
 
+# Agora que o setup está completo, os imports devem funcionar
 from projects.video_diffusion_sr.infer import VideoDiffusionInfer
 from common.config import load_config
 from common.seed import set_seed
+# ... (outros imports do seedvr)
 from torchvision.io.video import read_video
 from omegaconf import OmegaConf
-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
+# --- FIM DA SEÇÃO DE SETUP ---
+
 
 class SeedVrManager:
+    # ... (o resto do código permanece o mesmo da nossa última versão) ...
     def __init__(self, workspace_dir="deformes_workspace"):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self.runner = None
@@ -108,7 +115,6 @@ class SeedVrManager:
         logger.info("SeedVrManager initialized. Model will be loaded on demand.")
 
     def _patch_config_paths(self):
-        # ... (sem alterações) ...
         app_root = Path("/home/user/app")
         source_config_dir = SEEDVR_SPACE_DIR / "models" / "video_vae_v3"
         target_config_parent_dir = app_root / "models"
@@ -129,36 +135,38 @@ class SeedVrManager:
             raise IOError("Could not patch the required SeedVR configuration paths.")
 
     def _download_models(self):
-        logger.info("Verifying and downloading SeedVR2 3B model checkpoints...")
+        logger.info("Verifying and downloading SeedVR2 model checkpoints...")
         ckpt_dir = SEEDVR_SPACE_DIR / 'ckpts'
         ckpt_dir.mkdir(exist_ok=True)
-        
         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',
-            # 'dit_7b' REMOVIDO
+            'dit_7b': 'https://huggingface.co/ByteDance-Seed/SeedVR2-7B/resolve/main/seedvr2_ema_7b.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'
         }
         for key, url in pretrain_model_urls.items():
             _load_file_from_url(url=url, model_dir=str(ckpt_dir))
-        logger.info("SeedVR2 3B model checkpoints downloaded successfully.")
+        logger.info("SeedVR2 model checkpoints downloaded successfully.")
 
-    def _initialize_runner(self): # <--- REMOVIDO model_version
+    def _initialize_runner(self, model_version: str):
         if self.runner is not None: return
         self._patch_config_paths()
         self._download_models()
-
         if dist.is_available() and not dist.is_initialized():
+            logger.info("Applying patch to disable torch.distributed.barrier for single-GPU inference.")
             self._original_barrier = dist.barrier
             dist.barrier = lambda *args, **kwargs: None
-        
-        logger.info("Initializing SeedVR2 3B runner...")
-        config_path = SEEDVR_SPACE_DIR / 'configs_3b' / 'main.yaml'
-        checkpoint_path = SEEDVR_SPACE_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
-        
+        logger.info(f"Initializing SeedVR2 {model_version} runner...")
+        if model_version == '3B':
+            config_path = SEEDVR_SPACE_DIR / 'configs_3b' / 'main.yaml'
+            checkpoint_path = SEEDVR_SPACE_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
+        elif model_version == '7B':
+            config_path = SEEDVR_SPACE_DIR / 'configs_7b' / 'main.yaml'
+            checkpoint_path = SEEDVR_SPACE_DIR / 'ckpts' / 'seedvr2_ema_7b.pth'
+        else:
+            raise ValueError(f"Unsupported SeedVR model version: {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))
@@ -166,28 +174,27 @@ class SeedVrManager:
         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 for SeedVR2 3B initialized and ready.")
+        logger.info(f"Runner for SeedVR2 {model_version} initialized and ready.")
     
     def _unload_runner(self):
-        # ... (sem alterações) ...
         if self.runner is not None:
             del self.runner; self.runner = None
             gc.collect(); torch.cuda.empty_cache()
             self.is_initialized = False
             logger.info("SeedVR runner unloaded from VRAM.")
         if self._original_barrier is not None:
+            logger.info("Restoring original torch.distributed.barrier function.")
             dist.barrier = self._original_barrier
             self._original_barrier = None
 
     def process_video(self, input_video_path: str, output_video_path: str, prompt: str,
-                      steps: int = 50, seed: int = 666, 
-                      progress: gr.Progress = None) -> str: # <--- REMOVIDO model_version
+                      model_version: str = '3B', steps: int = 50, seed: int = 666, 
+                      progress: gr.Progress = None) -> str:
         try:
-            self._initialize_runner() # <--- REMOVIDO model_version
+            self._initialize_runner(model_version)
             set_seed(seed, same_across_ranks=True)
             self.runner.config.diffusion.timesteps.sampling.steps = steps
             self.runner.configure_diffusion()
-            # ... (resto da função sem alterações) ...
             video_tensor = read_video(input_video_path, output_format="TCHW")[0] / 255.0
             res_h, res_w = video_tensor.shape[-2:]
             video_transform = Compose([
@@ -231,7 +238,6 @@ class SeedVrManager:
             self._unload_runner()
 
 def _load_file_from_url(url, model_dir='./', file_name=None):
-    # ... (sem alterações) ...
     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))