# api/seedvr_server.py import os import sys import shutil import mimetypes import time import subprocess # Necessário para clonar o repositório na configuração inicial from pathlib import Path from typing import Optional, Callable from types import SimpleNamespace from huggingface_hub import hf_hub_download # Adiciona dinamicamente o caminho do repositório clonado ao sys.path. # Isso é crucial para que a importação do 'inference_cli' funcione. SEEDVR_REPO_PATH = Path(os.getenv("SEEDVR_ROOT", "/data/SeedVR")) if str(SEEDVR_REPO_PATH) not in sys.path: # Insere no início da lista para garantir prioridade de importação. sys.path.insert(0, str(SEEDVR_REPO_PATH)) # Tenta importar as funções necessárias APÓS a modificação do path. # Se falhar, a aplicação não pode continuar. #try: # from inference_cli import run_inference_logic, save_frames_to_video #except ImportError as e: # print(f"ERRO FATAL: Não foi possível importar de 'inference_cli.py'.") # print(f"Verifique se o repositório em '{SEEDVR_REPO_PATH}' está correto e completo.") # raise e #!/usr/bin/env python3 """ Standalone SeedVR2 Video Upscaler CLI Script (MODIFICADO PARA SUPORTE ROBUSTO A CALLBACKS EM MULTIPROCESSING) """ import sys import os import argparse import time import multiprocessing as mp import queue # Importa a classe de exceção para filas vazias # Garante o uso seguro de CUDA com multiprocessing, essencial para estabilidade. if mp.get_start_method(allow_none=True) != 'spawn': mp.set_start_method('spawn', force=True) # ------------------------------------------------------------- # 1) Configuração de alocação de memória da VRAM os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync") # 2) Pré-análise dos argumentos para configurar a visibilidade dos dispositivos CUDA _pre_parser = argparse.ArgumentParser(add_help=False) _pre_parser.add_argument("--cuda_device", type=str, default=None) _pre_args, _ = _pre_parser.parse_known_args() if _pre_args.cuda_device is not None: device_list_env = [x.strip() for x in _pre_args.cuda_device.split(',') if x.strip()!=''] if len(device_list_env) == 1: os.environ["CUDA_VISIBLE_DEVICES"] = device_list_env[0] # ------------------------------------------------------------- # 3) Importações pesadas (torch, etc.) são feitas após a configuração do ambiente. import torch import cv2 import numpy as np from datetime import datetime from pathlib import Path # Adiciona o diretório raiz do projeto ao path do sistema para permitir importações de `src` # Isso assume que o script está dentro do repositório clonado. script_dir = os.path.dirname(os.path.abspath(__file__)) if script_dir not in sys.path: sys.path.insert(0, script_dir) # Importa as funções do SeedVR DEPOIS de ajustar o path. from src.utils.downloads import download_weight def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load_cap=None): """ Extrai quadros de um vídeo e os converte para o formato de tensor. """ if debug: print(f"🎬 Extracting frames from video: {video_path}") if not os.path.exists(video_path): raise FileNotFoundError(f"Video file not found: {video_path}") cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError(f"Cannot open video file: {video_path}") fps = cap.get(cv2.CAP_PROP_FPS) frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) if debug: print(f"📊 Video info: {frame_count} frames, {width}x{height}, {fps:.2f} FPS") if skip_first_frames: print(f"⏭️ Will skip first {skip_first_frames} frames") if load_cap and load_cap > 0: print(f"🔢 Will load maximum {load_cap} frames") frames = [] frame_idx = 0 frames_loaded = 0 while True: ret, frame = cap.read() if not ret: break if frame_idx < skip_first_frames: frame_idx += 1 continue if load_cap is not None and load_cap > 0 and frames_loaded >= load_cap: break frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame = frame.astype(np.float32) / 255.0 frames.append(frame) frame_idx += 1 frames_loaded += 1 if debug and frames_loaded % 100 == 0: total_to_load = min(frame_count, load_cap) if (load_cap and load_cap > 0) else frame_count print(f"📹 Extracted {frames_loaded}/{total_to_load} frames") cap.release() if len(frames) == 0: raise ValueError(f"No frames extracted from video: {video_path}") if debug: print(f"✅ Extracted {len(frames)} frames") frames_tensor = torch.from_numpy(np.stack(frames)).to(torch.float16) if debug: print(f"📊 Frames tensor shape: {frames_tensor.shape}, dtype: {frames_tensor.dtype}") return frames_tensor, fps def save_frames_to_video(frames_tensor, output_path, fps=30.0, debug=False): """ Salva um tensor de quadros em um arquivo de vídeo. """ if debug: print(f"🎬 Saving {frames_tensor.shape[0]} frames to video: {output_path}") os.makedirs(os.path.dirname(output_path), exist_ok=True) frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8) T, H, W, C = frames_np.shape fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (W, H)) if not out.isOpened(): raise ValueError(f"Cannot create video writer for: {output_path}") for i, frame in enumerate(frames_np): frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) out.write(frame_bgr) if debug and (i + 1) % 100 == 0: print(f"💾 Saved {i + 1}/{T} frames") out.release() if debug: print(f"✅ Video saved successfully: {output_path}") def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, progress_queue=None): """ Processo filho (worker) que executa o upscaling em uma GPU dedicada. """ os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id) os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync") import torch from src.core.model_manager import configure_runner from src.core.generation import generation_loop frames_tensor = torch.from_numpy(frames_np).to(torch.float16) local_progress_callback = None if progress_queue: def callback_wrapper(batch_idx, total_batches, current_frames, message): # Envia uma tupla com informações de progresso para a fila progress_queue.put((proc_idx, batch_idx, total_batches, message)) local_progress_callback = callback_wrapper try: runner = configure_runner(shared_args["model"], shared_args["model_dir"], shared_args["preserve_vram"], shared_args["debug"]) result_tensor = generation_loop( runner=runner, images=frames_tensor, cfg_scale=shared_args["cfg_scale"], seed=shared_args["seed"], res_w=shared_args["res_w"], batch_size=shared_args["batch_size"], preserve_vram=shared_args["preserve_vram"], temporal_overlap=shared_args["temporal_overlap"], debug=shared_args["debug"], progress_callback=local_progress_callback ) return_queue.put((proc_idx, result_tensor.cpu().numpy())) except Exception as e: import traceback error_msg = f"ERROR in worker {proc_idx}: {e}\n{traceback.format_exc()}" print(error_msg) if progress_queue: progress_queue.put((proc_idx, -1, -1, error_msg)) return_queue.put((proc_idx, error_msg)) def _gpu_processing(frames_tensor, device_list, args, progress_callback=None): """ Divide os quadros, gerencia os workers e monitora o progresso de forma robusta. """ num_devices = len(device_list) chunks = torch.chunk(frames_tensor, num_devices, dim=0) manager = mp.Manager() return_queue = manager.Queue() progress_queue = manager.Queue() if progress_callback else None workers = [] shared_args = { "model": args.model, "model_dir": args.model_dir or "./models/SEEDVR2", "preserve_vram": args.preserve_vram, "debug": args.debug, "cfg_scale": 1.0, "seed": args.seed, "res_w": args.resolution, "batch_size": args.batch_size, "temporal_overlap": 0, } for idx, (device_id, chunk_tensor) in enumerate(zip(device_list, chunks)): p = mp.Process(target=_worker_process, args=(idx, device_id, chunk_tensor.cpu().numpy(), shared_args, return_queue, progress_queue)) p.start() workers.append(p) results_np = [None] * num_devices finished_workers_count = 0 worker_progress = [0.0] * num_devices while finished_workers_count < num_devices: if progress_queue: while not progress_queue.empty(): try: proc_idx, batch_idx, total_batches, message = progress_queue.get_nowait() if batch_idx == -1: raise RuntimeError(f"Worker {proc_idx} error: {message}") if total_batches > 0: worker_progress[proc_idx] = batch_idx / total_batches total_progress = sum(worker_progress) / num_devices progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: {message}") except queue.Empty: break try: proc_idx, result = return_queue.get(timeout=0.2) if isinstance(result, str) and result.startswith("ERROR"): raise RuntimeError(f"Worker {proc_idx} failed: {result}") results_np[proc_idx] = result worker_progress[proc_idx] = 1.0 finished_workers_count += 1 if progress_callback: total_progress = sum(worker_progress) / num_devices progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: Completed!") except queue.Empty: pass for p in workers: p.join() if any(r is None for r in results_np): raise RuntimeError("One or more workers failed to return a result.") return torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16) def parse_arguments(): """Analisa os argumentos da linha de comando.""" parser = argparse.ArgumentParser(description="SeedVR2 Video Upscaler CLI") parser.add_argument("--video_path", type=str, required=True, help="Path to input video file") parser.add_argument("--seed", type=int, default=100, help="Random seed for generation (default: 100)") parser.add_argument("--resolution", type=int, default=1072, help="Target resolution of the short side (default: 1072)") parser.add_argument("--batch_size", type=int, default=5, help="Number of frames per batch (default: 5)") parser.add_argument("--model", type=str, default="seedvr2_ema_3b_fp16.safetensors", choices=["seedvr2_ema_3b_fp16.safetensors", "seedvr2_ema_3b_fp8_e4m3fn.safetensors", "seedvr2_ema_7b_fp16.safetensors", "seedvr2_ema_7b_fp8_e4m3fn.safetensors"], help="Model to use") parser.add_argument("--model_dir", type=str, default=None, help="Directory containing the model files") parser.add_argument("--skip_first_frames", type=int, default=0, help="Skip the first frames during processing") parser.add_argument("--load_cap", type=int, default=0, help="Maximum number of frames to load from video (default: load all)") parser.add_argument("--output", type=str, default=None, help="Output path") parser.add_argument("--output_format", type=str, default="video", choices=["video", "png"], help="Output format: 'video' (mp4) or 'png' images") parser.add_argument("--preserve_vram", action="store_true", help="Enable VRAM preservation mode") parser.add_argument("--debug", action="store_true", help="Enable debug logging") parser.add_argument("--cuda_device", type=str, default=None, help="CUDA device id(s). e.g., '0' or '0,1' for multi-GPU") return parser.parse_args() def run_inference_logic(args, progress_callback=None): """ Função principal que executa o pipeline de upscaling. Pode ser importada e chamada por outros scripts. """ if args.debug: print(f"📋 Argumentos da Lógica de Inferência: {vars(args)}") if progress_callback: progress_callback(0.05, "Extracting frames...") print("🎬 Extraindo frames do vídeo...") start_time = time.time() frames_tensor, original_fps = extract_frames_from_video( args.video_path, args.debug, args.skip_first_frames, args.load_cap ) if args.debug: print(f"🔄 Tempo de extração de frames: {time.time() - start_time:.2f}s") device_list = [d.strip() for d in str(args.cuda_device).split(',') if d.strip()] if args.cuda_device else ["0"] if args.debug: print(f"🚀 Usando dispositivos: {device_list}") if progress_callback: progress_callback(0.1, "Starting generation...") processing_start = time.time() download_weight(args.model, args.model_dir or "seedvr_models") result_tensor = _gpu_processing(frames_tensor, device_list, args, progress_callback) generation_time = time.time() - processing_start if args.debug: print(f"🔄 Tempo de Geração: {generation_time:.2f}s") print(f"📊 Resultado: {result_tensor.shape}, dtype: {result_tensor.dtype}") return result_tensor, original_fps, generation_time, len(frames_tensor) def main(): """ Função principal para execução via linha de comando (CLI). """ print(f"🚀 SeedVR2 Video Upscaler CLI iniciado às {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") args = parse_arguments() try: result_tensor, original_fps, _, _ = run_inference_logic(args) if args.output is None: args.output = f"result_{Path(args.video_path).stem}.mp4" print(f"💾 Salvando vídeo em: {args.output}") save_frames_to_video(result_tensor, args.output, original_fps, args.debug) print("✅ Upscaling via CLI concluído com sucesso!") except Exception as e: print(f"❌ Erro durante o processamento via CLI: {e}") import traceback traceback.print_exc() sys.exit(1) if __name__ == "__main__": main() SeedVRServer.setup_dependencies() class SeedVRServer: INIT = False def __init__(self, **kwargs): """ Inicializa o servidor, define os caminhos e prepara o ambiente. """ self.SEEDVR_ROOT = SEEDVR_REPO_PATH self.CKPTS_ROOT = Path("/data/seedvr_models_fp16") self.OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/outputs")) self.INPUT_ROOT = Path(os.getenv("INPUT_ROOT", "/app/inputs")) self.HF_HOME_CACHE = Path(os.getenv("HF_HOME", "/data/.cache/huggingface")) self.REPO_URL = os.getenv("SEEDVR_GIT_URL", "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler") self.NUM_GPUS_TOTAL = int(os.getenv("NUM_GPUS", "4")) #print("🚀 SeedVRServer ja inicializando...") print("⚙️ SeedVRServer (Modo de Chamada Direta) inicializando...") for p in [self.CKPTS_ROOT, self.OUTPUT_ROOT, self.INPUT_ROOT, self.HF_HOME_CACHE]: p.mkdir(parents=True, exist_ok=True) self.setup_dependencies() print("📦 SeedVRServer (Modo de Chamada Direta) pronto.") INIT = True def setup_dependencies(self): """ Garante que o repositório e os modelos estão presentes. """ self._ensure_repo() self._ensure_model() def _ensure_repo(self) -> None: """ Clona o repositório do SeedVR se ele não existir. """ if not (self.SEEDVR_ROOT / ".git").exists(): print(f"[SeedVRServer] Clonando repositório para {self.SEEDVR_ROOT}...") # Usamos subprocess.run aqui porque é uma tarefa de inicialização única. subprocess.run(["git", "clone", "--depth", "1", self.REPO_URL, str(self.SEEDVR_ROOT)], check=True) else: print("[SeedVRServer] Repositório SeedVR já existe.") def _ensure_model(self) -> None: """ Baixa os checkpoints do Hugging Face se não existirem localmente. """ print(f"[SeedVRServer] Verificando checkpoints (FP16) em {self.CKPTS_ROOT}...") model_files = { "seedvr2_ema_3b_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses", "ema_vae_fp16.safetensors": "MonsterMMORPG/SeedVR2_SECourses", } for filename, repo_id in model_files.items(): if not (self.CKPTS_ROOT / filename).exists(): print(f"Baixando {filename} de {repo_id}...") hf_hub_download( repo_id=repo_id, filename=filename, local_dir=str(self.CKPTS_ROOT), cache_dir=str(self.HF_HOME_CACHE), token=os.getenv("HF_TOKEN") ) print("[SeedVRServer] Checkpoints (FP16) estão no local correto.") def run_inference_direct( self, file_path: str, *, seed: int, res_h: int, res_w: int, sp_size: int, fps: Optional[float] = None, progress: Optional[Callable] = None ) -> str: """ Executa a inferência diretamente no mesmo processo e retorna o caminho do arquivo de saída. """ # Cria um diretório de saída único para salvar o resultado. out_dir = self.OUTPUT_ROOT / f"run_{int(time.time())}_{Path(file_path).stem}" out_dir.mkdir(parents=True, exist_ok=True) output_filepath = out_dir / f"result_{Path(file_path).stem}.mp4" # Simula o objeto 'args' que a função de lógica do inference_cli espera. # Usamos SimpleNamespace para criar um objeto simples com atributos. args = SimpleNamespace( video_path=file_path, output=str(output_filepath), model_dir=str(self.CKPTS_ROOT), seed=seed, resolution=res_h, # O script do SeedVR usa a altura (lado menor) como referência. batch_size=4, model="seedvr2_ema_7b_sharp_fp16.safetensors", preserve_vram=True, debug=True, # Mantém o debug ativo para logs detalhados. cuda_device=",".join(map(str, range(self.NUM_GPUS_TOTAL))), skip_first_frames=0, load_cap=0, output_format='video' # Garante que sempre gere vídeo ) try: # Informa a UI que o processo começou. if progress: progress(0.01, "⌛ Initializing...") # Chama a função importada do script original, passando o callback de progresso. # Este callback será chamado de dentro da lógica de multi-processamento. result_tensor, original_fps, _, _ = run_inference_logic(args, progress_callback=progress) # Informa a UI que a inferência terminou e o salvamento vai começar. if progress: progress(0.95, "Saving the final video...") # Define o FPS final: usa o valor da UI ou o original do vídeo de entrada. final_fps = fps if fps and fps > 0 else original_fps save_frames_to_video(result_tensor, str(output_filepath), final_fps, args.debug) print(f"💾 Video saved successfully to: {output_filepath}") # Retorna o caminho do arquivo gerado para a UI. return str(output_filepath) except Exception as e: print(f"❌ Error during direct inference execution: {e}") import traceback traceback.print_exc() # Propaga o erro para a UI do Gradio, que o exibirá de forma amigável. raise