Update managers/audio_specialist.py

managers/audio_specialist.py

@@ -1,163 +1,242 @@
-# audio_specialist.py
-# Especialista ADUC para geração de áudio, com gerenciamento de memória GPU.
-# Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
+# hd_specialist.py
+#
+# 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 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 logging
-import subprocess
 import os
-import time
-import yaml
 import gc
+import logging
+import sys
+import subprocess
 from pathlib import Path
+from urllib.parse import urlparse
+from torch.hub import download_url_to_file
 import gradio as gr
-
-# Importa as classes e funções necessárias do MMAudio
-try:
-    from mmaudio.eval_utils import ModelConfig, all_model_cfg, generate as mmaudio_generate, load_video, make_video
-    from mmaudio.model.flow_matching import FlowMatching
-    from mmaudio.model.networks import MMAudio, get_my_mmaudio
-    from mmaudio.model.utils.features_utils import FeaturesUtils
-    from mmaudio.model.sequence_config import SequenceConfig
-except ImportError:
-    raise ImportError("MMAudio não foi encontrado. Por favor, instale-o a partir do GitHub: git+https://github.com/hkchengrex/MMAudio.git")
+import mediapy
+from einops import rearrange
 
 logger = logging.getLogger(__name__)
 
-class AudioSpecialist:
+# --- 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)
+    filename = file_name or os.path.basename(urlparse(url).path)
+    cached_file = os.path.abspath(os.path.join(model_dir, filename))
+    if not os.path.exists(cached_file):
+        logger.info(f'Downloading: "{url}" to {cached_file}')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=True)
+    return cached_file
+
+class HDSpecialist:
     """
-    Especialista responsável por gerar áudio para fragmentos de vídeo.
-    Gerencia o carregamento e descarregamento de modelos de áudio da VRAM.
+    Implements the HD Specialist (Δ+) using the SeedVR infrastructure.
+    Manages model loading, inference, and memory on demand.
     """
-    def __init__(self, workspace_dir):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.cpu_device = torch.device("cpu")
-        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32
+    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.")
         
-        self.model_config: ModelConfig = all_model_cfg['large_44k_v2']
-        self.net: MMAudio = None
-        self.feature_utils: FeaturesUtils = None
-        self.seq_cfg: SequenceConfig = None
+        # 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._load_models_to_cpu()
+        self._seedvr_modules_loaded = True
+        logger.info("SeedVR modules have been dynamically loaded.")
 
-    def _load_models_to_cpu(self):
-        """Carrega os modelos MMAudio para a memória da CPU na inicialização."""
-        try:
-            logger.info("Verificando e baixando modelos MMAudio, se necessário...")
-            self.model_config.download_if_needed()
-            
-            self.seq_cfg = self.model_config.seq_cfg
-            
-            logger.info(f"Carregando modelo MMAudio: {self.model_config.model_name} para a CPU...")
-            self.net = get_my_mmaudio(self.model_config.model_name).eval()
-            self.net.load_weights(torch.load(self.model_config.model_path, map_location=self.cpu_device, weights_only=True))
+    def _download_models(self):
+        """Downloads the necessary checkpoints for SeedVR2."""
+        logger.info("Verifying and downloading SeedVR2 models...")
+        ckpt_dir = SEEDVR_REPO_DIR / 'ckpts'
+        ckpt_dir.mkdir(exist_ok=True)
+
+        pretrain_model_urls = {
+            'vae': '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',
+            '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("Carregando utilitários de features do MMAudio para a CPU...")
-            self.feature_utils = FeaturesUtils(
-                tod_vae_ckpt=self.model_config.vae_path,
-                synchformer_ckpt=self.model_config.synchformer_ckpt,
-                enable_conditions=True,
-                mode=self.model_config.mode,
-                bigvgan_vocoder_ckpt=self.model_config.bigvgan_16k_path,
-                need_vae_encoder=False
-            )
-            self.feature_utils = self.feature_utils.eval()
-            self.net.to(self.cpu_device)
-            self.feature_utils.to(self.cpu_device)
-            logger.info("Especialista de áudio pronto na CPU.")
-        except Exception as e:
-            logger.error(f"Falha ao carregar modelos de áudio: {e}", exc_info=True)
-            self.net = None
-
-    def to_gpu(self):
-        """Move os modelos e utilitários para a GPU antes da inferência."""
-        if self.device == 'cpu': return
-        logger.info(f"Movendo especialista de áudio para a GPU ({self.device})...")
-        self.net.to(self.device, self.dtype)
-        self.feature_utils.to(self.device, self.dtype)
-
-    def to_cpu(self):
-        """Move os modelos de volta para a CPU e limpa a VRAM após a inferência."""
-        if self.device == 'cpu': return
-        logger.info("Descarregando especialista de áudio da GPU...")
-        self.net.to(self.cpu_device)
-        self.feature_utils.to(self.cpu_device)
-        gc.collect()
-        if torch.cuda.is_available(): torch.cuda.empty_cache()
-
-    def generate_audio_for_video(self, video_path: str, prompt: str, duration_seconds: float, output_path_override: str = None) -> str:
-        """
-        Gera áudio para um arquivo de vídeo, aplicando um prompt negativo para evitar fala.
+        logger.info("SeedVR2 models downloaded successfully.")
+
+    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()
+        self._download_models()
 
-        Args:
-            video_path (str): Caminho para o vídeo silencioso.
-            prompt (str): Descrição da cena para guiar a geração de SFX.
-            duration_seconds (float): Duração do áudio a ser gerado.
+        logger.info(f"Initializing SeedVR2 {model_version} runner...")
+        if model_version == '3B':
+            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 = 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}")
 
-        Returns:
-            str: Caminho para o novo arquivo de vídeo com áudio.
-        """
-        if self.net is None:
-            raise gr.Error("Modelo MMAudio não está carregado. Não é possível gerar áudio.")
-
-        logger.info("------------------------------------------------------")
-        logger.info("--- Gerando Áudio para Fragmento de Vídeo ---")
-        logger.info(f"--- Vídeo Fragmento: {os.path.basename(video_path)}")
-        logger.info(f"--- Duração: {duration_seconds:.2f}s")
-        logger.info(f"--- Prompt (Descrição da Cena): '{prompt}'")
+        config = load_config(str(config_path))
+        
+        self.runner = VideoDiffusionInfer(config)
+        OmegaConf.set_readonly(self.runner.config, False)
         
-        negative_prompt = "human voice"
-        logger.info(f"--- Negative Prompt: '{negative_prompt}'")
+        self.runner.configure_dit_model(device=self.device, checkpoint=str(checkpoint_path))
+        self.runner.configure_vae_model()
         
-        if duration_seconds < 1:
-            logger.warning("Fragmento muito curto (<1s). Retornando vídeo silencioso.")
-            logger.info("------------------------------------------------------")
-            return video_path
+        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.")
 
-        if self.device == 'cpu':
-            logger.warning("Gerando áudio na CPU. Isso pode ser muito lento.")
+    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()
+            self.is_initialized = False
+            logger.info("SeedVR2 runner unloaded from VRAM.")
 
+    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."""
         try:
-            self.to_gpu()
-            with torch.no_grad():
-                rng = torch.Generator(device=self.device).manual_seed(int(time.time()))
-                fm = FlowMatching(min_sigma=0, inference_mode='euler', num_steps=25)
-                
-                video_info = load_video(Path(video_path), duration_seconds)
-                self.seq_cfg.duration = video_info.duration_sec
-                self.net.update_seq_lengths(self.seq_cfg.latent_seq_len, self.seq_cfg.clip_seq_len, self.seq_cfg.sync_seq_len)
-                
-                audios = mmaudio_generate(
-                    clip_video=video_info.clip_frames.unsqueeze(0),
-                    sync_video=video_info.sync_frames.unsqueeze(0),
-                    text=[prompt],
-                    negative_text=[negative_prompt],
-                    feature_utils=self.feature_utils,
-                    net=self.net,
-                    fm=fm,
-                    rng=rng,
-                    cfg_strength=4.5
+            self._initialize_runner(model_version)
+            set_seed(seed, same_across_ranks=True)
+
+            self.runner.config.diffusion.timesteps.sampling.steps = steps
+            self.runner.configure_diffusion()
+
+            video_tensor = read_video(input_video_path, output_format="TCHW")[0] / 255.0
+            res_h, res_w = video_tensor.shape[-2:]
+            
+            video_transform = Compose([
+                NaResize(resolution=(res_h * res_w) ** 0.5, mode="area", downsample_only=False),
+                Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
+                DivisibleCrop((16, 16)),
+                Normalize(0.5, 0.5),
+                Rearrange("t c h w -> c t h w"),
+            ])
+            
+            cond_latents = [video_transform(video_tensor.to(self.device))]
+            input_videos = cond_latents
+
+            self.runner.dit.to("cpu")
+            self.runner.vae.to(self.device)
+            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'
+            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]
+            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,
                 )
-                audio_waveform = audios.float().cpu()[0]
-                
-                fragment_name = Path(video_path).stem
-                output_video_path = output_path_override if output_path_override else os.path.join(self.workspace_dir, f"{fragment_name}_com_audio.mp4")
-
-                make_video(video_info, Path(output_video_path), audio_waveform, sampling_rate=self.seq_cfg.sampling_rate)
-                logger.info(f"--- Fragmento com áudio salvo em: {os.path.basename(output_video_path)}")
-                logger.info("------------------------------------------------------")
-                return output_video_path
+            
+            self.runner.dit.to("cpu"); gc.collect(); torch.cuda.empty_cache()
+
+            self.runner.vae.to(self.device)
+            samples = self.runner.vae_decode(video_tensors)
+            
+            final_sample = samples[0]
+            input_video_sample = input_videos[0]
+
+            if final_sample.shape[1] < input_video_sample.shape[1]:
+                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(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()
+
+            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.to_cpu()
-
-# Singleton instantiation
-try:
-    with open("config.yaml", 'r') as f:
-        config = yaml.safe_load(f)
-    WORKSPACE_DIR = config['application']['workspace_dir']
-    audio_specialist_singleton = AudioSpecialist(workspace_dir=WORKSPACE_DIR)
-except Exception as e:
-    logger.error(f"Não foi possível inicializar o AudioSpecialist: {e}", exc_info=True)
-    audio_specialist_singleton = None
+            self._unload_runner()
+
+# Singleton instance
+hd_specialist_singleton = HDSpecialist()