api/ltx_pool_manager

# 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