Update api/ltx/vae_aduc_pipeline.py

api/ltx/vae_aduc_pipeline.py

@@ -1,177 +1,164 @@
 # FILE: api/ltx/vae_aduc_pipeline.py
-# DESCRIPTION: A high-level client for submitting VAE-related jobs to the LTXAducManager pool.
-# It handles encoding media to latents, decoding latents to pixels, and creating ConditioningItems.
+# DESCRIPTION: A dedicated, "hot" VAE service specialist.
+# It holds the VAE model on a dedicated GPU and provides high-level services
+# for encoding images/tensors into conditioning items and decoding latents back to pixels.
 
-import logging
-import time
-import torch
 import os
-import torchvision.transforms.functional as TVF
-from PIL import Image
-from typing import List, Union, Tuple, Literal, Optional
-from dataclasses import dataclass
-from pathlib import Path
 import sys
+import time
+import logging
+import threading
+from pathlib import Path
+from typing import List, Union, Tuple
 
-# O cliente importa o MANAGER para submeter os trabalhos ao pool de workers.
-from api.ltx.ltx_aduc_manager import ltx_aduc_manager
-
-# --- Adiciona o path do LTX-Video para importações de baixo nível ---
-LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
-def add_deps_to_path():
-    repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
-    if repo_path not in sys.path:
-        sys.path.insert(0, repo_path)
-add_deps_to_path()
-
-# Importações para anotação de tipos e para as funções de trabalho (jobs).
-from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
-from ltx_video.models.autoencoders.vae_encode import vae_encode, vae_decode
-import ltx_video.pipelines.crf_compressor as crf_compressor
+import torch
+import numpy as np
+from PIL import Image
 
 # ==============================================================================
-# --- DEFINIÇÕES DE ESTRUTURA E HELPERS ---
+# --- IMPORTAÇÕES DA ARQUITETURA E DO LTX ---
 # ==============================================================================
-
-@dataclass
-class LatentConditioningItem:
-    """
-    Estrutura de dados para passar latentes condicionados entre serviços.
-    O tensor latente é mantido na CPU para economizar VRAM entre as etapas.
-    """
-    latent_tensor: torch.Tensor
-    media_frame_number: int
-    conditioning_strength: float
-
-def load_image_to_tensor_with_resize_and_crop(
-    image_input: Union[str, Image.Image],
-    target_height: int,
-    target_width: int,
-) -> torch.Tensor:
-    """
-    Carrega e processa uma imagem para um tensor de pixel 5D, normalizado para [-1, 1],
-    pronto para ser enviado ao VAE para encoding.
-    """
-    if isinstance(image_input, str):
-        image = Image.open(image_input).convert("RGB")
-    elif isinstance(image_input, Image.Image):
-        image = image_input.convert("RGB")
-    else:
-        raise ValueError("image_input must be a file path or a PIL Image object")
-
-    # Lógica de corte e redimensionamento para manter a proporção
-    input_width, input_height = image.size
-    aspect_ratio_target = target_width / target_height
-    aspect_ratio_frame = input_width / input_height
-    if aspect_ratio_frame > aspect_ratio_target:
-        new_width, new_height = int(input_height * aspect_ratio_target), input_height
-        x_start = (input_width - new_width) // 2
-        image = image.crop((x_start, 0, x_start + new_width, new_height))
-    else:
-        new_height = int(input_width / aspect_ratio_target)
-        y_start = (input_height - new_height) // 2
-        image = image.crop((0, y_start, input_width, y_start + new_height))
-    
-    image = image.resize((target_width, target_height), Image.Resampling.LANCZOS)
-    
-    # Conversão para tensor e normalização
-    frame_tensor = TVF.to_tensor(image)
-    frame_tensor_hwc = frame_tensor.permute(1, 2, 0)
-    frame_tensor_hwc = crf_compressor.compress(frame_tensor_hwc)
-    frame_tensor = frame_tensor_hwc.permute(2, 0, 1)
+try:
+    from api.ltx.ltx_aduc_manager import LatentConditioningItem
+    from api.managers.gpu_manager import gpu_manager
+    # Adiciona o path para as bibliotecas do LTX
+    LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
+    if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
+        sys.path.insert(0, str(LTX_VIDEO_REPO_DIR.resolve()))
     
-    frame_tensor = (frame_tensor * 2.0) - 1.0
-    return frame_tensor.unsqueeze(0).unsqueeze(2)
+    from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
+    from ltx_video.models.autoencoders.vae_encode import vae_encode, vae_decode
+    # Nossos data classes customizados para condicionamento, importados do pool manager
+    except ImportError as e:
+    raise ImportError(f"A crucial import failed for VaeServer. Check dependencies. Error: {e}")
 
 # ==============================================================================
-# --- FUNÇÕES DE TRABALHO (Jobs a serem executados no Pool de VAE) ---
+# --- CLASSE DO SERVIÇO VAE ---
 # ==============================================================================
 
-def _job_encode_media(vae: CausalVideoAutoencoder, pixel_tensor: torch.Tensor) -> torch.Tensor:
-    """Job que codifica um tensor de pixel em um tensor latente."""
-    device = vae.device
-    dtype = vae.dtype
-    pixel_tensor_gpu = pixel_tensor.to(device, dtype=dtype)
-    latents = vae_encode(pixel_tensor_gpu, vae, vae_per_channel_normalize=True)
-    return latents.cpu()
+class VaeServer:
+    _instance = None
+    _lock = threading.Lock()
 
-def _job_decode_latent(vae: CausalVideoAutoencoder, latent_tensor: torch.Tensor) -> torch.Tensor:
-    """Job que decodifica um tensor latente em um tensor de pixels."""
-    device = vae.device
-    dtype = vae.dtype
-    latent_tensor_gpu = latent_tensor.to(device, dtype=dtype)
-    pixels = vae_decode(latent_tensor_gpu, vae, is_video=True, vae_per_channel_normalize=True)
-    return pixels.cpu()
+    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
 
-# ==============================================================================
-# --- A CLASSE CLIENTE (Interface Pública) ---
-# ==============================================================================
-
-class VaeAducPipeline:
-    """
-    Cliente de alto nível para orquestrar todas as tarefas relacionadas ao VAE.
-    Ele define a lógica de negócios e submete os trabalhos ao LTXAducManager.
-    """
     def __init__(self):
-        logging.info("✅ VAE ADUC Pipeline (Client) initialized and ready to submit jobs.")
-        pass
+        if self._initialized: return
+        with self._lock:
+            if self._initialized: return
+        
+            logging.info("⚙️ Initializing VaeServer Singleton...")
+            t0 = time.time()
+            
+            # 1. Obter o dispositivo VAE dedicado do gerenciador central
+            self.device = gpu_manager.get_ltx_vae_device()
+            
+            # 2. Obter o modelo VAE já carregado pelo LTXPoolManager
+            #    Isso garante consistência e evita carregar o modelo duas vezes.
+            try:
+                from api.ltx.ltx_aduc_manager import ltx_pool_manager
+                if ltx_pool_manager is None or ltx_pool_manager.get_pipeline() is None:
+                    raise RuntimeError("LTXPoolManager is not initialized yet. VaeServer must be initialized after.")
+                self.vae = ltx_pool_manager.get_pipeline().vae
+            except Exception as e:
+                logging.critical(f"Failed to get VAE from LTXPoolManager. Error: {e}", exc_info=True)
+                raise
+
+            # 3. Garante que o VAE está no dispositivo correto e em modo de avaliação
+            self.vae.to(self.device)
+            self.vae.eval()
+            self.dtype = self.vae.dtype
+            
+            self._initialized = True
+            logging.info(f"✅ VaeServer ready. VAE model is 'hot' on {self.device} with dtype {self.dtype}. Startup time: {time.time() - t0:.2f}s")
+
+    def _cleanup_gpu(self):
+        """Limpa a VRAM da GPU do VAE."""
+        if torch.cuda.is_available():
+            with torch.cuda.device(self.device):
+                torch.cuda.empty_cache()
+
+    def _preprocess_input(self, item: Union[Image.Image, torch.Tensor], target_resolution: Tuple[int, int]) -> torch.Tensor:
+        """Prepara uma imagem PIL ou um tensor para o formato de pixel que o VAE espera para encodar."""
+        if isinstance(item, Image.Image):
+            from PIL import ImageOps
+            img = item.convert("RGB")
+            processed_img = ImageOps.fit(img, target_resolution, Image.Resampling.LANCZOS)
+            image_np = np.array(processed_img).astype(np.float32) / 255.0
+            tensor = torch.from_numpy(image_np).permute(2, 0, 1) # HWC -> CHW
+        elif isinstance(item, torch.Tensor):
+            if item.ndim == 4 and item.shape[0] == 1: tensor = item.squeeze(0)
+            elif item.ndim == 3: tensor = item
+            else: raise ValueError(f"Input tensor must have 3 or 4 dimensions (CHW or BCHW), but got {item.ndim}")
+        else:
+            raise TypeError(f"Input must be a PIL Image or a torch.Tensor, but got {type(item)}")
+
+        # Converte para 5D (B, C, F, H, W) e normaliza para [-1, 1]
+        tensor_5d = tensor.unsqueeze(0).unsqueeze(2)
+        return (tensor_5d * 2.0) - 1.0
 
-    def __call__(
+    @torch.no_grad()
+    def generate_conditioning_items(
         self,
-        media: Union[torch.Tensor, List[Union[Image.Image, str]]],
-        task: Literal['encode', 'decode', 'create_conditioning_items'],
-        target_resolution: Optional[Tuple[int, int]] = (512, 512),
-        conditioning_params: Optional[List[Tuple[int, float]]] = None
-    ) -> Union[List[torch.Tensor], torch.Tensor, List[LatentConditioningItem]]:
+        media_items: List[Union[Image.Image, torch.Tensor]],
+        target_frames: List[int],
+        strengths: List[float],
+        target_resolution: Tuple[int, int]
+    ) -> List[LatentConditioningItem]:
         """
-        Ponto de entrada principal para executar tarefas de VAE.
-
-        Args:
-            media: O dado de entrada.
-            task: A tarefa a executar ('encode', 'decode', 'create_conditioning_items').
-            target_resolution: A resolução (altura, largura) para o pré-processamento.
-            conditioning_params: Para 'create_conditioning_items', uma lista de tuplas
-                                 (frame_number, strength) para cada item de mídia.
-
-        Returns:
-            O resultado da tarefa, sempre na CPU.
+        [FUNÇÃO PRINCIPAL] Converte uma lista de imagens (PIL ou tensores de pixel)
+        em uma lista de LatentConditioningItem, pronta para a pipeline LTX corrigida.
         """
         t0 = time.time()
-        logging.info(f"VAE Client received a '{task}' job.")
-
-        if task == 'encode':
-            if not isinstance(media, list): media = [media]
-            pixel_tensors = [load_image_to_tensor_with_resize_and_crop(m, target_resolution[0], target_resolution[1]) for m in media]
-            results = [ltx_aduc_manager.submit_job(job_type='vae', job_func=_job_encode_media, pixel_tensor=pt) for pt in pixel_tensors]
-            return results
-
-        elif task == 'decode':
-            if not isinstance(media, torch.Tensor):
-                raise TypeError("Para a tarefa 'decode', 'media' deve ser um único tensor latente.")
-            return ltx_aduc_manager.submit_job(job_type='vae', job_func=_job_decode_latent, latent_tensor=media)
-
-        elif task == 'create_conditioning_items':
-            if not isinstance(media, list) or not isinstance(conditioning_params, list) or len(media) != len(conditioning_params):
-                raise ValueError("Para 'create_conditioning_items', 'media' e 'conditioning_params' devem ser listas de mesmo tamanho.")
-            
-            pixel_tensors = [load_image_to_tensor_with_resize_and_crop(m, target_resolution[0], target_resolution[1]) for m in media]
-            conditioning_items = []
-            for i, pt in enumerate(pixel_tensors):
-                latent_tensor = ltx_aduc_manager.submit_job(job_type='vae', job_func=_job_encode_media, pixel_tensor=pt)
-                frame_number, strength = conditioning_params[i]
-                conditioning_items.append(LatentConditioningItem(
-                    latent_tensor=latent_tensor,
-                    media_frame_number=frame_number,
-                    conditioning_strength=strength
-                ))
+        logging.info(f"VaeServer: Generating {len(media_items)} latent conditioning items...")
+        
+        if not (len(media_items) == len(target_frames) == len(strengths)):
+            raise ValueError("Input lists for conditioning items must have the same length.")
+        
+        conditioning_items = []
+        try:
+            for item, frame, strength in zip(media_items, target_frames, strengths):
+                pixel_tensor = self._preprocess_input(item, target_resolution)
+                pixel_tensor_gpu = pixel_tensor.to(self.device, dtype=self.dtype)
+                latents = vae_encode(pixel_tensor_gpu, self.vae, vae_per_channel_normalize=True)
+                conditioning_items.append(LatentConditioningItem(latents.cpu(), frame, strength))
+
+            logging.info(f"VaeServer: Generated {len(conditioning_items)} items in {time.time() - t0:.2f}s.")
             return conditioning_items
-            
-        else:
-            raise ValueError(f"Tarefa desconhecida: '{task}'. Opções: 'encode', 'decode', 'create_conditioning_items'.")
+        finally:
+            self._cleanup_gpu()
 
-# --- INSTÂNCIA SINGLETON DO CLIENTE ---
+    @torch.no_grad()
+    def decode_to_pixels(self, latent_tensor: torch.Tensor, decode_timestep: float = 0.05) -> torch.Tensor:
+        """Decodifica um tensor latente para um tensor de pixels, retornando na CPU."""
+        t0 = time.time()
+        try:
+            latent_tensor_gpu = latent_tensor.to(self.device, dtype=self.dtype)
+            num_items_in_batch = latent_tensor_gpu.shape[0]
+            timestep_tensor = torch.tensor([decode_timestep] * num_items_in_batch, device=self.device, dtype=self.dtype)
+            
+            pixels = vae_decode(
+                latent_tensor_gpu, self.vae, is_video=True, 
+                timestep=timestep_tensor, vae_per_channel_normalize=True
+            )
+            logging.info(f"VaeServer: Decoded latents with shape {latent_tensor.shape} in {time.time() - t0:.2f}s.")
+            return pixels.cpu()
+        finally:
+            self._cleanup_gpu()
+
+# --- Instância Singleton ---
 try:
-    vae_aduc_pipeline = VaeAducPipeline()
+    # A inicialização depende do LTXPoolManager para obter o VAE
+    from api.ltx.ltx_aduc_manager import ltx_pool_manager
+    if ltx_pool_manager:
+        vae_server_singleton = VaeServer()
+    else:
+        raise RuntimeError("LTXPoolManager failed to initialize, cannot start VaeServer.")
 except Exception as e:
-    logging.critical("CRITICAL: Failed to initialize the VaeAducPipeline client.", exc_info=True)
-    vae_aduc_pipeline = None
+    logging.critical("CRITICAL: Failed to initialize VaeServer singleton.", exc_info=True)
+    vae_server_singleton = None