Update api/ltx/ltx_aduc_manager.py

api/ltx/ltx_aduc_manager.py

@@ -1,148 +1,107 @@
 # FILE: api/ltx/ltx_aduc_manager.py
-# DESCRIPTION: A singleton pool manager for the LTX-Video pipeline.
-# This module is the "secret weapon": it handles loading, device placement,
-# and applies a runtime monkey patch to the LTX pipeline for full control
-# and compatibility with the ADUC-SDR architecture, especially for latent conditioning.
-
+# DESCRIPTION: A singleton manager for the LTX-Video pipeline.
+# This module loads the pipeline, places it on the correct devices, and applies a
+# targeted runtime monkey patch to delegate conditioning tasks to the specialized
+# VaeAducPipeline service, enabling full control for the ADUC-SDR architecture.
 
 import time
-import os
 import yaml
-import json
 from pathlib import Path
 from typing import List, Optional, Tuple, Union, Dict
-from dataclasses import dataclass
 import threading
 import sys
-from pathlib import Path
 import torch
-from diffusers.utils.torch_utils import randn_tensor
-from huggingface_hub import hf_hub_download
 
-# --- Importações da nossa arquitetura ---
+# --- Importações da arquitetura ADUC-SDR ---
 from managers.gpu_manager import gpu_manager
 from api.ltx.ltx_utils import build_ltx_pipeline_on_cpu
 from utils.debug_utils import log_function_io
-
-LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
-LTX_REPO_ID = "Lightricks/LTX-Video"
-CACHE_DIR = os.environ.get("HF_HOME")
+# Importa o serviço VAE que fará o trabalho real
+from api.ltx.vae_aduc_pipeline import vae_aduc_pipeline, LatentConditioningItem
 
 # --- Importações da biblioteca LTX-Video ---
+LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
 repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
 if repo_path not in sys.path:
    sys.path.insert(0, repo_path)
 
 from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline
-from ltx_video.models.autoencoders.vae_encode import vae_encode, latent_to_pixel_coords
+# Importa o tipo original de conditioning item para type hinting
+from ltx_video.pipelines.pipeline_ltx_video import ConditioningItem as PipelineConditioningItem
 import logging
 
 import warnings
 warnings.filterwarnings("ignore", category=UserWarning)
 warnings.filterwarnings("ignore", category=FutureWarning)
 warnings.filterwarnings("ignore", message=".*")
-from huggingface_hub import logging as ll
-ll.set_verbosity_error()
-ll.set_verbosity_warning()
-ll.set_verbosity_info()
-ll.set_verbosity_debug()
+
+try:
+    from huggingface_hub import logging as hf_logging
+    hf_logging.set_verbosity_error()
+except ImportError:
+    pass
 
 logger = logging.getLogger("AducDebug")
 logging.basicConfig(level=logging.DEBUG)
 logger.setLevel(logging.DEBUG)
 
-# ==============================================================================
-# --- DEFINIÇÃO DOS DATACLASSES DE CONDICIONAMENTO ADUC-SDR ---
-# ==============================================================================
-
-@dataclass
-class ConditioningItem:
-    """Nosso Data Class para condicionamento com TENSORES DE PIXEL (de imagens)."""
-    pixel_tensor: torch.Tensor
-    media_frame_number: int
-    conditioning_strength: float
-
-@dataclass
-class LatentConditioningItem:
-    """Nossa "arma secreta": um Data Class para condicionamento com TENSORES LATENTES (de overlap)."""
-    latent_tensor: torch.Tensor
-    media_frame_number: int
-    conditioning_strength: float
 
 # ==============================================================================
-# --- O MONKEY PATCH ---
-# Nossa versão customizada de `prepare_conditioning` que entende ambos os Data Classes.
+# --- O MONKEY PATCH DIRECIONADO E SIMPLES ---
 # ==============================================================================
 
 @log_function_io
 def _aduc_prepare_conditioning_patch(
     self: "LTXVideoPipeline",
-    conditioning_items: Optional[List[Union[ConditioningItem, LatentConditioningItem]]],
+    conditioning_items: Optional[List[Union[PipelineConditioningItem, LatentConditioningItem]]],
     init_latents: torch.Tensor,
-    num_frames: int, height: int, width: int, # Assinatura mantida para compatibilidade
+    num_frames: int,
+    height: int,
+    width: int,
     vae_per_channel_normalize: bool = False,
     generator=None,
-) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
-    
-    if not conditioning_items:
-        latents, latent_coords = self.patchifier.patchify(latents=init_latents)
-        pixel_coords = latent_to_pixel_coords(latent_coords, self.vae, causal_fix=self.transformer.config.causal_temporal_positioning)
-        return latents, pixel_coords, None, 0
-    
-    init_conditioning_mask = torch.zeros_like(init_latents[:, 0, ...], dtype=torch.float32, device=init_latents.device)
-    extra_conditioning_latents, extra_conditioning_pixel_coords, extra_conditioning_mask = [], [], []
-    extra_conditioning_num_latents = 0
-    
-    for item in conditioning_items:
-        strength = item.conditioning_strength
-        media_frame_number = item.media_frame_number
-        
-        if isinstance(item, ConditioningItem):
-            logging.debug("Patch ADUC: Processando ConditioningItem (pixels).")
-            pixel_tensor_on_vae_device = item.pixel_tensor.to(device=self.vae.device, dtype=self.vae.dtype)
-            media_item_latents = vae_encode(pixel_tensor_on_vae_device, self.vae, vae_per_channel_normalize=vae_per_channel_normalize)
-            media_item_latents = media_item_latents.to(device=init_latents.device, dtype=init_latents.dtype)
-        elif isinstance(item, LatentConditioningItem):
-            logging.debug("Patch ADUC: Processando LatentConditioningItem (latentes).")
-            media_item_latents = item.latent_tensor.to(device=init_latents.device, dtype=init_latents.dtype)
-        else:
-            logging.warning(f"Patch ADUC: Item de condicionamento de tipo desconhecido '{type(item)}' será ignorado.")
-            continue
-        
-        if media_frame_number == 0:
-            f_l, h_l, w_l = media_item_latents.shape[-3:]
-            init_latents[..., :f_l, :h_l, :w_l] = torch.lerp(init_latents[..., :f_l, :h_l, :w_l], media_item_latents, strength)
-            init_conditioning_mask[..., :f_l, :h_l, :w_l] = strength
-        else:
-            noise = randn_tensor(media_item_latents.shape, generator=generator, device=media_item_latents.device, dtype=media_item_latents.dtype)
-            media_item_latents = torch.lerp(noise, media_item_latents, strength)
-            patched_latents, latent_coords = self.patchifier.patchify(latents=media_item_latents)
-            pixel_coords = latent_to_pixel_coords(latent_coords, self.vae, causal_fix=self.transformer.config.causal_temporal_positioning)
-            pixel_coords[:, 0] += media_frame_number
-            extra_conditioning_num_latents += patched_latents.shape[1]
-            new_mask = torch.full(patched_latents.shape[:2], strength, dtype=torch.float32, device=init_latents.device)
-            extra_conditioning_latents.append(patched_latents)
-            extra_conditioning_pixel_coords.append(pixel_coords)
-            extra_conditioning_mask.append(new_mask)
-
-    init_latents, init_latent_coords = self.patchifier.patchify(latents=init_latents)
-    init_pixel_coords = latent_to_pixel_coords(init_latent_coords, self.vae, causal_fix=self.transformer.config.causal_temporal_positioning)
-    init_conditioning_mask, _ = self.patchifier.patchify(latents=init_conditioning_mask.unsqueeze(1))
-    init_conditioning_mask = init_conditioning_mask.squeeze(-1)
-    
-    if extra_conditioning_latents:
-        init_latents = torch.cat([*extra_conditioning_latents, init_latents], dim=1)
-        init_pixel_coords = torch.cat([*extra_conditioning_pixel_coords, init_pixel_coords], dim=2)
-        init_conditioning_mask = torch.cat([*extra_conditioning_mask, init_conditioning_mask], dim=1)
-    
-    return init_latents, init_pixel_coords, init_conditioning_mask, extra_conditioning_num_latents
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], int]:
+    """
+    [PATCH] Substitui o método `prepare_conditioning` original da LTXVideoPipeline.
+
+    Esta função atua como um proxy (intermediário). Ela não contém lógica de processamento.
+    Em vez disso, ela delega 100% do trabalho para o `vae_aduc_pipeline`, que é o nosso
+    serviço especializado e otimizado para essa tarefa.
+    """
+    logging.debug(f"Patch ADUC: Interceptado 'prepare_conditioning'. Delegando para o serviço VaeAducPipeline.")
+
+    # 1. Chama o serviço especializado para fazer todo o trabalho pesado.
+    #    O serviço VAE processa na sua própria GPU dedicada e retorna os tensores na CPU.
+    latents_cpu, coords_cpu, mask_cpu, num_latents = vae_aduc_pipeline.prepare_conditioning(
+        conditioning_items=conditioning_items,
+        init_latents=init_latents,
+        num_frames=num_frames,
+        height=height,
+        width=width,
+        vae_per_channel_normalize=vae_per_channel_normalize,
+        generator=generator,
+    )
+
+    # 2. Move os resultados da CPU para o dispositivo correto que a pipeline principal espera.
+    #    O `init_latents.device` garante que estamos usando o dispositivo principal da pipeline (ex: 'cuda:0').
+    device = init_latents.device
+    latents = latents_cpu.to(device)
+    pixel_coords = coords_cpu.to(device)
+    conditioning_mask = mask_cpu.to(device) if mask_cpu is not None else None
+
+    # 3. Retorna os tensores prontos. A pipeline principal continua sua execução normalmente,
+    #    sem saber que a lógica de condicionamento foi executada por um serviço externo.
+    return latents, pixel_coords, conditioning_mask, num_latents
 
 # ==============================================================================
 # --- LTX WORKER E POOL MANAGER ---
 # ==============================================================================
 
 class LTXWorker:
-    """Gerencia uma instância do LTX Pipeline em um par de GPUs (main + vae)."""
+    """
+    Gerencia uma instância única da LTXVideoPipeline, aplicando o patch
+    necessário durante a inicialização.
+    """
     def __init__(self, main_device_str: str, vae_device_str: str, config: dict):
         self.main_device = torch.device(main_device_str)
         self.vae_device = torch.device(vae_device_str)
@@ -152,16 +111,27 @@ class LTXWorker:
 
     @log_function_io
     def _load_and_patch_pipeline(self):
+        """
+        Orquestra o carregamento da pipeline e a aplicação do monkey patch.
+        """
         logging.info(f"[LTXWorker-{self.main_device}] Carregando pipeline LTX para a CPU...")
         self.pipeline, _ = build_ltx_pipeline_on_cpu(self.config)
+        
         logging.info(f"[LTXWorker-{self.main_device}] Movendo pipeline para GPUs (Main: {self.main_device}, VAE: {self.vae_device})...")
-        self.pipeline.to(self.main_device)
-        self.pipeline.vae.to(self.vae_device)
-        logging.info(f"[LTXWorker-{self.main_device}] Aplicando patch ADUC-SDR na função 'prepare_conditioning'...")
+        self.pipeline.to(self.main_device)  # Move a maioria dos componentes
+        self.pipeline.vae.to(self.vae_device) # Move o VAE para sua GPU dedicada
+        
+        logging.info(f"[LTXWorker-{self.main_device}] Aplicando patch ADUC-SDR em 'prepare_conditioning'...")
+        # A "mágica" simples e eficaz acontece aqui:
         self.pipeline.prepare_conditioning = _aduc_prepare_conditioning_patch.__get__(self.pipeline, LTXVideoPipeline)
-        logging.info(f"[LTXWorker-{self.main_device}] ✅ Pipeline 'quente', corrigido e pronto para uso.")
+        
+        logging.info(f"[LTXWorker-{self.main_device}] ✅ Pipeline 'quente', corrigida e pronta para uso.")
 
 class LtxAducManager:
+    """
+    Implementa o padrão Singleton para garantir que a pipeline LTX seja
+    carregada e corrigida apenas uma vez durante a vida útil da aplicação.
+    """
     _instance = None
     _lock = threading.Lock()
 
@@ -173,28 +143,38 @@ class LtxAducManager:
             return cls._instance
 
     def __init__(self):
-        if self._initialized: return
+        if hasattr(self, '_initialized') and self._initialized:
+            return
         with self._lock:
-            if self._initialized: return
-            logging.info("⚙️ Inicializando LTXPoolManager Singleton...")
+            if hasattr(self, '_initialized') and self._initialized:
+                return
+            logging.info("⚙️ Inicializando LtxAducManager Singleton...")
             self.config = self._load_config()
             main_device_str = str(gpu_manager.get_ltx_device())
             vae_device_str = str(gpu_manager.get_ltx_vae_device())
+            
+            # Cria o worker que irá carregar e patchear a pipeline
             self.worker = LTXWorker(main_device_str, vae_device_str, self.config)
+            
             self._initialized = True
-            logging.info("✅ LTXPoolManager pronto.")
+            logging.info("✅ LtxAducManager pronto.")
 
-    @log_function_io
     def _load_config(self) -> Dict:
         """Carrega a configuração YAML principal do LTX."""
+        # TODO: Considerar mover o path da configuração para uma variável de ambiente ou config central
         config_path = Path("/data/LTX-Video/configs/ltxv-13b-0.9.8-dev-fp8.yaml")
         with open(config_path, "r") as file:
             return yaml.safe_load(file)
 
-    @log_function_io
     def get_pipeline(self) -> LTXVideoPipeline:
-        """Retorna a instância do pipeline, já carregada e corrigida."""
+        """
+        Ponto de acesso principal para obter a instância da pipeline.
+
+        Returns:
+            LTXVideoPipeline: A instância única, carregada e já corrigida.
+        """
         return self.worker.pipeline
 
 # --- Instância Singleton Global ---
-ltx_aduc_manager = LtxAducManager()
+# Outras partes do código importarão esta instância para interagir com a pipeline.
+ltx_aduc_manager = LtxAducManager()