api/seedvr_server.py

# FILE: api/seedvr_server.py
# DESCRIPTION: Backend service for SeedVR video upscaling.
# Features multi-GPU processing, memory swapping with other services,
# and detailed debug logging.

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

# ==============================================================================
# --- IMPORTAÇÃO DOS MÓDulos Compartilhados ---
# ==============================================================================
try:
    from api.gpu_manager import gpu_manager
    from api.ltx_server_refactored_complete import video_generation_service
    from api.utils.debug_utils import log_function_io
except ImportError:
    # Fallback para o decorador caso o import falhe
    def log_function_io(func):
        return func
    logging.critical("CRITICAL: Failed to import shared modules like gpu_manager or video_generation_service.", exc_info=True)
    # Em um cenário real, poderíamos querer sair aqui ou desativar o servidor.
    # Por enquanto, a aplicação pode tentar continuar sem o SeedVR.
    raise

# ==============================================================================
# --- CONFIGURAÇÃO DE AMBIENTE ---
# ==============================================================================
if mp.get_start_method(allow_none=True) != 'spawn':
    try:
        mp.set_start_method('spawn', force=True)
    except RuntimeError:
        logging.warning("Multiprocessing context is already set. Skipping.")

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

# Adiciona o caminho do repositório SeedVR ao sys.path
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 após a configuração de path e multiprocessing
import torch
import cv2
import numpy as np
from datetime import datetime

# ==============================================================================
# --- FUNÇÕES WORKER E AUXILIARES (I/O de Vídeo) ---
# ==============================================================================
# (Estas funções são de baixo nível e não precisam do decorador de log principal)

def extract_frames_from_video(video_path, debug=False):
    if debug: logging.debug(f"🎬 [SeedVR] Extracting frames from: {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 IOError(f"Cannot open video file: {video_path}")
    
    fps = cap.get(cv2.CAP_PROP_FPS)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret: break
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        frames.append(frame.astype(np.float32) / 255.0)
    cap.release()
    if not frames: raise ValueError(f"No frames extracted from: {video_path}")
    if debug: logging.debug(f"✅ [SeedVR] {len(frames)} frames extracted successfully.")
    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: logging.debug(f"💾 [SeedVR] Saving {frames_tensor.shape[0]} frames to: {output_path}")
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8)
    _, H, W, _ = frames_np.shape
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
    if not out.isOpened(): raise IOError(f"Cannot create video writer for: {output_path}")
    for frame in frames_np:
        out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
    out.release()
    if debug: logging.debug(f"✅ [SeedVR] 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)
    # É importante reimportar torch aqui para que ele respeite a variável de ambiente
    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('cuda', dtype=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_h=shared_args["resolution"], # Assumindo que a UI passa a altura
            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"ERROR in worker {proc_idx} (GPU {device_id}): {e}\n{traceback.format_exc()}"
        logging.error(error_msg)
        if progress_queue: progress_queue.put((proc_idx, -1, -1, error_msg))
        return_queue.put((proc_idx, error_msg))

# ==============================================================================
# --- CLASSE DO SERVIDOR PRINCIPAL ---
# ==============================================================================

class SeedVRServer:
    @log_function_io
    def __init__(self, **kwargs):
        """Inicializa o servidor, define os caminhos e prepara o ambiente."""
        logging.info("⚙️ SeedVRServer initializing...")
        self.OUTPUT_ROOT = Path(os.getenv("OUTPUT_ROOT", "/app/output"))
        
        self.device_list = gpu_manager.get_seedvr_devices()
        self.num_gpus = len(self.device_list)
        logging.info(f"[SeedVR] Allocated to use {self.num_gpus} GPU(s): {self.device_list}")

        # O setup de dependências já é feito pelo setup.py principal, então aqui apenas verificamos
        if not SEEDVR_REPO_PATH.is_dir():
            raise NotADirectoryError(f"SeedVR repository not found at {SEEDVR_REPO_PATH}. Run setup.py first.")
        
        logging.info("📦 SeedVRServer ready.")

    @log_function_io
    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, "⌛ Initializing SeedVR inference...")

        if gpu_manager.requires_memory_swap():
            logging.warning("[SWAP] Memory swapping is active. Moving LTX service to CPU to free VRAM for SeedVR.")
            if progress: progress(0.02, "🔄 Freeing VRAM for SeedVR...")
            video_generation_service.move_to_cpu()

        try:
            if progress: progress(0.05, "🎬 Extracting frames from video...")
            frames_tensor, original_fps = extract_frames_from_video(file_path, debug)

            if self.num_gpus == 0:
                raise RuntimeError("SeedVR requires at least 1 allocated GPU, but found none.")
            
            logging.info(f"[SeedVR] Splitting {frames_tensor.shape[0]} frames into {self.num_gpus} chunks for parallel processing.")
            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": "/data/models/SeedVR", "preserve_vram": preserve_vram,
                "debug": debug, "seed": seed, "resolution": resolution, "batch_size": batch_size
            }

            if progress: progress(0.1, f"🚀 Starting generation on {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)
            
            results_np = [None] * self.num_gpus
            finished_workers = 0
            # (Loop de monitoramento de progresso e coleta de resultados)
            # ...

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

            result_tensor = torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)
            if progress: progress(0.95, "💾 Saving final video...")
            
            out_dir = self.OUTPUT_ROOT / f"seedvr_run_{int(time.time())}"
            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)
            
            logging.info(f"✅ Video successfully saved to: {output_filepath}")
            return str(output_filepath)

        finally:
            # --- CORREÇÃO IMPORTANTE ---
            # Restaura o LTX para seus dispositivos corretos (main e vae)
            if gpu_manager.requires_memory_swap():
                logging.warning("[SWAP] SeedVR inference finished. Moving LTX service back to GPU(s)...")
                if progress: progress(0.99, "🔄 Restoring LTX environment...")
                ltx_main_device = gpu_manager.get_ltx_device()
                ltx_vae_device = gpu_manager.get_ltx_vae_device()
                # Chama a função move_to_device com os dois dispositivos
                video_generation_service.move_to_device(
                    main_device_str=str(ltx_main_device),
                    vae_device_str=str(ltx_vae_device)
                )
                logging.info(f"[SWAP] LTX service restored to Main: {ltx_main_device}, VAE: {ltx_vae_device}.")

# --- PONTO DE ENTRADA E INSTANCIAÇÃO ---
# A instância é criada na primeira importação.
try:
    # A classe é instanciada globalmente para ser usada pela UI
    seedvr_server_singleton = SeedVRServer()
except Exception as e:
    logging.critical("Failed to initialize SeedVRServer singleton.", exc_info=True)
    seedvr_server_singleton = None