Update managers/hd_specialist.py

managers/hd_specialist.py

@@ -2,43 +2,34 @@
 #
 # Copyright (C) 2025 Carlos Rodrigues dos Santos
 #
+# Version: 2.2.0
+#
 # This file implements the HD Specialist (Δ+), which uses the SeedVR model
-# for video super-resolution. It's designed to be called by the ADUC orchestrator
-# to perform the final HD mastering pass on a generated video. It manages the
-# loading/unloading of the heavy SeedVR models to conserve VRAM and can switch
-# between different model sizes (e.g., 3B and 7B).
+# 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.
 
 import torch
-import gradio as gr
-import imageio
 import os
 import gc
 import logging
-import numpy as np
-from PIL import Image
-from tqdm import tqdm
-import shlex
+import sys
 import subprocess
 from pathlib import Path
 from urllib.parse import urlparse
 from torch.hub import download_url_to_file
-from omegaconf import OmegaConf
+import gradio as gr
 import mediapy
 from einops import rearrange
 
-# Assuming these files are in the project structure
-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
-
 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"
+
 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)
@@ -59,12 +50,62 @@ class HDSpecialist:
         self.runner = None
         self.workspace_dir = workspace_dir
         self.is_initialized = False
-        logger.info("HD Specialist (SeedVR) initialized. Model will be loaded on demand.")
+        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.")
 
     def _download_models(self):
         """Downloads the necessary checkpoints for SeedVR2."""
         logger.info("Verifying and downloading SeedVR2 models...")
-        ckpt_dir = Path('./ckpts')
+        ckpt_dir = SEEDVR_REPO_DIR / 'ckpts'
         ckpt_dir.mkdir(exist_ok=True)
 
         pretrain_model_urls = {
@@ -76,7 +117,7 @@ class HDSpecialist:
         }
         
         for key, url in pretrain_model_urls.items():
-            _load_file_from_url(url=url, model_dir='./ckpts/')
+            _load_file_from_url(url=url, model_dir=str(ckpt_dir))
             
         logger.info("SeedVR2 models downloaded successfully.")
 
@@ -84,25 +125,26 @@ class HDSpecialist:
         """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()
         self._download_models()
 
         logger.info(f"Initializing SeedVR2 {model_version} runner...")
         if model_version == '3B':
-            config_path = os.path.join('./configs_3b', 'main.yaml')
-            checkpoint_path = './ckpts/seedvr2_ema_3b.pth'
+            config_path = SEEDVR_REPO_DIR / 'configs_3b' / 'main.yaml'
+            checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_3b.pth'
         elif model_version == '7B':
-            config_path = os.path.join('./configs_7b', 'main.yaml')
-            checkpoint_path = './ckpts/seedvr2_ema_7b.pth'
+            config_path = SEEDVR_REPO_DIR / 'configs_7b' / 'main.yaml'
+            checkpoint_path = SEEDVR_REPO_DIR / 'ckpts' / 'seedvr2_ema_7b.pth'
         else:
             raise ValueError(f"Unsupported SeedVR model version: {model_version}")
 
-        config = load_config(config_path)
+        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=checkpoint_path)
+        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"):
@@ -129,7 +171,6 @@ class HDSpecialist:
             self._initialize_runner(model_version)
             set_seed(seed, same_across_ranks=True)
 
-            # --- Adapted inference logic from SeedVR scripts ---
             self.runner.config.diffusion.timesteps.sampling.steps = steps
             self.runner.configure_diffusion()
 
@@ -153,8 +194,10 @@ class HDSpecialist:
             self.runner.vae.to("cpu"); gc.collect(); torch.cuda.empty_cache()
             self.runner.dit.to(self.device)
 
-            text_pos_embeds = torch.load('./ckpts/pos_emb.pt').to(self.device)
-            text_neg_embeds = torch.load('./ckpts/neg_emb.pt').to(self.device)
+            pos_emb_path = SEEDVR_REPO_DIR / 'ckpts' / 'pos_emb.pt'
+            neg_emb_path = SEEDVR_REPO_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]}
 
             noises = [torch.randn_like(latent) for latent in cond_latents]
@@ -176,7 +219,7 @@ class HDSpecialist:
             final_sample = samples[0]
             input_video_sample = input_videos[0]
 
-            if final_sample.shape[1] < input_video_sample.shape[1]: # if generated frames are less
+            if final_sample.shape[1] < input_video_sample.shape[1]:
                 input_video_sample = input_video_sample[:, :final_sample.shape[1]]
 
             final_sample = wavelet_reconstruction(