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 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 (Δ+) 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 = workspace_dir
        self.is_initialized = False
        logger.info("Especialista HD (SeedVR) inicializado. Modelo será carregado sob demanda.")

    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)

        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'
        }

        # 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'])
        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._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)
        
        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"),
            ])
            
            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()