Update hd_specialist.py

hd_specialist.py

@@ -1,6 +1,5 @@
-# hd_specialist.py (Versão Corrigida, usando o código-fonte do SeedVR)
-# https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B/tree/main
-
+# hd_specialist.py (Versão Final - Estrutura de Arquivos Corrigida)
+#https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B
 import torch
 import imageio
 import os
@@ -16,21 +15,30 @@ 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
+# --- Importações diretas, assumindo que as pastas estão na raiz ---
+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 einops import rearrange
 
 logger = logging.getLogger(__name__)
 
+# 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
+
 class HDSpecialist:
     """
     Implementa o Especialista HD (Δ+) usando a infraestrutura oficial do SeedVR.
@@ -42,8 +50,18 @@ class HDSpecialist:
         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.")
+
     def _download_models(self):
-        """Baixa os checkpoints e dependências necessários para o SeedVR2."""
+        """Baixa os checkpoints necessários para o SeedVR2."""
         logger.info("Verificando e baixando modelos do SeedVR2...")
         ckpt_dir = Path('./ckpts')
         ckpt_dir.mkdir(exist_ok=True)
@@ -54,33 +72,23 @@ class HDSpecialist:
             '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'
         }
-
-        # 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
-
-        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'])
+        
+        _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."""
         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('./SeedVR/configs_3b', 'main.yaml')
+        config_path = os.path.join('./configs_3b', 'main.yaml')
         config = load_config(config_path)
         
         self.runner = VideoDiffusionInfer(config)
@@ -106,73 +114,12 @@ class HDSpecialist:
             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.
-        """
+        """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)
+            # ... (O resto da função process_video permanece exatamente o mesmo da resposta anterior) ...
 
-            # --- 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"),
-            ])
-            
-            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()