Update hd_specialist.py

hd_specialist.py

@@ -1,97 +1,180 @@
-# hd_specialist.py
+# hd_specialist.py (Versão Corrigida, usando o código-fonte do SeedVR)
+# https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B/tree/main
+
 import torch
 import imageio
 import os
+import gc
 import logging
+import numpy as np
 from PIL import Image
-from diffusers import SeedVR2Pipeline
 from tqdm import tqdm
+import shlex
+import subprocess
+from pathlib import Path
+from urllib.parse import urlparse
+from torch.hub import download_url_to_file, get_dir
+from omegaconf import OmegaConf
+
+# --- Importações do código-fonte do SeedVR ---
+# Certifique-se de que a pasta 'SeedVR' está no seu projeto
+from SeedVR.projects.video_diffusion_sr.infer import VideoDiffusionInfer
+from SeedVR.common.config import load_config
+from SeedVR.common.seed import set_seed
+from SeedVR.data.image.transforms.divisible_crop import DivisibleCrop
+from SeedVR.data.image.transforms.na_resize import NaResize
+from SeedVR.data.video.transforms.rearrange import Rearrange
+from SeedVR.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__)
 
 class HDSpecialist:
     """
-    Implementa o Especialista HD (Δ+) da arquitetura ADUC-SDR.
-    Utiliza o modelo SeedVR2 para realizar a restauração e o aprimoramento
-    de vídeo, atuando como a etapa final de pós-produção.
+    Implementa o Especialista HD (Δ+) usando a infraestrutura oficial do SeedVR.
     """
-    def __init__(self):
+    def __init__(self, workspace_dir="deformes_workspace"):
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.pipeline = None
-        self.model_id = "ByteDance-Seed/SeedVR2-3B"
-        logger.info(f"Especialista HD inicializado. Aguardando carregamento do modelo {self.model_id}...")
-
-    def _load_pipeline(self):
-        """Carrega o pipeline do modelo sob demanda para economizar memória."""
-        if self.pipeline is None:
-            logger.info("Carregando o pipeline SeedVR2... Isso pode levar alguns minutos.")
-            try:
-                self.pipeline = SeedVR2Pipeline.from_pretrained(
-                    self.model_id,
-                    torch_dtype=torch.float16,
-                    variant="fp16"
-                ).to(self.device)
-                logger.info("Pipeline SeedVR2 carregado com sucesso.")
-            except Exception as e:
-                logger.error(f"Falha ao carregar o pipeline SeedVR2: {e}")
-                raise
-
-    def process_video(self, input_video_path: str, output_video_path: str, prompt: str, strength: float = 0.8, batch_size: int = 8) -> str:
-        """
-        Aplica o aprimoramento HD a um vídeo.
+        self.runner = None
+        self.workspace_dir = workspace_dir
+        self.is_initialized = False
+        logger.info("Especialista HD (SeedVR) inicializado. Modelo será carregado sob demanda.")
 
-        Args:
-            input_video_path (str): Caminho para o vídeo de entrada.
-            output_video_path (str): Caminho para salvar o vídeo aprimorado.
-            prompt (str): Um prompt para guiar a restauração (pode ser o prompt global).
-            strength (float): Força do efeito de restauração.
-            batch_size (int): Número de frames para processar por lote para gerenciar a VRAM.
+    def _download_models(self):
+        """Baixa os checkpoints e dependências necessários para o SeedVR2."""
+        logger.info("Verificando e baixando modelos do SeedVR2...")
+        ckpt_dir = Path('./ckpts')
+        ckpt_dir.mkdir(exist_ok=True)
 
-        Returns:
-            str: O caminho para o vídeo aprimorado.
-        """
-        self._load_pipeline()
-        if not os.path.exists(input_video_path):
-            logger.error(f"Vídeo de entrada não encontrado em: {input_video_path}")
-            raise FileNotFoundError(f"Vídeo de entrada não encontrado: {input_video_path}")
+        pretrain_model_url = {
+            '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',
+            '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'
+        }
 
-        logger.info(f"Iniciando processo HD para: {input_video_path}")
+        # Função auxiliar para download
+        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)
+            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}')
+                download_url_to_file(url, cached_file, hash_prefix=None, progress=True)
+            return cached_file
 
-        # 1. Ler os frames do vídeo de entrada
-        reader = imageio.get_reader(input_video_path)
-        fps = reader.get_meta_data()['fps']
-        input_frames = [Image.fromarray(frame) for frame in reader]
-        reader.close()
+        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.")
 
-        logger.info(f"Vídeo lido. Total de {len(input_frames)} frames a {fps} FPS.")
 
-        # 2. Processar frames em lotes
-        processed_frames = []
+    def _initialize_runner(self):
+        """Carrega e configura o modelo SeedVR sob demanda."""
+        if self.runner is not None:
+            return
+
+        self._download_models()
+
+        logger.info("Inicializando o runner do SeedVR2...")
+        config_path = os.path.join('./SeedVR/configs_3b', 'main.yaml')
+        config = load_config(config_path)
+        
+        self.runner = VideoDiffusionInfer(config)
+        OmegaConf.set_readonly(self.runner.config, False)
         
-        # Usamos tqdm para ter uma barra de progresso no console
-        for i in tqdm(range(0, len(input_frames), batch_size), desc="Aprimorando frames com SeedVR2"):
-            batch = input_frames[i:i + batch_size]
+        self.runner.configure_dit_model(device=self.device, checkpoint='./ckpts/seedvr2_ema_3b.pth')
+        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.")
+
+    def _unload_runner(self):
+        """Remove o runner da VRAM para liberar recursos."""
+        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.")
+
+    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.
+        """
+        try:
+            self._initialize_runner()
+            set_seed(seed, same_across_ranks=True)
+
+            # --- Configuração do Pipeline (adaptado de app.py) ---
+            self.runner.config.diffusion.cfg.scale = 1.0 # cfg_scale
+            self.runner.config.diffusion.cfg.rescale = 0.0 # cfg_rescale
+            self.runner.config.diffusion.timesteps.sampling.steps = 1 # sample_steps (one-step model)
+            self.runner.configure_diffusion()
+
+            # --- Preparação do Vídeo de Entrada ---
+            logger.info(f"Processando vídeo de entrada: {input_video_path}")
+            video_tensor = read_video(input_video_path, output_format="TCHW")[0] / 255.0
+            if video_tensor.size(0) > 121:
+                logger.warning(f"Vídeo com {video_tensor.size(0)} frames. Truncando para 121 frames.")
+                video_tensor = video_tensor[:121]
+            
+            video_transform = Compose([
+                NaResize(resolution=(1280 * 720)**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"),
+            ])
             
-            # O pipeline retorna uma lista de imagens PIL
-            enhanced_batch = self.pipeline(
-                prompt=prompt,
-                frames=batch,
-                num_inference_steps=20,
-                guidance_scale=7.5,
-                strength=strength,
-            ).frames
-
-            processed_frames.extend(enhanced_batch)
-
-        # 3. Salvar os frames aprimorados em um novo vídeo
-        logger.info(f"Salvando vídeo aprimorado em: {output_video_path}")
-        writer = imageio.get_writer(output_video_path, fps=fps, codec='libx264', quality=9)
-        for frame in processed_frames:
-            writer.append_data(np.array(frame))
-        writer.close()
-
-        logger.info("Processo HD concluído com sucesso.")
-        return output_video_path
-
-# Instância Singleton para ser usada em toda a aplicação
+            cond_latent = video_transform(video_tensor.to(self.device))
+            input_video_for_colorfix = cond_latent.clone() # Salva para o color fix
+            ori_length = cond_latent.size(1)
+
+            # --- Codificação VAE e Geração ---
+            logger.info("Codificando vídeo para o espaço latente...")
+            cond_latent = self.runner.vae_encode([cond_latent])[0]
+
+            text_pos_embeds = torch.load('pos_emb.pt').to(self.device)
+            text_neg_embeds = torch.load('neg_emb.pt').to(self.device)
+            text_embeds_dict = {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]}
+
+            noise = torch.randn_like(cond_latent)
+            
+            logger.info(f"Iniciando a geração de restauração para {ori_length} frames...")
+            with torch.no_grad(), torch.autocast("cuda", torch.bfloat16, enabled=True):
+                video_tensor_out = self.runner.inference(
+                    noises=[noise],
+                    conditions=[self.runner.get_condition(noise, task="sr", latent_blur=cond_latent)],
+                    dit_offload=False,
+                    **text_embeds_dict,
+                )[0]
+
+            sample = rearrange(video_tensor_out, "c t h w -> t c h w")
+            
+            # --- Pós-processamento e Salvamento ---
+            if ori_length < sample.shape[0]:
+                sample = sample[:ori_length]
+            
+            input_video_for_colorfix = rearrange(input_video_for_colorfix, "c t h w -> t c h w")
+            sample = wavelet_reconstruction(sample.cpu(), input_video_for_colorfix[:sample.size(0)].cpu())
+            
+            sample = rearrange(sample, "t c h w -> t h w c")
+            sample = sample.clip(-1, 1).mul_(0.5).add_(0.5).mul_(255).round().to(torch.uint8).numpy()
+
+            logger.info(f"Salvando vídeo aprimorado em: {output_video_path}")
+            imageio.get_writer(output_video_path, fps=fps_out, codec='libx264', quality=9).extend(sample)
+
+            return output_video_path
+        finally:
+            self._unload_runner()
+
+# Instância Singleton
 hd_specialist_singleton = HDSpecialist()