Update managers/seedvr_manager.py

managers/seedvr_manager.py

@@ -2,11 +2,11 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 2.3.3
+# Version: 2.3.4
 #
-# This version adds a monkey patch to disable torch.distributed.barrier calls
-# within the SeedVR library, allowing it to run in a single-GPU inference mode
-# without raising a "process group not initialized" error.
+# This version is optimized for Hugging Face Spaces environments. It now clones
+# the dependency directly from the official SeedVR HF Space, which is faster,
+# lighter, and more reliable than cloning from GitHub.
 
 import torch
 import torch.distributed as dist
@@ -22,37 +22,46 @@ import gradio as gr
 import mediapy
 from einops import rearrange
 
+# Internalized utility for color correction, ensuring stability.
 from tools.tensor_utils import wavelet_reconstruction
 
 logger = logging.getLogger(__name__)
 
 # --- Dependency Management ---
 DEPS_DIR = Path("./deps")
-SEEDVR_REPO_DIR = DEPS_DIR / "SeedVR"
-SEEDVR_REPO_URL = "https://github.com/ByteDance-Seed/SeedVR.git"
-VAE_CONFIG_URL = "https://raw.githubusercontent.com/ByteDance-Seed/SeedVR/main/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml"
+# Renamed to reflect the new source
+SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
+# NEW: Cloning from the HF Space directly is much more efficient
+SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"
 
 def setup_seedvr_dependencies():
-    """Ensures the SeedVR repository is cloned and available in the sys.path."""
-    if not SEEDVR_REPO_DIR.exists():
-        logger.info(f"SeedVR repository not found at '{SEEDVR_REPO_DIR}'. Cloning from GitHub...")
+    """
+    Ensures the SeedVR Space repository is cloned and available in the sys.path.
+    """
+    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)
-            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.")
+            # We clone the entire space repo to get its file structure
+            subprocess.run(
+                ["git", "clone", SEEDVR_SPACE_URL, str(SEEDVR_SPACE_DIR)],
+                check=True, capture_output=True, text=True
+            )
+            logger.info("SeedVR Space cloned successfully.")
         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"Failed to clone SeedVR Space. Git stderr: {e.stderr}")
+            raise RuntimeError("Could not clone the required SeedVR dependency from Hugging Face.")
     else:
-        logger.info("Found local SeedVR repository.")
+        logger.info("Found local SeedVR Space repository.")
     
-    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.")
+    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()
 
-from projects.video_diffusion_sr.infer import VideoDiffusionInfer
+# The imports from a Space are often directly from the root
+from infer import VideoDiffusionInfer
 from common.config import load_config
 from common.seed import set_seed
 from data.image.transforms.divisible_crop import DivisibleCrop
@@ -62,7 +71,6 @@ from torchvision.transforms import Compose, Lambda, Normalize
 from torchvision.io.video import read_video
 from omegaconf import OmegaConf
 
-
 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)
@@ -79,17 +87,14 @@ class SeedVrManager:
         self.runner = None
         self.workspace_dir = workspace_dir
         self.is_initialized = False
-        self._original_barrier = None # To store the original distributed barrier function
+        self._original_barrier = None
         logger.info("SeedVrManager initialized. Model will be loaded on demand.")
 
-    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...")
-        ckpt_dir = SEEDVR_REPO_DIR / 'ckpts'
-        config_dir = SEEDVR_REPO_DIR / 'configs' / 'vae'
+    def _download_models(self):
+        """Downloads the necessary checkpoints for SeedVR2."""
+        logger.info("Verifying and downloading SeedVR2 models...")
+        ckpt_dir = SEEDVR_SPACE_DIR / 'ckpt' # Note: Path in Space repo might be different
         ckpt_dir.mkdir(exist_ok=True)
-        config_dir.mkdir(parents=True, exist_ok=True)
-        _load_file_from_url(url=VAE_CONFIG_URL, model_dir=str(config_dir))
         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',
@@ -99,55 +104,53 @@ 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("SeedVR2 models downloaded successfully.")
 
     def _initialize_runner(self, model_version: str):
-        """Loads and configures the SeedVR model, with patches for single-GPU inference."""
+        """Loads and configures the SeedVR model."""
         if self.runner is not None: return
-        self._download_models_and_configs()
-
+        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(f"Initializing SeedVR2 {model_version} runner...")
+        logger.info(f"Initializing SeedVR2 {model_version} runner from Space repo...")
         if model_version == '3B':
-            config_path = SEEDVR_REPO_DIR / 'configs_3b' / 'main.yaml'
-            checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
+            config_path = SEEDVR_SPACE_DIR / 'configs' / 'generate.yaml' # Typical path in a Space
+            checkpoint_path = SEEDVR_SPACE_DIR / 'ckpt' / 'VINCIE-3B' / 'dit.pth'
         elif model_version == '7B':
-            config_path = SEEDVR_REPO_DIR / 'configs_7b' / 'main.yaml'
-            checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_7b.pth'
+            # Assuming a similar structure for a 7B space if it existed
+            config_path = SEEDVR_SPACE_DIR / 'configs' / 'generate_7b.yaml' 
+            checkpoint_path = SEEDVR_SPACE_DIR / 'ckpt' / 'VINCIE-7B' / 'dit.pth'
         else:
             raise ValueError(f"Unsupported SeedVR model version: {model_version}")
 
-        try:
-            config = load_config(str(config_path))
-        except FileNotFoundError:
-            logger.warning("Caught expected FileNotFoundError. Loading config manually.")
-            config = OmegaConf.load(str(config_path))
-            correct_vae_config_path = SEEDVR_REPO_DIR / 'configs' / 'vae' / 's8_c16_t4_inflation_sd3.yaml'
-            vae_config = OmegaConf.load(str(correct_vae_config_path))
-            config.vae = vae_config
-            logger.info("Configuration loaded and patched manually.")
-
+        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))
+        # Manually set the correct checkpoint paths since the config inside the space might be relative
+        self.runner.config.dit.checkpoint = str(checkpoint_path)
+        self.runner.config.vae.checkpoint = str(SEEDVR_SPACE_DIR / 'ckpt' / 'VINCIE-3B' / 'vae.pth')
+        self.runner.config.text.models[0].path = str(SEEDVR_SPACE_DIR / 'ckpt' / 'VINCIE-3B' / 'llm14b')
+        
+        self.runner.configure_dit_model(device=self.device, checkpoint=self.runner.config.dit.checkpoint)
         self.runner.configure_vae_model()
+        
         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.")
     
     def _unload_runner(self):
-        """Unloads the runner from VRAM and restores any applied patches."""
+        """Unloads the runner from VRAM and restores patches."""
         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
@@ -156,7 +159,7 @@ 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:
-        """Applies HD enhancement to a video using the SeedVR logic."""
+        """Applies HD enhancement to a video."""
         try:
             self._initialize_runner(model_version)
             set_seed(seed, same_across_ranks=True)
@@ -178,8 +181,8 @@ class SeedVrManager:
             cond_latents = self.runner.vae_encode(cond_latents)
             self.runner.vae.to("cpu"); gc.collect(); torch.cuda.empty_cache()
             self.runner.dit.to(self.device)
-            pos_emb_path = SEEDVR_REPO_DIR / 'ckpts' / 'pos_emb.pt'
-            neg_emb_path = SEEDVR_REPO_DIR / 'ckpts' / 'neg_emb.pt'
+            pos_emb_path = SEEDVR_SPACE_DIR / 'ckpt' / 'pos_emb.pt'
+            neg_emb_path = SEEDVR_SPACE_DIR / 'ckpt' / 'neg_emb.pt'
             text_pos_embeds = torch.load(pos_emb_path).to(self.device)
             text_neg_embeds = torch.load(neg_emb_path).to(self.device)
             text_embeds_dict = {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}