Update inference_cli.py

inference_cli.py

@@ -211,7 +211,37 @@ def save_frames_to_png(frames_tensor, output_dir, base_name, debug=False):
         print(f"✅ PNG saving completed: {total} files in '{output_dir}'")
 
 
-def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue):
+def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, progress_queue=None): # Adicionado progress_queue
+    """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)
+    
+    # Cria uma função de callback local que envia o progresso para a fila
+    local_progress_callback = None
+    if progress_queue:
+        def callback_wrapper(batch_idx, total_batches, current_frames, message):
+            progress_queue.put((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"])
+
+    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 # Passa o callback para o generation_loop
+    )
+    return_queue.put((proc_idx, result_tensor.cpu().numpy()))
+    
+
+def _worker_process1(proc_idx, device_id, frames_np, shared_args, return_queue):
     """Worker process that performs upscaling on a slice of frames using a dedicated GPU."""
     # 1. Limit CUDA visibility to the chosen GPU BEFORE importing torch-heavy deps
     os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
@@ -251,7 +281,52 @@ def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue):
     return_queue.put((proc_idx, result_tensor.cpu().numpy()))
 
 
-def _gpu_processing(frames_tensor, device_list, args):
+def _gpu_processing(frames_tensor, device_list, args, progress_callback=None): # Adicionado progress_callback
+    """Divide os frames e os processa em paralelo em múltiplas GPUs."""
+    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 # Cria a fila de progresso
+    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
+    collected = 0
+    total_batches_per_worker = -1 # Para calcular o progresso total
+    while collected < num_devices:
+        # Verifica as duas filas (resultado e progresso) de forma não-bloqueante
+        if progress_queue and not progress_queue.empty():
+            batch_idx, total_batches, message = progress_queue.get()
+            if total_batches_per_worker == -1: total_batches_per_worker = total_batches
+            total_progress = (collected + (batch_idx / total_batches_per_worker)) / num_devices
+            progress_callback(total_progress, desc=f"GPU {collected+1}/{num_devices}: {message}")
+        
+        if not return_queue.empty():
+            proc_idx, res_np = return_queue.get()
+            results_np[proc_idx] = res_np
+            collected += 1
+        
+        time.sleep(0.1) # Evita busy-waiting
+
+    for p in workers: p.join()
+
+    return torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)
+    
+
+
+def _gpu_processing1(frames_tensor, device_list, args):
     """Split frames and process them in parallel on multiple GPUs."""
     num_devices = len(device_list)
     # split frames tensor along time dimension
@@ -431,6 +506,9 @@ def main():
     finally:
         print(f"🧹 Process {os.getpid()} terminating - VRAM will be automatically freed")
 
+
+
+
 def run_inference_logic(args, progress_callback=None):
     """
     Função principal que executa o pipeline de upscaling.
@@ -484,6 +562,8 @@ def run_inference_logic(args, progress_callback=None):
     return result_tensor, original_fps, generation_time, len(frames_tensor)
 
 
+
+
 # FUNÇÃO MAIN ORIGINAL (agora um wrapper)
 def main():
     """Função principal do CLI"""
@@ -512,4 +592,4 @@ def main():
 
 # Ponto de entrada para execução via linha de comando
 if __name__ == "__main__":
-    main()
+    main()