Update api/seedvr_server.py

api/seedvr_server.py

@@ -21,12 +21,342 @@ if str(SEEDVR_REPO_PATH) not in sys.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
+#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()
+
 
 
 
@@ -46,16 +376,16 @@ class SeedVRServer:
         self.NUM_GPUS_TOTAL = int(os.getenv("NUM_GPUS", "4"))
 
 
-        if INIT:
-            print("🚀 SeedVRServer ja inicializando...")
-        else:
-            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)
+    
+        #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
+        self.setup_dependencies()
+        print("📦 SeedVRServer (Modo de Chamada Direta) pronto.")
+        INIT = True
             
 
     def setup_dependencies(self):