Update managers/seedvr_manager.py

managers/seedvr_manager.py

@@ -2,11 +2,11 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 2.3.4
+# Version: 2.3.5
 #
-# 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.
+# This version uses the optimal strategy of cloning the self-contained Hugging Face
+# Space repository and uses the full, correct import paths to resolve all
+# ModuleNotFoundErrors, while retaining necessary runtime patches.
 
 import torch
 import torch.distributed as dist
@@ -22,17 +22,15 @@ 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")
-# 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"
+VAE_CONFIG_URL = "https://raw.githubusercontent.com/ByteDance-Seed/SeedVR/main/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml"
 
 def setup_seedvr_dependencies():
     """
@@ -42,7 +40,6 @@ def setup_seedvr_dependencies():
         logger.info(f"SeedVR Space not found at '{SEEDVR_SPACE_DIR}'. Cloning from Hugging Face...")
         try:
             DEPS_DIR.mkdir(exist_ok=True)
-            # 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
@@ -60,8 +57,8 @@ def setup_seedvr_dependencies():
 
 setup_seedvr_dependencies()
 
-# The imports from a Space are often directly from the root
-from infer import VideoDiffusionInfer
+# Use full import paths relative to the root of the cloned repository
+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
@@ -71,6 +68,7 @@ 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)
@@ -90,11 +88,14 @@ class SeedVrManager:
         self._original_barrier = None
         logger.info("SeedVrManager initialized. Model will be loaded on demand.")
 
-    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
+    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_SPACE_DIR / 'ckpts'
+        config_dir = SEEDVR_SPACE_DIR / 'configs' / 'vae'
         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',
@@ -104,41 +105,42 @@ 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 downloaded successfully.")
+        logger.info("SeedVR2 models and configs downloaded successfully.")
 
     def _initialize_runner(self, model_version: str):
-        """Loads and configures the SeedVR model."""
+        """Loads and configures the SeedVR model, with patches for single-GPU inference."""
         if self.runner is not None: return
-        self._download_models()
-        
+        self._download_models_and_configs()
+
         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 from Space repo...")
+        logger.info(f"Initializing SeedVR2 {model_version} runner...")
         if model_version == '3B':
-            config_path = SEEDVR_SPACE_DIR / 'configs' / 'generate.yaml' # Typical path in a Space
-            checkpoint_path = SEEDVR_SPACE_DIR / 'ckpt' / 'VINCIE-3B' / 'dit.pth'
+            config_path = SEEDVR_SPACE_DIR / 'configs_3b' / 'main.yaml'
+            checkpoint_path = SEEDVR_SPACE_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
         elif model_version == '7B':
-            # 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'
+            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))
-        
+        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_SPACE_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.")
+
         self.runner = VideoDiffusionInfer(config)
         OmegaConf.set_readonly(self.runner.config, False)
-        # 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_dit_model(device=self.device, checkpoint=str(checkpoint_path))
         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
@@ -181,8 +183,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_SPACE_DIR / 'ckpt' / 'pos_emb.pt'
-            neg_emb_path = SEEDVR_SPACE_DIR / 'ckpt' / 'neg_emb.pt'
+            pos_emb_path = SEEDVR_SPACE_DIR / 'ckpts' / 'pos_emb.pt'
+            neg_emb_path = SEEDVR_SPACE_DIR / 'ckpts' / '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]}
@@ -208,4 +210,4 @@ class SeedVrManager:
             self._unload_runner()
 
 # --- Singleton Instance ---
-seedvr_manager_singleton = SeedVrManager()
+seedvr_manager_singleton = SeedVrManager()```