Delete aduc_framework/managers/seedvr_manager.py

aduc_framework/managers/seedvr_manager.py

@@ -1,243 +0,0 @@
-# hd_specialist.py (Versão Final - Corrigindo o Contexto de Execução de Caminhos)
-# https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B
-
-import torch
-import imageio
-import os
-import gc
-import logging
-import numpy as np
-import shlex
-import subprocess
-from pathlib import Path
-from urllib.parse import urlparse
-from torch.hub import download_url_to_file
-from omegaconf import OmegaConf
-import sys
-
-# --- Configuração do Logging ---
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
-
-# --- Constantes de Caminho ---
-# Define a raiz do projeto (onde este script está) e cria um diretório para dependências
-# Usamos 'Path.cwd()' para ter certeza que é relativo ao diretório de execução do projeto principal.
-PROJECT_ROOT = Path.cwd() 
-DEPS_DIR = PROJECT_ROOT / "deps"
-SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
-SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"
-
-def setup_environment():
-    """
-    Clona o repositório SeedVR se não existir e o adiciona ao sys.path
-    para que seus módulos (common, projects, etc.) possam ser importados.
-    """
-    if not SEEDVR_SPACE_DIR.is_dir():
-        logger.info(f"Repositório SeedVR não encontrado. Clonando de '{SEEDVR_SPACE_URL}'...")
-        try:
-            DEPS_DIR.mkdir(exist_ok=True)
-            subprocess.run(
-                ["git", "clone", "--depth", "1", SEEDVR_SPACE_URL, str(SEEDVR_SPACE_DIR)],
-                check=True, capture_output=True, text=True
-            )
-            logger.info(f"✅ Repositório clonado com sucesso em '{SEEDVR_SPACE_DIR}'")
-        except subprocess.CalledProcessError as e:
-            logger.error(f"❌ Falha ao clonar o repositório. Erro do Git: {e.stderr}")
-            raise RuntimeError("Não foi possível clonar a dependência SeedVR do Hugging Face.")
-    else:
-        logger.info(f"Repositório SeedVR já existe em '{SEEDVR_SPACE_DIR}'.")
-
-    resolved_path = str(SEEDVR_SPACE_DIR.resolve())
-    if resolved_path not in sys.path:
-        sys.path.insert(0, resolved_path)
-        logger.info(f"Adicionado '{resolved_path}' ao sys.path.")
-
-# Executa a configuração do ambiente assim que o módulo é carregado
-setup_environment()
-
-# Função auxiliar de download (permanece a mesma)
-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
-
-# --- Importações do Repositório Clonado ---
-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
-
-class SeedVrManager:
-    """
-    Implementa o Especialista HD (Δ+) usando a infraestrutura oficial do SeedVR.
-    """
-    def __init__(self, workspace_dir="deformes_workspace"):
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.runner = None
-        self.workspace_dir = Path(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, model_dir=str(DEPS_DIR))
-        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 necessários para o SeedVR2 DENTRO do repositório clonado."""
-        logger.info("Verificando e baixando modelos do SeedVR2...")
-        ckpt_dir = SEEDVR_SPACE_DIR / 'ckpts'
-        
-        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'
-        }
-        
-        _load_file_from_url(url=pretrain_model_url['dit'], model_dir=str(ckpt_dir))
-        _load_file_from_url(url=pretrain_model_url['vae'], model_dir=str(ckpt_dir))
-        _load_file_from_url(url=pretrain_model_url['pos_emb'], model_dir=str(SEEDVR_SPACE_DIR))
-        _load_file_from_url(url=pretrain_model_url['neg_emb'], model_dir=str(SEEDVR_SPACE_DIR))
-        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...")
-        
-        # --- CORREÇÃO CRÍTICA: MUDANÇA DE DIRETÓRIO DE TRABALHO ---
-        original_cwd = Path.cwd()
-        try:
-            # Muda para o diretório do repositório clonado. Isso é essencial para que
-            # o `load_config` encontre os arquivos .yaml herdados (como os da pasta `models`).
-            os.chdir(SEEDVR_SPACE_DIR)
-            
-            logger.info(f"Diretório de trabalho alterado para: {SEEDVR_SPACE_DIR}")
-
-            # Agora todos os caminhos são relativos à raiz do repositório
-            config_path = './configs_3b/main.yaml'
-            dit_checkpoint_path = './ckpts/seedvr2_ema_3b.pth'
-            
-            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=dit_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.")
-
-        finally:
-            # Garante que o diretório de trabalho original seja restaurado,
-            # não importa se a inicialização foi bem-sucedida ou falhou.
-            os.chdir(original_cwd)
-            logger.info(f"Diretório de trabalho restaurado para: {original_cwd}")
-        # --- FIM DA CORREÇÃO ---
-
-    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) -> str:
-        """Aplica o aprimoramento HD a um vídeo usando a lógica oficial do SeedVR."""
-        original_cwd = Path.cwd() # Salva o diretório original
-        try:
-            self._initialize_runner()
-            
-            # --- CORREÇÃO CRÍTICA 2: MUDAR DIRETÓRIO TAMBÉM DURANTE O PROCESSAMENTO ---
-            # O código pode precisar acessar arquivos de embedding novamente
-            os.chdir(SEEDVR_SPACE_DIR)
-
-            set_seed(seed, same_across_ranks=True)
-
-            self.runner.config.diffusion.cfg.scale = 1.0
-            self.runner.config.diffusion.cfg.rescale = 0.0
-            self.runner.config.diffusion.timesteps.sampling.steps = 1
-            self.runner.configure_diffusion()
-
-            logger.info(f"Processando vídeo de entrada: {input_video_path}")
-            # Garante que os caminhos de entrada/saída sejam absolutos para não quebrar com a mudança de CWD
-            abs_input_path = original_cwd / input_video_path
-            abs_output_path = original_cwd / output_video_path
-
-            video_tensor = read_video(str(abs_input_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()
-            ori_length = cond_latent.size(1)
-
-            logger.info("Codificando vídeo para o espaço latente...")
-            cond_latent = self.runner.vae_encode([cond_latent])[0]
-
-            # Carrega embeddings com caminhos relativos ao diretório do repo
-            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")
-            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: {abs_output_path}")
-            self.workspace_dir.mkdir(parents=True, exist_ok=True)
-            imageio.get_writer(str(abs_output_path), fps=fps_out, codec='libx264', quality=9).extend(sample)
-
-            return str(abs_output_path)
-        finally:
-            os.chdir(original_cwd) # Restaura o diretório de trabalho original
-            self._unload_runner()
-
-# Instância Singleton
-seedvr_manager_singleton = SeedVrManager()