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)
+(MODIFICADO PARA SER IMPORTÁVEL E SUPORTAR CALLBACKS DE PROGRESSO EM MULTIPROCESSING)
 """
 
 import sys
@@ -9,12 +9,14 @@ 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 (essencial para performance)
+# 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
@@ -24,7 +26,6 @@ _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:
-        # Se apenas uma GPU for especificada, restringe a visibilidade do PyTorch a ela.
         os.environ["CUDA_VISIBLE_DEVICES"] = device_list_env[0]
 
 # -------------------------------------------------------------
@@ -48,101 +49,45 @@ def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load
     """
     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}")
-    
+    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}")
-    
+    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:
-            print(f"🔢 Will load maximum {load_cap} frames")
-    
-    frames = []
-    frame_idx = 0
-    frames_loaded = 0
-    
+    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
-        
+        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 else frame_count
-            print(f"📹 Extracted {frames_loaded}/{total_to_load} frames")
-    
+        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")
-    
+    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}")
-    
+    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}")
-    
+    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}")
-    
+    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}")
+    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):
     """
@@ -157,30 +102,33 @@ def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, p
     
     frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
     
-    # Cria uma função de callback local que envia o progresso para a fila de comunicação
     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
 
-    runner = configure_runner(shared_args["model"], shared_args["model_dir"], shared_args["preserve_vram"], shared_args["debug"])
-
-    # Passa o callback local para o generation_loop, que sabe como usá-lo
-    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 
-    )
-    # Envia o resultado final de volta para o processo pai
-    return_queue.put((proc_idx, result_tensor.cpu().numpy()))
+    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 entre as GPUs e gerencia os processos filhos, monitorando o progresso.
+    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)
@@ -202,40 +150,54 @@ def _gpu_processing(frames_tensor, device_list, args, progress_callback=None):
         workers.append(p)
 
     results_np = [None] * num_devices
-    collected_workers = 0
-    worker_progress = [0.0] * num_devices # Armazena o progresso individual de cada worker
+    finished_workers = [False] * num_devices
+    worker_progress = [0.0] * num_devices
     
-    while collected_workers < num_devices:
-        # 1. Processa todas as mensagens de progresso na fila de forma não-bloqueante
+    while not all(finished_workers):
         if progress_queue:
             while not progress_queue.empty():
-                proc_idx, batch_idx, total_batches, message = progress_queue.get()
-                if total_batches > 0:
-                    worker_progress[proc_idx] = batch_idx / total_batches
+                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
+                    
+                    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
                 
-                # Calcula o progresso geral como a média do progresso de todos os workers
                 total_progress = sum(worker_progress) / num_devices
-                
-                # Chama o callback principal (do Gradio) com a informação formatada
-                progress_callback(total_progress, desc=f"GPU {proc_idx+1}: {message}")
-
-        # 2. Verifica se algum worker terminou e enviou seu resultado
-        if not return_queue.empty():
-            proc_idx, res_np = return_queue.get()
-            results_np[proc_idx] = res_np
-            worker_progress[proc_idx] = 1.0 # Marca este worker como 100% concluído
-            collected_workers += 1
+                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) # Pequena pausa para evitar uso excessivo da CPU no loop
+        time.sleep(0.2)
 
     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.")
+
     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)")
@@ -262,6 +224,7 @@ def run_inference_logic(args, progress_callback=None):
     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(
@@ -274,10 +237,10 @@ def run_inference_logic(args, progress_callback=None):
     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)
     
-    # Passa o callback para a função de processamento, que o gerenciará
     result_tensor = _gpu_processing(frames_tensor, device_list, args, progress_callback)
 
     generation_time = time.time() - processing_start
@@ -285,7 +248,6 @@ def run_inference_logic(args, progress_callback=None):
         print(f"🔄 Tempo de Geração: {generation_time:.2f}s")
         print(f"📊 Resultado: {result_tensor.shape}, dtype: {result_tensor.dtype}")
         
-    # Retorna o tensor e metadados em memória para o chamador
     return result_tensor, original_fps, generation_time, len(frames_tensor)
 
 def main():
@@ -295,10 +257,8 @@ def main():
     print(f"🚀 SeedVR2 Video Upscaler CLI iniciado às {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
     args = parse_arguments()
     try:
-        # Chama a função de lógica principal
         result_tensor, original_fps, _, _ = run_inference_logic(args)
         
-        # Salva o resultado no disco, como esperado pelo modo CLI
         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!")
@@ -309,6 +269,5 @@ def main():
         traceback.print_exc()
         sys.exit(1)
 
-# Ponto de entrada para quando o script é executado diretamente
 if __name__ == "__main__":
     main()