aduc_framework/managers/seedvr_manager.py

# hd_specialist.py (Versão Final - Estrutura de Arquivos Corrigida e Autossuficiente)
# https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B

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
import sys


logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

APP_ROOT = Path("/home/user/app")
DEPS_DIR = APP_ROOT / "deps"
SEEDVR_SPACE_DIR = DEPS_DIR / "SeedVR_Space"
SEEDVR_SPACE_URL = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B"


def setup_dependencies():
    """
    Ensures the SEEDVR repository is cloned and available in the sys.path.
    This function is run once when the module is first imported.
    """
    if not SEEDVR_SPACE_DIR.exists():
        logger.info(f"SEEDVR repository not found at '{SEEDVR_SPACE_DIR}'. Cloning from GitHub...")
        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("SEEDVR repository cloned successfully.")
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to clone SEEDVR repository. Git stderr: {e.stderr}")
            raise RuntimeError("Could not clone the required SEEDVR dependency from GitHub.")
    else:
        logger.info("Found local SEEDVR repository.")
    
    if str(SEEDVR_SPACE_DIR.resolve()) not in sys.path:
        sys.path.insert(0, str(SEEDVR_SPACE_DIR.resolve()))
        logger.info(f"Added '{SEEDVR_SPACE_DIR.resolve()}' to sys.path.")

setup_dependencies()



# 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



# --- 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

    

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 = 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.")

    def _download_models(self):
        """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)

        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='./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...")
        
        # --- NOVO: Verificação e download automático do main.yaml ---
        config_dir = Path('./configs_3b')
        config_path = config_dir / 'main.yaml'
        config_url = "https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B/resolve/main/configs_3b/main.yaml"
        
        if not config_path.is_file():
            logger.warning(f"Arquivo de configuração '{config_path}' não encontrado.")
            logger.info("Tentando baixar automaticamente...")
            try:
                # Reutiliza a função de download existente
                _load_file_from_url(url=config_url, model_dir=str(config_dir), file_name='main.yaml')
                logger.info(f"✅ Arquivo de configuração baixado com sucesso para '{config_path}'")
            except Exception as e:
                logger.error(f"Falha ao baixar o arquivo de configuração. Por favor, baixe-o manualmente de {config_url} e coloque-o em {config_dir}. Erro: {e}")
                raise
        else:
            logger.info(f"Arquivo de configuração encontrado em '{config_path}'.")
        # --- FIM DA MODIFICAÇÃO ---

        config = load_config(str(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) -> 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 para facilitar o uso (opcional)
seedvr_manager_singleton = SeedVrManager()