Update inference_cli.py

inference_cli.py

@@ -35,15 +35,15 @@ import cv2
 import numpy as np
 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`
+# 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)
-root_dir = os.path.join(script_dir, '..', '..')
-if root_dir not in sys.path:
-    sys.path.insert(0, root_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):
     """
@@ -68,7 +68,7 @@ def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load
         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:
+        if load_cap and load_cap > 0:
             print(f"🔢 Will load maximum {load_cap} frames")
     
     frames = []
@@ -95,7 +95,7 @@ def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load
         frames_loaded += 1
         
         if debug and frames_loaded % 100 == 0:
-            total_to_load = min(frame_count, load_cap) if load_cap else frame_count
+            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()
@@ -211,39 +211,34 @@ def _gpu_processing(frames_tensor, device_list, args, progress_callback=None):
     worker_progress = [0.0] * num_devices
     
     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
+                    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
 
-        # 2. Verifica se algum worker terminou
         try:
-            proc_idx, result = return_queue.get(timeout=0.1) # Usa um timeout curto
+            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 # Marca como 100%
+            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 # Continua o loop se não houver resultados ainda
+            pass
 
     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("One or more workers failed to return a result.")
 
@@ -293,7 +288,7 @@ def run_inference_logic(args, progress_callback=None):
     
     if progress_callback: progress_callback(0.1, "Starting generation...")
     processing_start = time.time()
-    download_weight(args.model, args.model_dir)
+    download_weight(args.model, args.model_dir or "seedvr_models")
     
     result_tensor = _gpu_processing(frames_tensor, device_list, args, progress_callback)
 
@@ -313,6 +308,9 @@ def main():
     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!")