Create ltx_aduc_orchestrator.py

api/ltx/ltx_aduc_orchestrator.py

@@ -0,0 +1,151 @@
+# FILE: api/ltx_aduc_orchestrator.py
+# DESCRIPTION: The main workflow orchestrator for the ADUC-SDR LTX suite.
+# It acts as the primary entry point for the UI, coordinating the specialized
+# LTX and VAE clients to execute a complete video generation pipeline from prompt to MP4.
+
+import logging
+import time
+from PIL import Image
+from typing import Optional, Dict
+
+# O Orquestrador importa os CLIENTES especialistas que ele vai coordenar.
+# Estes clientes são responsáveis por submeter os trabalhos ao pool de workers.
+from api.ltx.ltx_aduc_pipeline import ltx_aduc_pipeline
+from api.ltx.vae_aduc_pipeline import vae_aduc_pipeline
+
+# O Orquestrador importa as FERRAMENTAS de que precisa para as tarefas finais.
+from tools.video_encode_tool import video_encode_tool_singleton
+
+# ==============================================================================
+# --- A CLASSE ORQUESTRADORA (Cérebro do Workflow) ---
+# ==============================================================================
+
+class LtxAducOrchestrator:
+    """
+    Orquestra o fluxo de trabalho completo de geração de vídeo,
+    coordenando os clientes LTX e VAE. É o ponto de entrada principal para a UI.
+    """
+    def __init__(self):
+        """
+        Inicializa o orquestrador. A inicialização é leve, pois os modelos
+        pesados são gerenciados pelo LTXAducManager em segundo plano.
+        """
+        self.output_dir = "/app/output"  # Diretório padrão para salvar os vídeos
+        logging.info("✅ LTX ADUC Orchestrator initialized and ready.")
+
+    def __call__(
+        self,
+        prompt: str,
+        initial_image: Optional[Image.Image] = None,
+        height: int = 432,
+        width: int = 768,
+        duration_in_seconds: float = 4.0,
+        ltx_configs: Optional[Dict] = None,
+        output_filename_base: str = "ltx_aduc_video"
+    ) -> Optional[str]:
+        """
+        Ponto de entrada principal do Orquestrador. Executa o pipeline completo.
+
+        Args:
+            prompt (str): O prompt de texto completo. Cada nova linha é tratada como uma cena.
+            initial_image (Optional[Image.Image]): Uma imagem PIL para condicionar a primeira cena.
+            height (int): Altura do vídeo final.
+            width (int): Largura do vídeo final.
+            duration_in_seconds (float): Duração total desejada do vídeo.
+            ltx_configs (Optional[Dict]): Configurações avançadas para a geração LTX (steps, guidance, etc.).
+            output_filename_base (str): O nome base para o arquivo de vídeo final.
+
+        Returns:
+            Optional[str]: O caminho do arquivo de vídeo .mp4 gerado, ou None em caso de falha.
+        """
+        t0 = time.time()
+        logging.info(f"Orchestrator starting new job for prompt: '{prompt.splitlines()[0]}...'")
+
+        try:
+            # =================================================================
+            # --- ETAPA 1: PREPARAÇÃO DO INPUT ---
+            # =================================================================
+            # Converte a string do prompt em uma lista de cenas.
+            prompt_list = [line.strip() for line in prompt.splitlines() if line.strip()]
+            if not prompt_list:
+                raise ValueError("O prompt está vazio ou não contém linhas válidas.")
+
+            # Prepara o item de condicionamento inicial, se uma imagem for fornecida.
+            initial_conditioning_items = []
+            if initial_image:
+                logging.info("Preparing initial conditioning item via VAE client...")
+                # Define os parâmetros: aplicar no frame 0 com força total (1.0).
+                conditioning_params = [(0, 1.0)] 
+                # Chama o cliente VAE para fazer o trabalho pesado de conversão de imagem para LatentConditioningItem.
+                initial_conditioning_items = vae_aduc_pipeline(
+                    media=[initial_image],
+                    task='create_conditioning_items',
+                    target_resolution=(height, width),
+                    conditioning_params=conditioning_params
+                )
+                logging.info(f"Successfully created {len(initial_conditioning_items)} conditioning item(s).")
+
+            # =================================================================
+            # --- ETAPA 2: GERAÇÃO DO VÍDEO LATENTE ---
+            # =================================================================
+            logging.info("Submitting job to LTX client for latent video generation...")
+            # Chama o cliente LTX para gerar o tensor latente completo.
+            final_latents, used_seed = ltx_aduc_pipeline(
+                prompt_list=prompt_list,
+                initial_conditioning_items=initial_conditioning_items,
+                height=height,
+                width=width,
+                duration_in_seconds=duration_in_seconds,
+                ltx_configs=ltx_configs
+            )
+
+            if final_latents is None:
+                raise RuntimeError("LTX client failed to generate a latent tensor.")
+            logging.info(f"LTX client returned latent tensor with shape: {final_latents.shape}")
+
+            # =================================================================
+            # --- ETAPA 3: DECODIFICAÇÃO DO LATENTE PARA PIXELS ---
+            # =================================================================
+            logging.info("Submitting job to VAE client for latent-to-pixel decoding...")
+            # Chama o cliente VAE para converter o resultado em um vídeo visível (tensor de pixels).
+            pixel_tensor = vae_aduc_pipeline(
+                media=final_latents,
+                task='decode'
+            )
+
+            if pixel_tensor is None:
+                raise RuntimeError("VAE client failed to decode the latent tensor.")
+            logging.info(f"VAE client returned pixel tensor with shape: {pixel_tensor.shape}")
+
+            # =================================================================
+            # --- ETAPA 4: CODIFICAÇÃO PARA ARQUIVO DE VÍDEO MP4 ---
+            # =================================================================
+            video_filename = f"{output_filename_base}_{int(time.time())}_{used_seed}.mp4"
+            output_path = f"{self.output_dir}/{video_filename}"
+            
+            logging.info(f"Submitting job to VideoEncodeTool to save final MP4 to: {output_path}")
+            # Usa a ferramenta de vídeo para salvar o tensor de pixels no arquivo final.
+            video_encode_tool_singleton.save_video_from_tensor(
+                pixel_5d=pixel_tensor,
+                path=output_path,
+                fps=24
+            )
+
+            total_time = time.time() - t0
+            logging.info(f"🚀🚀🚀 Orchestrator job complete! Video saved to {output_path}. Total time: {total_time:.2f}s")
+            
+            return output_path
+
+        except Exception:
+            logging.error("ORCHESTRATOR FAILED! A critical error occurred during the workflow.", exc_info=True)
+            return None
+
+# ==============================================================================
+# --- INSTÂNCIA SINGLETON DO ORQUESTRADOR ---
+# Este é o ponto de entrada principal que a UI (app.py) irá chamar.
+# ==============================================================================
+try:
+    ltx_aduc_orchestrator = LtxAducOrchestrator()
+except Exception as e:
+    logging.critical("CRITICAL: Failed to initialize the LtxAducOrchestrator.", exc_info=True)
+    ltx_aduc_orchestrator = None