Create ltx_pool_manager

api/ltx_pool_manager

@@ -0,0 +1,208 @@
+# FILE: api/ltx_pool_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.
+
+import logging
+import time
+import os
+import yaml
+import json
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+from dataclasses import dataclass
+
+import torch
+from diffusers.utils.torch_utils import randn_tensor
+from huggingface_hub import hf_hub_download
+
+# --- Importações da nossa arquitetura ---
+from api.gpu_manager import gpu_manager
+
+# --- Importações da biblioteca LTX-Video e Utilitários ---
+from api.ltx.ltx_utils import build_ltx_pipeline_on_cpu
+from ltx_video.pipelines.pipeline_ltx_video import LTXVideoPipeline
+from ltx_video.models.autoencoders.vae_encode import vae_encode, latent_to_pixel_coords
+
+# ==============================================================================
+# --- DEFINIÇÃO DOS NOSSOS DATACLASSES DE CONDICIONAMENTO ---
+# ==============================================================================
+
+@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.
+# ==============================================================================
+
+def _aduc_prepare_conditioning_patch(
+    self: "LTXVideoPipeline",
+    conditioning_items: Optional[List[Union[ConditioningItem, LatentConditioningItem]]],
+    init_latents: torch.Tensor,
+    num_frames: int, height: int, width: int, # Assinatura mantida para compatibilidade
+    vae_per_channel_normalize: bool = False,
+    generator=None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+    
+    # Se não houver itens, apenas "patchify" os latentes iniciais e retorna.
+    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
+    
+    # Prepara máscaras e listas para acumular os tensores de condição.
+    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
+        
+        # --- LÓGICA PRINCIPAL DO PATCH ---
+        if isinstance(item, ConditioningItem):
+            # Item é um tensor de PIXEL (ex: imagem inicial).
+            logging.debug("Patch ADUC: Processando ConditioningItem (pixels).")
+            # Encodifica o tensor de pixel para o espaço latente usando o VAE.
+            # Garante que a operação ocorra no dispositivo do VAE para evitar erros.
+            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)
+            # Traz o resultado de volta para o dispositivo principal (do Transformer).
+            media_item_latents = media_item_latents.to(device=init_latents.device, dtype=init_latents.dtype)
+
+        elif isinstance(item, LatentConditioningItem):
+            # Item já é um tensor LATENTE (ex: overlap de chunks).
+            logging.debug("Patch ADUC: Processando LatentConditioningItem (latentes).")
+            # Apenas garante que o tensor está no dispositivo e tipo corretos.
+            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
+        
+        # Lógica original da pipeline, agora operando sobre `media_item_latents` garantido.
+        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: # Condicionamento em frames intermediários
+            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)
+
+    # Finaliza o processo de patchifying e concatenação dos tensores.
+    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
+
+# ==============================================================================
+# --- LTX WORKER E POOL MANAGER ---
+# ==============================================================================
+
+class LTXWorker:
+    """Gerencia uma instância do LTX Pipeline em um par de GPUs (main + vae)."""
+    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)
+        self.config = config
+        self.pipeline: LTXVideoPipeline = None
+        
+        self._load_and_patch_pipeline()
+
+    def _load_and_patch_pipeline(self):
+        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'...")
+        # Substitui o método da instância pelo nosso patch
+        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.")
+
+class LTXPoolManager:
+    _instance = None
+    _lock = threading.Lock()
+
+    def __new__(cls, *args, **kwargs):
+        with cls._lock:
+            if cls._instance is None:
+                cls._instance = super().__new__(cls)
+                cls._instance._initialized = False
+            return cls._instance
+
+    def __init__(self):
+        if self._initialized: return
+        with self._lock:
+            if self._initialized: return
+            
+            logging.info("⚙️ Inicializando LTXPoolManager Singleton...")
+            self.config = self._load_config()
+            self._resolve_model_paths_from_cache()
+
+            main_device_str = str(gpu_manager.get_ltx_device())
+            vae_device_str = str(gpu_manager.get_ltx_vae_device())
+            
+            self.worker = LTXWorker(main_device_str, vae_device_str, self.config)
+            
+            self._initialized = True
+            logging.info("✅ LTXPoolManager pronto.")
+
+    def _load_config(self) -> Dict:
+        """Carrega a configuração YAML principal do LTX."""
+        config_path = Path("/data/LTX-Video/configs/ltxv-13b-0.9.8-distilled-fp8.yaml")
+        with open(config_path, "r") as file:
+            return yaml.safe_load(file)
+
+    def _resolve_model_paths_from_cache(self):
+        """Garante que a configuração em memória tenha os caminhos absolutos para os modelos no cache."""
+        try:
+            main_ckpt_path = hf_hub_download(repo_id="Lightricks/LTX-Video", filename=self.config["checkpoint_path"])
+            self.config["checkpoint_path"] = main_ckpt_path
+            if self.config.get("spatial_upscaler_model_path"):
+                upscaler_path = hf_hub_download(repo_id="Lightricks/LTX-Video", filename=self.config["spatial_upscaler_model_path"])
+                self.config["spatial_upscaler_model_path"] = upscaler_path
+        except Exception as e:
+            logging.critical(f"Falha ao resolver caminhos de modelo LTX. O setup.py foi executado? Erro: {e}", exc_info=True)
+            raise
+
+    def get_pipeline(self) -> LTXVideoPipeline:
+        """Retorna a instância do pipeline, já carregada e corrigida."""
+        return self.worker.pipeline
+
+# --- Instância Singleton Global ---
+# A aplicação importará esta instância para interagir com o LTX.
+try:
+    ltx_pool_manager = LTXPoolManager()
+except Exception as e:
+    logging.critical("Falha crítica ao inicializar o LTXPoolManager.", exc_info=True)
+    ltx_pool_manager = None