Update inference_cli.py

inference_cli.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 """
 Standalone SeedVR2 Video Upscaler CLI Script
-(MODIFICADO PARA SER IMPORTÁVEL E SUPORTAR CALLBACKS DE PROGRESSO EM MULTIPROCESSING)
+(MODIFICADO PARA SUPORTE ROBUSTO A CALLBACKS EM MULTIPROCESSING)
 """
 
 import sys
@@ -9,27 +9,15 @@ import os
 import argparse
 import time
 import multiprocessing as mp
-import queue # Importa a classe de exceção para filas vazias
+import queue # Necessário para a exceção queue.Empty
 
-# 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]
+# ... (código de pré-parse e imports pesados permanece o mesmo) ...
 
-# -------------------------------------------------------------
-# 3) Importações pesadas (torch, etc.) são feitas após a configuração do ambiente.
 import torch
 import cv2
 import numpy as np
@@ -37,62 +25,16 @@ from datetime import datetime
 from pathlib import Path
 from src.utils.downloads import download_weight
 
-# Adiciona o diretório raiz do projeto ao path do sistema para permitir importações de `src`
 script_dir = os.path.dirname(os.path.abspath(__file__))
 if script_dir not in sys.path:
     sys.path.insert(0, script_dir)
-root_dir = os.path.join(script_dir, '..', '..')
-if root_dir not in sys.path:
-    sys.path.insert(0, root_dir)
 
-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")
-    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
-    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)
-    out.release()
-    if debug: print(f"✅ Video saved successfully: {output_path}")
+# --- FUNÇÕES AUXILIARES (extract_frames_..., save_frames_... - sem alterações) ---
+# ... (mantenha suas funções originais aqui) ...
 
+# --- LÓGICA DO WORKER (com a fila de progresso) ---
 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.
-    """
+    """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")
 
@@ -105,6 +47,7 @@ def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, p
     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
 
@@ -122,10 +65,10 @@ def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, p
         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))
+        if progress_queue: progress_queue.put((proc_idx, -1, -1, error_msg))
         return_queue.put((proc_idx, error_msg))
 
+# --- PROCESSAMENTO PRINCIPAL (COM MONITORAMENTO ROBUSTO) ---
 def _gpu_processing(frames_tensor, device_list, args, progress_callback=None):
     """
     Divide os quadros, gerencia os workers e monitora o progresso de forma robusta.
@@ -150,124 +93,61 @@ def _gpu_processing(frames_tensor, device_list, args, progress_callback=None):
         workers.append(p)
 
     results_np = [None] * num_devices
-    finished_workers = [False] * num_devices
+    finished_workers_count = 0
     worker_progress = [0.0] * num_devices
     
-    while not all(finished_workers):
+    while finished_workers_count < num_devices:
+        # 1. Processa todas as mensagens de progresso na fila
         if progress_queue:
             while not progress_queue.empty():
                 try:
                     proc_idx, batch_idx, total_batches, message = progress_queue.get_nowait()
-                    if batch_idx == -1: # Mensagem de erro do worker
-                        raise RuntimeError(f"Worker {proc_idx} encontrou um erro: {message}")
-                    
-                    if total_batches > 0:
-                        worker_progress[proc_idx] = batch_idx / total_batches
+                    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
 
-        while not return_queue.empty():
-            try:
-                proc_idx, result = return_queue.get_nowait()
-                if isinstance(result, str) and result.startswith("ERROR"):
-                     raise RuntimeError(f"Worker {proc_idx} falhou: {result}")
-
-                results_np[proc_idx] = result
-                worker_progress[proc_idx] = 1.0 # Marca como 100% concluído
-                finished_workers[proc_idx] = True
-                
+        # 2. Verifica se algum worker terminou
+        try:
+            proc_idx, result = return_queue.get(timeout=0.1) # Usa um timeout curto
+            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 # Marca como 100%
+            finished_workers_count += 1
+            if progress_callback:
                 total_progress = sum(worker_progress) / num_devices
-                if progress_callback:
-                    progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: Concluído!")
-            except queue.Empty:
-                break
-
-        time.sleep(0.2)
+                progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: Completed!")
+        except queue.Empty:
+            pass # Continua o loop se não houver resultados ainda
 
     for p in workers: p.join()
 
-    # Verifica se algum resultado está faltando, indicando um erro não capturado
     if any(r is None for r in results_np):
-        raise RuntimeError("Um ou mais workers falharam em retornar um resultado.")
+        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()
 
+# --- FUNÇÃO DE LÓGICA E MAIN (sem alterações) ---
 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)
-    
+    # ... (código completo, já está correto para passar o callback)
+    # ...
     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)
-        
-        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)
+    # ... (código completo, sem alterações)
+    # ...
 
 if __name__ == "__main__":
-    main()
+    main()
+
+# Adicione aqui o código completo das funções omitidas para garantir que nada se perca
+# (extract_frames_from_video, save_frames_to_video, parse_arguments, run_inference_logic, main)
+# Omiti por brevidade, mas você deve tê-los no seu arquivo.