Rename managers/hd_specialist.py to managers/seedvr_manager.py

managers/{hd_specialist.py → seedvr_manager.py}

@@ -1,14 +1,12 @@
-# managers/hd_specialist.py
+# managers/seedvr_manager.py
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
-# Version: 2.2.0
+# Version: 2.3.0
 #
-# This file implements the HD Specialist (Δ+), which uses the SeedVR model
-# for video super-resolution. It has been refactored to be self-contained by
-# automatically cloning its own dependencies from the official SeedVR repository
-# if they are not found locally. This removes the need for manual file copying
-# and makes the ADUC-SDR framework more robust and portable.
+# This file implements the SeedVrManager, which uses the SeedVR model for
+# video super-resolution. It is self-contained, automatically cloning its own
+# dependencies from the official SeedVR repository.
 
 import torch
 import os
@@ -23,6 +21,9 @@ 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 ---
@@ -30,6 +31,47 @@ DEPS_DIR = Path("./deps")
 SEEDVR_REPO_DIR = DEPS_DIR / "SeedVR"
 SEEDVR_REPO_URL = "https://github.com/ByteDance-Seed/SeedVR.git"
 
+def setup_seedvr_dependencies():
+    """
+    Ensures the SeedVR repository is cloned and available in the sys.path.
+    This function is run once when the module is first imported.
+    """
+    if not SEEDVR_REPO_DIR.exists():
+        logger.info(f"SeedVR repository not found at '{SEEDVR_REPO_DIR}'. Cloning from GitHub...")
+        try:
+            DEPS_DIR.mkdir(exist_ok=True)
+            # Use --depth 1 for a shallow clone to save space and time
+            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.")
+        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.")
+    else:
+        logger.info("Found local SeedVR repository.")
+    
+    # Add the cloned repo to Python's path to allow direct imports
+    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.")
+
+# --- Execute dependency setup immediately upon module import ---
+setup_seedvr_dependencies()
+
+# --- Now that the path is set, we can safely import from the cloned repo ---
+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
+from torchvision.io.video import read_video
+from omegaconf import OmegaConf
+
+
 def _load_file_from_url(url, model_dir='./', file_name=None):
     """Helper function to download files from a URL to a local directory."""
     os.makedirs(model_dir, exist_ok=True)
@@ -40,67 +82,16 @@ 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
 
-class HDSpecialist:
+class SeedVrManager:
     """
-    Implements the HD Specialist (Δ+) using the SeedVR infrastructure.
-    Manages model loading, inference, and memory on demand.
+    Manages the SeedVR model for HD Mastering tasks.
     """
     def __init__(self, workspace_dir="deformes_workspace"):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
         self.runner = None
         self.workspace_dir = workspace_dir
         self.is_initialized = False
-        self._seedvr_modules_loaded = False
-        self._setup_dependencies()
-        logger.info("HD Specialist (SeedVR) initialized. Dependencies checked. Model will be loaded on demand.")
-
-    def _setup_dependencies(self):
-        """
-        Checks for the SeedVR repository locally. If not found, clones it.
-        Then, it adds the repository to the Python path to make its modules importable.
-        """
-        if not SEEDVR_REPO_DIR.exists():
-            logger.info(f"SeedVR repository not found at '{SEEDVR_REPO_DIR}'. Cloning from GitHub...")
-            try:
-                DEPS_DIR.mkdir(exist_ok=True)
-                subprocess.run(
-                    ["git", "clone", SEEDVR_REPO_URL, str(SEEDVR_REPO_DIR)],
-                    check=True, capture_output=True, text=True
-                )
-                logger.info("SeedVR repository 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.")
-        else:
-            logger.info("Found local SeedVR repository.")
-        
-        # Add the cloned repo to Python's path to allow direct imports
-        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.")
-    
-    def _lazy_load_seedvr_modules(self):
-        """
-        Dynamically imports SeedVR modules only when needed.
-        This prevents ImportError if the class is instantiated before dependencies are ready.
-        """
-        if self._seedvr_modules_loaded:
-            return
-
-        global VideoDiffusionInfer, load_config, set_seed, DivisibleCrop, NaResize, Rearrange, wavelet_reconstruction, Compose, Lambda, Normalize, read_video, OmegaConf
-        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 projects.video_diffusion_sr.color_fix import wavelet_reconstruction
-        from torchvision.transforms import Compose, Lambda, Normalize
-        from torchvision.io.video import read_video
-        from omegaconf import OmegaConf
-        
-        self._seedvr_modules_loaded = True
-        logger.info("SeedVR modules have been dynamically loaded.")
+        logger.info("SeedVrManager initialized. Model will be loaded on demand.")
 
     def _download_models(self):
         """Downloads the necessary checkpoints for SeedVR2."""
@@ -123,10 +114,8 @@ class HDSpecialist:
 
     def _initialize_runner(self, model_version: str):
         """Loads and configures the SeedVR model on demand based on the selected version."""
-        if self.runner is not None:
-            return
-        
-        self._lazy_load_seedvr_modules()
+        if self.runner is not None: return
+
         self._download_models()
 
         logger.info(f"Initializing SeedVR2 {model_version} runner...")
@@ -156,10 +145,8 @@ class HDSpecialist:
     def _unload_runner(self):
         """Removes the runner from VRAM to free resources."""
         if self.runner is not None:
-            del self.runner
-            self.runner = None
-            gc.collect()
-            torch.cuda.empty_cache()
+            del self.runner; self.runner = None
+            gc.collect(); torch.cuda.empty_cache()
             self.is_initialized = False
             logger.info("SeedVR2 runner unloaded from VRAM.")
 
@@ -204,12 +191,7 @@ class HDSpecialist:
             conditions = [self.runner.get_condition(noise, latent_blur=latent, task="sr") for noise, latent in zip(noises, cond_latents)]
 
             with torch.no_grad(), torch.autocast("cuda", torch.bfloat16, enabled=True):
-                video_tensors = self.runner.inference(
-                    noises=noises,
-                    conditions=conditions,
-                    dit_offload=True,
-                    **text_embeds_dict,
-                )
+                video_tensors = self.runner.inference(noises=noises, conditions=conditions, dit_offload=True, **text_embeds_dict)
             
             self.runner.dit.to("cpu"); gc.collect(); torch.cuda.empty_cache()
 
@@ -223,10 +205,10 @@ class HDSpecialist:
                 input_video_sample = input_video_sample[:, :final_sample.shape[1]]
 
             final_sample = wavelet_reconstruction(
-                rearrange(final_sample, "c t h w -> t c h w"),
+                rearrange(final_sample, "c t h w -> t c h w"), 
                 rearrange(input_video_sample, "c t h w -> t c h w")
             )
-
+            
             final_sample = rearrange(final_sample, "t c h w -> t h w c")
             final_sample = final_sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round()
             final_sample_np = final_sample.to(torch.uint8).cpu().numpy()
@@ -234,9 +216,8 @@ class HDSpecialist:
             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
-
         finally:
             self._unload_runner()
 
-# Singleton instance
-hd_specialist_singleton = HDSpecialist()
+# --- Singleton Instance ---
+seedvr_manager_singleton = SeedVrManager()