api/seedvr_server.py

# api/seedvr_server.py

import os
import sys
import time
import subprocess
import queue
import multiprocessing as mp
from pathlib import Path
from typing import Optional, Callable

# --- 1. Import dos Módulos Compartilhados ---
# É crucial que estes imports venham antes dos imports pesados (torch, etc.)
# para que o ambiente de multiprocessing seja configurado corretamente.

try:
    # Importa o gerenciador de GPUs que centraliza a lógica de alocação
    from api.gpu_manager import gpu_manager
    # Importa o serviço do LTX para podermos comandá-lo a liberar a VRAM
    from api.ltx_server_refactored import video_generation_service
except ImportError:
    print("ERRO FATAL: Não foi possível importar `gpu_manager` ou `video_generation_service`.")
    print("Certifique-se de que os arquivos `gpu_manager.py` e `ltx_server_refactored.py` existem em `api/`.")
    sys.exit(1)


# --- 2. Configuração de Ambiente e CUDA ---
if mp.get_start_method(allow_none=True) != 'spawn':
    mp.set_start_method('spawn', force=True)

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")

# Adiciona o caminho do repositório SeedVR
SEEDVR_REPO_PATH = Path(os.getenv("SEEDVR_ROOT", "/data/SeedVR"))
if str(SEEDVR_REPO_PATH) not in sys.path:
    sys.path.insert(0, str(SEEDVR_REPO_PATH))

# Imports pesados
import torch
import cv2
import numpy as np
from datetime import datetime


# --- 3. Funções Auxiliares de Processamento (Workers e I/O) ---
# (Estas funções não precisam de alteração)

def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load_cap=None):
    if debug: print(f"🎬 [SeedVR] Extraindo frames de: {video_path}")
    if not os.path.exists(video_path): raise FileNotFoundError(f"Arquivo de vídeo não encontrado: {video_path}")
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened(): raise ValueError(f"Não foi possível abrir o vídeo: {video_path}")
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    frames = []
    frames_loaded = 0
    for i in range(frame_count):
        ret, frame = cap.read()
        if not ret: break
        if i < skip_first_frames: continue
        if load_cap and frames_loaded >= load_cap: break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame.astype(np.float32) / 255.0)
        frames_loaded += 1
    cap.release()
    if not frames: raise ValueError(f"Nenhum frame extraído de: {video_path}")
    if debug: print(f"✅ [SeedVR] {len(frames)} frames extraídos com sucesso.")
    return torch.from_numpy(np.stack(frames)).to(torch.float16), fps

def save_frames_to_video(frames_tensor, output_path, fps=30.0, debug=False):
    if debug: print(f"💾 [SeedVR] Salvando {frames_tensor.shape[0]} frames em: {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, _ = frames_np.shape
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
    if not out.isOpened(): raise ValueError(f"Não foi possível criar o vídeo: {output_path}")
    for frame in frames_np:
        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    out.release()
    if debug: print(f"✅ [SeedVR] Vídeo salvo com sucesso: {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
    
    try:
        frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
        callback = (lambda b, t, _, m: progress_queue.put((proc_idx, b, t, m))) if progress_queue else None

        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=1.0, seed=shared_args["seed"],
            res_w=shared_args["resolution"], batch_size=shared_args["batch_size"],
            preserve_vram=shared_args["preserve_vram"], temporal_overlap=0,
            debug=shared_args["debug"], progress_callback=callback
        )
        return_queue.put((proc_idx, result_tensor.cpu().numpy()))
    except Exception as e:
        import traceback
        error_msg = f"ERRO no worker {proc_idx} (GPU {device_id}): {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))

# --- 4. CLASSE DO SERVIDOR PRINCIPAL ---

class SeedVRServer:
    def __init__(self, **kwargs):
        """Inicializa o servidor, define os caminhos e prepara o ambiente."""
        print("⚙️ SeedVRServer inicializando...")
        self.SEEDVR_ROOT = SEEDVR_REPO_PATH
        self.CKPTS_ROOT = Path("/data/seedvr_models_fp16")
        self.OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/output"))
        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")
        
        # OBTÉM AS GPUS ALOCADAS PELO GERENCIADOR CENTRAL
        self.device_list = gpu_manager.get_seedvr_devices()
        self.num_gpus = len(self.device_list)
        print(f"[SeedVR] Alocado para usar {self.num_gpus} GPU(s): {self.device_list}")

        for p in [self.CKPTS_ROOT, self.OUTPUT_ROOT, self.HF_HOME_CACHE]:
            p.mkdir(parents=True, exist_ok=True)

        self.setup_dependencies()
        print("📦 SeedVRServer pronto.")

    def setup_dependencies(self):
        """Garante que o repositório e os modelos estão presentes."""
        if not (self.SEEDVR_ROOT / ".git").exists():
            print(f"[SeedVR] Clonando repositório para {self.SEEDVR_ROOT}...")
            subprocess.run(["git", "clone", "--depth", "1", self.REPO_URL, str(self.SEEDVR_ROOT)], check=True)
        
        model_files = {
            "seedvr2_ema_7b_sharp_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}...")
                from huggingface_hub import hf_hub_download
                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("[SeedVR] Checkpoints verificados.")

    def run_inference(
        self,
        file_path: str, *,
        seed: int,
        resolution: int,
        batch_size: int,
        model: str = "seedvr2_ema_7b_sharp_fp16.safetensors",
        fps: Optional[float] = None,
        debug: bool = True,
        preserve_vram: bool = True,
        progress: Optional[Callable] = None
    ) -> str:
        """
        Executa o pipeline completo de upscaling de vídeo, gerenciando a memória da GPU.
        """
        if progress: progress(0.01, "⌛ Inicializando inferência SeedVR...")

        # --- NÓ 1: GERENCIAMENTO DE MEMÓRIA (SWAP) ---
        if gpu_manager.requires_memory_swap():
            print("[SWAP] SeedVR precisa da GPU. Movendo LTX para a CPU...")
            if progress: progress(0.02, "🔄 Liberando VRAM para o SeedVR...")
            video_generation_service.move_to_cpu()
            print("[SWAP] LTX movido para a CPU. VRAM liberada.")

        try:
            # --- NÓ 2: EXTRAÇÃO DE FRAMES ---
            if progress: progress(0.05, "🎬 Extraindo frames do vídeo...")
            frames_tensor, original_fps = extract_frames_from_video(file_path, debug)

            # --- NÓ 3: DIVISÃO PARA MULTI-GPU ---
            if self.num_gpus == 0:
                raise RuntimeError("SeedVR requer pelo menos 1 GPU alocada, mas não encontrou nenhuma.")
            
            print(f"[SeedVR] Dividindo {frames_tensor.shape[0]} frames em {self.num_gpus} chunks para processamento paralelo.")
            chunks = torch.chunk(frames_tensor, self.num_gpus, dim=0)
            
            manager = mp.Manager()
            return_queue = manager.Queue()
            progress_queue = manager.Queue() if progress else None
            
            shared_args = {
                "model": model, "model_dir": str(self.CKPTS_ROOT), "preserve_vram": preserve_vram,
                "debug": debug, "seed": seed, "resolution": resolution, "batch_size": batch_size
            }

            # --- NÓ 4: INÍCIO DOS WORKERS ---
            if progress: progress(0.1, f"🚀 Iniciando geração em {self.num_gpus} GPU(s)...")
            workers = []
            for idx, device_id in enumerate(self.device_list):
                p = mp.Process(target=_worker_process, args=(idx, device_id, chunks[idx].cpu().numpy(), shared_args, return_queue, progress_queue))
                p.start()
                workers.append(p)
            
            # --- NÓ 5: COLETA DE RESULTADOS E MONITORAMENTO ---
            results_np = [None] * self.num_gpus
            finished_workers = 0
            worker_progress = [0.0] * self.num_gpus
            while finished_workers < self.num_gpus:
                if progress_queue:
                    while not progress_queue.empty():
                        try:
                            p_idx, b_idx, b_total, msg = progress_queue.get_nowait()
                            if b_idx == -1: raise RuntimeError(f"Erro no Worker {p_idx}: {msg}")
                            if b_total > 0: worker_progress[p_idx] = b_idx / b_total
                            total_progress = sum(worker_progress) / self.num_gpus
                            progress(0.1 + total_progress * 0.85, desc=f"GPU {p_idx+1}/{self.num_gpus}: {msg}")
                        except queue.Empty: pass
                
                try:
                    proc_idx, result = return_queue.get(timeout=0.2)
                    if isinstance(result, str): raise RuntimeError(f"Worker {proc_idx} falhou: {result}")
                    results_np[proc_idx] = result
                    worker_progress[proc_idx] = 1.0
                    finished_workers += 1
                except queue.Empty: pass

            for p in workers: p.join()

            # --- NÓ 6: FINALIZAÇÃO ---
            if any(r is None for r in results_np):
                raise RuntimeError("Um ou mais workers falharam ao retornar um resultado.")

            result_tensor = torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)
            if progress: progress(0.95, "💾 Salvando o vídeo final...")
            
            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"

            final_fps = fps if fps and fps > 0 else original_fps
            save_frames_to_video(result_tensor, str(output_filepath), final_fps, debug)
            
            print(f"✅ Vídeo salvo com sucesso em: {output_filepath}")
            return str(output_filepath)

        finally:
            # --- NÓ 7: RESTAURAÇÃO DE MEMÓRIA (SWAP BACK) ---
            if gpu_manager.requires_memory_swap():
                print("[SWAP] Inferência do SeedVR concluída. Movendo LTX de volta para a GPU...")
                if progress: progress(0.99, "🔄 Restaurando o ambiente LTX...")
                ltx_device = gpu_manager.get_ltx_device()
                video_generation_service.move_to_device(ltx_device)
                print(f"[SWAP] LTX de volta em {ltx_device}.")

# --- PONTO DE ENTRADA ---
if __name__ == "__main__":
    print("🚀 Executando o servidor SeedVR em modo autônomo para inicialização...")
    try:
        server = SeedVRServer()
        print("✅ Servidor inicializado com sucesso. Pronto para receber chamadas.")
    except Exception as e:
        print(f"❌ Falha ao inicializar o servidor SeedVR: {e}")
        traceback.print_exc()
        sys.exit(1)