Update aduc_framework/orchestrator.py

aduc_framework/orchestrator.py

@@ -2,12 +2,14 @@
 #
 # Copyright (C) August 4, 2025  Carlos Rodrigues dos Santos
 #
-# Versão 3.2.0 (Com Orquestração de Pós-Produção)
+# Versão 4.0.0 (Multi-Pool Orchestration)
 #
-# Esta versão representa a camada de orquestração do Aduc Framework.
-# Ela é agnóstica a qualquer interface (UI ou API) e opera com
-# tipos de dados bem definidos (Pydantic) e um estado de geração central.
-# Agora, é responsável por inicializar os singletons e orquestrar todas as etapas.
+# Esta versão do orquestrador opera com uma arquitetura de múltiplos pools de
+# especialistas. A responsabilidade de gerenciar a VRAM e o ciclo de vida dos
+# modelos foi delegada aos seus respectivos managers (LtxPoolManager,
+# SeedVrPoolManager, MMAudioPoolManager), cada um operando em um conjunto
+# dedicado de GPUs. O orquestrador agora foca exclusivamente na lógica de
+# alto nível do fluxo de trabalho.
 
 import logging
 from typing import List, Dict, Any, Tuple, Callable, Optional, Generator
@@ -20,15 +22,20 @@ import time
 import gc
 import torch
 
-# Importa componentes internos do framework
+# Componentes internos do framework
 from .director import AducDirector
 from .types import GenerationState, PreProductionParams, ProductionParams
+
+# Engenheiros de alto nível que definem a lógica do fluxo
 from .engineers import deformes2d_thinker_singleton, deformes3d_engine_singleton, Deformes4DEngine
-# Importa managers diretamente para as tarefas de pós-produção
+
+# Managers (Pools) de especialistas que executam as tarefas em hardware dedicado
 from .managers.latent_enhancer_manager import latent_enhancer_specialist_singleton
 from .managers.seedvr_manager import seedvr_manager_singleton
 from .managers.mmaudio_manager import mmaudio_manager_singleton
 from .managers.vae_manager import vae_manager_singleton
+
+# Ferramentas de utilidade
 from .tools.video_encode_tool import video_encode_tool_singleton
 
 logger = logging.getLogger(__name__)
@@ -39,7 +46,7 @@ class AducOrchestrator:
     """
     Implementa o Maestro (Γ), a camada de orquestração central do Aduc Framework.
     Ele recebe solicitações, atualiza o estado de geração, delega tarefas para os
-    engenheiros especialistas e retorna o estado atualizado.
+    engenheiros e seus pools de especialistas, e retorna o estado atualizado.
     """
     def __init__(self, workspace_dir: str):
         self.director = AducDirector(workspace_dir)
@@ -48,12 +55,14 @@ class AducOrchestrator:
         self.painter = deformes3d_engine_singleton
         self.painter.initialize(workspace_dir)
         self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        logger.info("ADUC Maestro (Framework Core) está no pódio. Engenheiros especialistas prontos.")
+        logger.info("ADUC Maestro (Framework Core) pronto para reger a orquestra de especialistas.")
 
     def get_current_state(self) -> GenerationState:
+        """Retorna o estado de geração atual."""
         return self.director.get_full_state()
 
     def process_image_for_story(self, image_path: str, size: int, filename: str) -> str:
+        """Processa uma imagem de referência para o formato quadrado padrão."""
         img = Image.open(image_path).convert("RGB")
         img_square = ImageOps.fit(img, (size, size), Image.Resampling.LANCZOS)
         processed_path = os.path.join(self.director.workspace_dir, filename)
@@ -61,35 +70,48 @@ class AducOrchestrator:
         logger.info(f"Imagem de referência processada e salva em: {processed_path}")
         return processed_path
 
+    # --- ETAPA 1: PRÉ-PRODUÇÃO ---
     def task_pre_production(self, params: PreProductionParams, progress_callback: ProgressCallback = None) -> Tuple[List[str], List[str], GenerationState]:
+        """Orquestra a criação do storyboard e dos keyframes visuais."""
         logger.info("Maestro: Iniciando tarefa de Pré-Produção.")
         self.director.update_parameters("pre_producao", params)
+        
         if progress_callback: progress_callback(0.1, "Gerando storyboard...")
         storyboard_list = deformes2d_thinker_singleton.generate_storyboard(prompt=params.prompt, num_keyframes=params.num_keyframes, ref_image_paths=params.ref_paths)
         self.director.update_pre_production_state(params.prompt, params.ref_paths, storyboard_list)
+        
         if progress_callback: progress_callback(0.2, "Iniciando geração de keyframes...")
+        # O 'painter' (Deformes3D) internamente usará o LTX Pool Manager para gerar as imagens
         keyframes_detailed_data = self.painter.generate_keyframes_from_storyboard(generation_state=self.director.get_full_state_as_dict(), progress_callback=progress_callback)
         self.director.update_keyframes_state(keyframes_detailed_data)
+        
         final_keyframe_paths = [kf["caminho_pixel"] for kf in keyframes_detailed_data]
         final_state = self.director.get_full_state()
         logger.info("Maestro: Tarefa de Pré-Produção concluída.")
         return storyboard_list, final_keyframe_paths, final_state
 
+    # --- ETAPA 2: PRODUÇÃO ---
     def task_produce_original_movie(self, params: ProductionParams, progress_callback: ProgressCallback = None) -> Tuple[str, List[str], GenerationState]:
+        """Orquestra a geração do vídeo principal a partir dos keyframes."""
         logger.info("Maestro: Iniciando tarefa de Produção do Filme Original.")
         self.director.update_parameters("producao", params)
+        
+        # O 'editor' (Deformes4D) usa intensivamente o LTX Pool Manager
         result_data = self.editor.generate_original_movie(full_generation_state=self.director.get_full_state_as_dict(), progress_callback=progress_callback)
         self.director.update_video_state(result_data["video_data"])
+        
         final_video_path = result_data["final_path"]
         latent_paths = result_data["latent_paths"]
         final_state = self.director.get_full_state()
         logger.info("Maestro: Tarefa de Produção do Filme Original concluída.")
         return final_video_path, latent_paths, final_state
 
+    # --- ETAPA 3: PÓS-PRODUÇÃO (Cadeia de Efeitos) ---
+
     def task_run_latent_upscaler(self, latent_paths: List[str], chunk_size: int, progress_callback: ProgressCallback = None) -> Generator[Dict[str, Any], None, None]:
-        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado corretamente.")
-        if not latent_paths:
-            raise ValueError("Não é possível fazer o upscale: nenhum caminho de latente foi fornecido.")
+        """Aplica upscale 2x nos latentes e os decodifica para um novo vídeo."""
+        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado.")
+        if not latent_paths: raise ValueError("Nenhum caminho de latente fornecido para o upscale.")
         
         logger.info("--- ORQUESTRADOR: Tarefa de Upscaling de Latentes ---")
         run_timestamp = int(time.time())
@@ -97,34 +119,26 @@ class AducOrchestrator:
         os.makedirs(temp_dir, exist_ok=True)
         
         final_upscaled_clip_paths = []
-        num_chunks = -(-len(latent_paths) // chunk_size) if chunk_size > 0 else 0
+        num_chunks = -(-len(latent_paths) // chunk_size)
         
         for i in range(num_chunks):
-            chunk_start_index = i * chunk_size
-            chunk_end_index = chunk_start_index + chunk_size
-            chunk_paths = latent_paths[chunk_start_index:chunk_end_index]
-            
+            chunk_paths = latent_paths[i * chunk_size:(i + 1) * chunk_size]
             if progress_callback: progress_callback(i / num_chunks, f"Upscalando & Decodificando Lote {i+1}/{num_chunks}")
 
             tensors_in_chunk = [torch.load(p, map_location=self.device) for p in chunk_paths]
             sub_group_latent = torch.cat(tensors_in_chunk, dim=2)
-            del tensors_in_chunk; gc.collect(); torch.cuda.empty_cache()
             
             upscaled_latent_chunk = latent_enhancer_specialist_singleton.upscale(sub_group_latent)
-            del sub_group_latent; gc.collect(); torch.cuda.empty_cache()
-            
             pixel_tensor = vae_manager_singleton.decode(upscaled_latent_chunk)
-            del upscaled_latent_chunk; gc.collect(); torch.cuda.empty_cache()
 
-            base_name = f"upscaled_clip_{i:04d}_{run_timestamp}"
-            current_clip_path = os.path.join(temp_dir, f"{base_name}.mp4")
+            current_clip_path = os.path.join(temp_dir, f"upscaled_clip_{i:04d}.mp4")
             self.editor.save_video_from_tensor(pixel_tensor, current_clip_path, fps=24)
             final_upscaled_clip_paths.append(current_clip_path)
-            del pixel_tensor; gc.collect(); torch.cuda.empty_cache()
             
-            yield {"progress": (i + 1) / num_chunks, "desc": f"Lote {i+1}/{num_chunks} completo."}
+            del tensors_in_chunk, sub_group_latent, upscaled_latent_chunk, pixel_tensor
+            gc.collect(); torch.cuda.empty_cache()
+            yield {"progress": (i + 1) / num_chunks}
 
-        if progress_callback: progress_callback(0.98, "Montando vídeo com upscale...")
         final_video_path = os.path.join(self.director.workspace_dir, f"upscaled_movie_{run_timestamp}.mp4")
         video_encode_tool_singleton.concatenate_videos(final_upscaled_clip_paths, final_video_path, self.director.workspace_dir)
         
@@ -133,22 +147,29 @@ class AducOrchestrator:
         yield {"final_path": final_video_path}
 
     def task_run_hd_mastering(self, source_video_path: str, model_version: str, steps: int, prompt: str, progress_callback: ProgressCallback = None) -> Generator[Dict[str, Any], None, None]:
-        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado corretamente.")
+        """Aplica masterização em HD usando o pool de GPUs do SeedVR."""
+        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado.")
         logger.info(f"--- ORQUESTRADOR: Tarefa de Masterização HD com SeedVR {model_version} ---")
         
+        # Com pools dedicados, não há mais necessidade de descarregar modelos LTX.
+        # A chamada �� direta e o SeedVrPoolManager gerencia seus próprios recursos.
+        
         run_timestamp = int(time.time())
         output_path = os.path.join(self.director.workspace_dir, f"hd_mastered_movie_{model_version}_{run_timestamp}.mp4")
         
         final_path = seedvr_manager_singleton.process_video(
             input_video_path=source_video_path,
             output_video_path=output_path,
-            prompt=prompt
+            prompt=prompt,
+            model_version=model_version,
+            steps=steps
         )
         logger.info(f"Masterização HD completa! Vídeo final em: {final_path}")
         yield {"final_path": final_path}
 
     def task_run_audio_generation(self, source_video_path: str, audio_prompt: str, progress_callback: ProgressCallback = None) -> Generator[Dict[str, Any], None, None]:
-        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado corretamente.")
+        """Gera e adiciona áudio ao vídeo usando o pool de GPUs do MMAudio."""
+        if not self.director.workspace_dir: raise RuntimeError("Orchestrator não inicializado.")
         logger.info(f"--- ORQUESTRADOR: Tarefa de Geração de Áudio ---")
 
         if progress_callback: progress_callback(0.1, "Preparando para geração de áudio...")
@@ -157,14 +178,20 @@ class AducOrchestrator:
         source_name = Path(source_video_path).stem
         output_path = os.path.join(self.director.workspace_dir, f"{source_name}_with_audio_{run_timestamp}.mp4")
 
-        result = subprocess.run(
-            ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", source_video_path],
-            capture_output=True, text=True, check=True
-        )
-        duration = float(result.stdout.strip())
+        try:
+            result = subprocess.run(
+                ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", source_video_path],
+                capture_output=True, text=True, check=True
+            )
+            duration = float(result.stdout.strip())
+        except Exception as e:
+            logger.error(f"Não foi possível obter a duração do vídeo '{source_video_path}': {e}", exc_info=True)
+            yield {"error": "Falha ao obter duração do vídeo."}
+            return
 
         if progress_callback: progress_callback(0.5, "Gerando trilha de áudio...")
         
+        # A chamada é direta. O MMAudioPoolManager gerencia suas GPUs dedicadas.
         final_path = mmaudio_manager_singleton.generate_audio_for_video(
             video_path=source_video_path,
             prompt=audio_prompt,