File size: 5,380 Bytes
478560e 5137a03 3dce029 5137a03 3dce029 5137a03 3dce029 478560e 3dce029 5137a03 478560e 5137a03 3dce029 5137a03 3dce029 5137a03 3dce029 5137a03 478560e 3dce029 478560e 3dce029 5137a03 3dce029 478560e 3dce029 5137a03 3dce029 478560e 3dce029 478560e 3dce029 5137a03 3dce029 478560e 3dce029 478560e 3dce029 5137a03 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
# hd_specialist.py (Versão Final - Estrutura de Arquivos Corrigida)
#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
# --- 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.
"""
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...")
config_path = os.path.join('./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)
# ... (O resto da função process_video permanece exatamente o mesmo da resposta anterior) ...
finally:
self._unload_runner()
# Instância Singleton
hd_specialist_singleton = HDSpecialist() |