Update hd_specialist.py

hd_specialist.py

@@ -1,5 +1,13 @@
-# hd_specialist.py (Versão Final - Estrutura de Arquivos Corrigida)
-#https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B
+# hd_specialist.py
+#
+# Copyright (C) 2025 Carlos Rodrigues dos Santos
+#
+# 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).
+
 import torch
 import imageio
 import os
@@ -12,10 +20,12 @@ import shlex
 import subprocess
 from pathlib import Path
 from urllib.parse import urlparse
-from torch.hub import download_url_to_file, get_dir
+from torch.hub import download_url_to_file
 from omegaconf import OmegaConf
+import mediapy
+from einops import rearrange
 
-# --- Importações diretas, assumindo que as pastas estão na raiz ---
+# 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
@@ -25,103 +35,164 @@ 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 einops import rearrange
 
 logger = logging.getLogger(__name__)
 
-# Função auxiliar para download
 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'Baixando: "{url}" para {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:
     """
-    Implementa o Especialista HD (Δ+) usando a infraestrutura oficial do SeedVR.
+    Implements the HD Specialist (Δ+) using the SeedVR infrastructure.
+    Manages model loading, inference, and memory on demand.
     """
     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
-        logger.info("Especialista HD (SeedVR) inicializado. Modelo será carregado sob demanda.")
-
-    def _setup_dependencies(self):
-        """Instala dependências complexas como Apex."""
-        logger.info("Configurando dependências do SeedVR (Apex)...")
-        apex_url = 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x86_64.whl'
-        apex_wheel_path = _load_file_from_url(url=apex_url)
-        
-        # Instala a roda do Apex baixada
-        subprocess.run(shlex.split(f"pip install {apex_wheel_path}"), check=True)
-        logger.info("✅ Dependência Apex instalada com sucesso.")
+        logger.info("HD Specialist (SeedVR) initialized. Model will be loaded on demand.")
 
     def _download_models(self):
-        """Baixa os checkpoints necessários para o SeedVR2."""
-        logger.info("Verificando e baixando modelos do SeedVR2...")
+        """Downloads the necessary checkpoints for SeedVR2."""
+        logger.info("Verifying and downloading SeedVR2 models...")
         ckpt_dir = Path('./ckpts')
         ckpt_dir.mkdir(exist_ok=True)
 
-        pretrain_model_url = {
+        pretrain_model_urls = {
             'vae': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/ema_vae.pth',
-            'dit': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/seedvr2_ema_3b.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'
         }
         
-        _load_file_from_url(url=pretrain_model_url['dit'], model_dir='./ckpts/')
-        _load_file_from_url(url=pretrain_model_url['vae'], model_dir='./ckpts/')
-        _load_file_from_url(url=pretrain_model_url['pos_emb'])
-        _load_file_from_url(url=pretrain_model_url['neg_emb'])
-        logger.info("Modelos do SeedVR2 baixados com sucesso.")
-
-    def _initialize_runner(self):
-        """Carrega e configura o modelo SeedVR sob demanda."""
+        for key, url in pretrain_model_urls.items():
+            _load_file_from_url(url=url, model_dir='./ckpts/')
+            
+        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._setup_dependencies()
         self._download_models()
 
-        logger.info("Inicializando o runner do SeedVR2...")
-        config_path = os.path.join('./configs_3b', 'main.yaml')
+        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'
+        elif model_version == '7B':
+            config_path = os.path.join('./configs_7b', 'main.yaml')
+            checkpoint_path = './ckpts/seedvr2_ema_7b.pth'
+        else:
+            raise ValueError(f"Unsupported SeedVR model version: {model_version}")
+
         config = load_config(config_path)
         
         self.runner = VideoDiffusionInfer(config)
         OmegaConf.set_readonly(self.runner.config, False)
         
-        self.runner.configure_dit_model(device=self.device, checkpoint='./ckpts/seedvr2_ema_3b.pth')
+        self.runner.configure_dit_model(device=self.device, checkpoint=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
-        logger.info("Runner do SeedVR2 inicializado e pronto.")
+        logger.info(f"Runner for SeedVR2 {model_version} initialized and ready.")
 
     def _unload_runner(self):
-        """Remove o runner da VRAM para liberar recursos."""
+        """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("Runner do SeedVR2 descarregado da VRAM.")
+            logger.info("SeedVR2 runner unloaded from VRAM.")
 
-    def process_video(self, input_video_path: str, output_video_path: str, prompt: str, seed: int = 666, fps_out: int = 24) -> str:
-        """Aplica o aprimoramento HD a um vídeo usando a lógica oficial do SeedVR."""
+    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._initialize_runner()
+            self._initialize_runner(model_version)
             set_seed(seed, same_across_ranks=True)
-            # ... (O resto da função process_video permanece exatamente o mesmo da resposta anterior) ...
+
+            # --- Adapted inference logic from SeedVR scripts ---
+            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)
+
+            text_pos_embeds = torch.load('./ckpts/pos_emb.pt').to(self.device)
+            text_neg_embeds = torch.load('./ckpts/neg_emb.pt').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,
+                )
+            
+            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]: # if generated frames are less
+                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._unload_runner()
 
-# Instância Singleton
+# Singleton instance
 hd_specialist_singleton = HDSpecialist()