Update managers/seedvr_manager.py

managers/seedvr_manager.py

@@ -1,14 +1,12 @@
 # managers/seedvr_manager.py
 #
-# Copyright (C) 2025 Carlos Rodrigues dos Santos
+# Version: 3.2.0 (3B Model Focus)
 #
-# Version: 3.0.1 (Config Path Patch)
-#
-# 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.
+# 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.
 
+# ... (imports permanecem os mesmos) ...
 import torch
 import torch.distributed as dist
 import os
@@ -22,21 +20,54 @@ from torch.hub import download_url_to_file
 import gradio as gr
 import mediapy
 from einops import rearrange
-import shutil # <--- NOVO IMPORT
-
+import shutil
 from tools.tensor_utils import wavelet_reconstruction
 
 logger = logging.getLogger(__name__)
-
-# --- Gerenciamento de Dependências (sem alterações) ---
+# ... (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"
 SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"
 
-def setup_seedvr_dependencies():
-    # ... (sem alterações aqui)
+def setup_seedvr_environment_and_dependencies():
+    """
+    Performs all one-time setup tasks for SeedVR:
+    1. Sets torch.distributed environment variables.
+    2. Forces the installation of flash-attn.
+    3. Clones the SeedVR repository for its code modules.
+    4. Adds the repository to the Python path.
+    """
+    # 1. Configurar variáveis de ambiente para torch.distributed
+    if "MASTER_ADDR" not in os.environ:
+        os.environ["MASTER_ADDR"] = "127.0.0.1"
+        os.environ["MASTER_PORT"] = "12355" # Porta aleatória
+        os.environ["RANK"] = str(0)
+        os.environ["WORLD_SIZE"] = str(1)
+        logger.info("Set up environment variables for torch.distributed.")
+
+    # 2. Forçar a instalação do flash-attn
+    try:
+        import flash_attn
+        logger.info("flash-attn is already installed.")
+    except ImportError:
+        logger.info("Attempting to install flash-attn...")
+        try:
+            subprocess.run(
+                "pip install flash-attn --no-build-isolation",
+                env={**os.environ, "FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
+                shell=True,
+                check=True,
+                capture_output=True, 
+                text=True
+            )
+            logger.info("flash-attn installed successfully.")
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Failed to install flash-attn. Stderr: {e.stderr}")
+
+    # 3. Clonar o repositório do SeedVR Space
     if not SEEDVR_SPACE_DIR.exists():
-        logger.info(f"SeedVR Space not found at '{SEEDVR_SPACE_DIR}'. Cloning from Hugging Face...")
+        logger.info(f"SeedVR Space not found at '{SEEDVR_SPACE_DIR}'. Cloning...")
         try:
             DEPS_DIR.mkdir(exist_ok=True)
             subprocess.run(
@@ -46,28 +77,29 @@ def setup_seedvr_dependencies():
             logger.info("SeedVR Space repository cloned successfully.")
         except subprocess.CalledProcessError as e:
             logger.error(f"Failed to clone SeedVR Space. Git stderr: {e.stderr}")
-            raise RuntimeError("Could not clone the required SeedVR dependency from Hugging Face.")
+            raise RuntimeError("Could not clone the required SeedVR dependency.")
     else:
         logger.info("Found local SeedVR Space repository.")
     
+    # 4. Adicionar o repositório ao path do Python
     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.")
 
-setup_seedvr_dependencies()
+setup_seedvr_environment_and_dependencies()
 
 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
-# ... (outros imports do seedvr sem alterações)
 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
 
 class SeedVrManager:
     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
@@ -75,80 +107,56 @@ 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.
-        """
+        # ... (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"
         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):
-        # ... (sem alterações aqui)
-        logger.info("Verifying and downloading SeedVR2 model checkpoints...")
+        logger.info("Verifying and downloading SeedVR2 3B 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': 'https://huggingface.co/ByteDance-Seed/SeedVR2-7B/resolve/main/seedvr2_ema_7b.pth',
+            # 'dit_7b' REMOVIDO
             '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 model checkpoints downloaded successfully.")
+        logger.info("SeedVR2 3B model checkpoints downloaded successfully.")
 
-    def _initialize_runner(self, model_version: str):
-        """Loads and configures the SeedVR model."""
+    def _initialize_runner(self): # <--- REMOVIDO model_version
         if self.runner is not None: return
-        
-        # 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
         
-        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}")
-
-        # Agora, quando `load_config` for chamado, ele encontrará o arquivo no caminho esperado.
+        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'
+        
         config = load_config(str(config_path))
         
         self.runner = VideoDiffusionInfer(config)
@@ -158,28 +166,28 @@ 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(f"Runner for SeedVR2 {model_version} initialized and ready.")
+        logger.info("Runner for SeedVR2 3B initialized and ready.")
     
-    # ... (o resto da classe e do arquivo permanece o mesmo)
     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,
-                      model_version: str = '3B', steps: int = 50, seed: int = 666, 
-                      progress: gr.Progress = None) -> str:
+                      steps: int = 50, seed: int = 666, 
+                      progress: gr.Progress = None) -> str: # <--- REMOVIDO model_version
         try:
-            self._initialize_runner(model_version)
+            self._initialize_runner() # <--- REMOVIDO 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([
@@ -223,6 +231,7 @@ 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))