Update api/ltx/vae_aduc_pipeline.py

api/ltx/vae_aduc_pipeline.py

@@ -1,164 +1,175 @@
 # FILE: api/ltx/vae_aduc_pipeline.py
-# DESCRIPTION: A dedicated, "hot" VAE service specialist.
-# It loads the VAE model onto a dedicated GPU (managed by GPUManager)
-# and keeps it in memory to handle all encoding and decoding requests
-# with minimal latency, using the instance pre-loaded by LTXAducManager.
+# 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.
 
-import os
-import sys
-import time
 import logging
-from pathlib import Path
-from typing import List, Union, Tuple
-
+import time
 import torch
-import numpy as np
+import torchvision.transforms.functional as TVF
 from PIL import Image
+from typing import List, Union, Tuple, Literal
+from dataclasses import dataclass
+import os
+import subprocess
+import sys
+from pathlib import Path
 
-# Importa o gerenciador de GPUs e o gerenciador principal do LTX
-from managers.gpu_manager import gpu_manager
-from api.ltx.ltx_aduc_manager import LatentConditioningItem, ltx_aduc_manager
+from api.ltx.ltx_aduc_manager import ltx_aduc_manager
 
-# --- Importações da Arquitetura e do LTX ---
-try:
-    # 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()))
-    
+DEPS_DIR = Path("/data")
+LTX_VIDEO_REPO_DIR = DEPS_DIR / "LTX-Video"
+repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
+if str(LTX_VIDEO_REPO_DIR.resolve()) not in sys.path:
+    sys.path.insert(0, repo_path)
+    print(f"[DEBUG] Repo adicionado ao sys.path: {repo_path}")
     from ltx_video.models.autoencoders.causal_video_autoencoder import CausalVideoAutoencoder
     from ltx_video.models.autoencoders.vae_encode import vae_encode, vae_decode
-except ImportError as e:
-    raise ImportError(f"A crucial import failed for VaeLtxAducPipeline. Check dependencies. Error: {e}")
-
-
-class VaeLtxAducPipeline:
-    _instance = None
-
-    def __new__(cls, *args, **kwargs):
-        if cls._instance is None:
-            cls._instance = super().__new__(cls)
-            cls._instance._initialized = False
-        return cls._instance
-
+    import ltx_video.pipelines.crf_compressor as crf_compressor
+
+# ==============================================================================
+# --- DEFINIÇÕES DE ESTRUTURA E HELPERS (Importadas ou movidas para cá) ---
+# ==============================================================================
+
+@dataclass
+class LatentConditioningItem:
+    """
+    Estrutura de dados para passar latentes condicionados entre serviços.
+    O tensor latente é mantido na CPU para economizar VRAM.
+    """
+    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.
+    """
+    if isinstance(image_input, str):
+        image = Image.open(image_input).convert("RGB")
+    elif isinstance(image_input, Image.Image):
+        image = image_input
+    else:
+        raise ValueError("image_input must be a file path or a PIL Image object")
+
+    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, y_start = (input_width - new_width) // 2, 0
+    else:
+        new_width, new_height = input_width, int(input_width / aspect_ratio_target)
+        x_start, y_start = 0, (input_height - new_height) // 2
+
+    image = image.crop((x_start, y_start, x_start + new_width, y_start + new_height))
+    image = image.resize((target_width, target_height), Image.Resampling.LANCZOS)
+
+    frame_tensor = TVF.to_tensor(image)
+    frame_tensor = TVF.gaussian_blur(frame_tensor, kernel_size=(3, 3))
+    
+    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)
+    
+    frame_tensor = (frame_tensor * 2.0) - 1.0
+    return frame_tensor.unsqueeze(0).unsqueeze(2)
+
+
+# ==============================================================================
+# --- FUNÇÕES DE TRABALHO (Jobs a serem executados no Pool) ---
+# ==============================================================================
+
+def _job_encode_media(vae: CausalVideoAutoencoder, pixel_tensor: torch.Tensor) -> torch.Tensor:
+    """Função de trabalho genérica para codificar um tensor de pixel."""
+    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()
+
+def _job_decode_latent_to_pixels(vae: CausalVideoAutoencoder, latent_tensor: torch.Tensor) -> torch.Tensor:
+    """Função de trabalho para decodificar um tensor latente."""
+    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()
+
+# ==============================================================================
+# --- A CLASSE CLIENTE (Interface Pública) ---
+# ==============================================================================
+
+class VaeAducPipeline:
+    """Cliente de alto nível para orquestrar todas as tarefas de VAE."""
     def __init__(self):
-        if self._initialized: return
-        
-        logging.info("⚙️ Initializing VaeLtxAducPipeline Singleton...")
-        t0 = time.time()
-        
-        # 1. Obter o dispositivo VAE dedicado do gerenciador central
-        self.device = gpu_manager.get_ltx_vae_device()
-        
-        # 2. Obter a referência ao modelo VAE já carregado e posicionado pelo LTXAducManager
-        try:
-            # Esta é a etapa crucial: reutilizamos o pipeline já existente.
-            self.vae = ltx_aduc_manager.get_pipeline().vae
-        except Exception as e:
-            logging.critical(f"Failed to get VAE from LTXAducManager. Is it initialized first? Error: {e}", exc_info=True)
-            raise
-
-        # 3. Confirmação: Garante que o VAE está no dispositivo correto.
-        #    O LTXAducManager já deve ter feito isso, mas esta é uma verificação de segurança.
-        if self.vae.device != self.device:
-            logging.warning(f"VAE device mismatch! Expected {self.device} but found {self.vae.device}. Forcing move.")
-            self.vae.to(self.device)
-        
-        self.vae.eval()
-        self.dtype = self.vae.dtype
-        
-        self._initialized = True
-        logging.info(f"✅ VaeLtxAducPipeline 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."""
-        if isinstance(item, Image.Image):
-            from PIL import ImageOps
-            img = item.convert("RGB")
-            # Redimensiona mantendo a proporção e cortando o excesso
-            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):
-            # Se já for um tensor, apenas garante que está no formato CHW
-            if item.ndim == 4 and item.shape[0] == 1: # Remove dimensão de batch se houver
-                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) # Adiciona B=1 e F=1
-        return (tensor_5d * 2.0) - 1.0
+        logging.info("✅ VAE ADUC Pipeline (Client) initialized and ready to submit jobs.")
+        pass
 
-    @torch.no_grad()
-    def generate_conditioning_items(
+    def __call__(
         self,
-        media_items: List[Union[Image.Image, torch.Tensor]],
-        target_frames: List[int],
-        strengths: List[float],
-        target_resolution: Tuple[int, int]
-    ) -> List[LatentConditioningItem]:
+        media: Union[torch.Tensor, List[Union[Image.Image, torch.Tensor]]],
+        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]]:
         """
-        [FUNÇÃO PRINCIPAL]
-        Converte uma lista de imagens (PIL ou tensores de pixel) em uma lista de
-        LatentConditioningItem, pronta para ser usada pelo pipeline LTX corrigido.
+        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) correspondente a cada item de mídia.
+
+        Returns:
+            O resultado da tarefa, sempre na CPU.
         """
         t0 = time.time()
-        logging.info(f"Generating {len(media_items)} latent conditioning items on device {self.device}...")
-        
-        if not (len(media_items) == len(target_frames) == len(strengths)):
-            raise ValueError("As listas de media_items, target_frames e strengths devem ter o mesmo tamanho.")
-        
-        conditioning_items = []
-        try:
-            for item, frame, strength in zip(media_items, target_frames, strengths):
-                # 1. Prepara a imagem/tensor para o formato de pixel correto
-                pixel_tensor = self._preprocess_input(item, target_resolution)
-                
-                # 2. Move o tensor de pixel para a GPU do VAE e encoda para latente
-                pixel_tensor_gpu = pixel_tensor.to(self.device, dtype=self.dtype)
-                latents = vae_encode(pixel_tensor_gpu, self.vae, vae_per_channel_normalize=True)
-                
-                # 3. Cria o LatentConditioningItem com o latente (movido para CPU para evitar manter na VRAM)
-                conditioning_items.append(LatentConditioningItem(latents.cpu(), frame, strength))
-
-            logging.info(f"Generated {len(conditioning_items)} items in {time.time() - t0:.2f}s.")
+        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 = []
+            for pt in pixel_tensors:
+                latent = ltx_aduc_manager.submit_job(job_type='vae', job_func=_job_encode_media, pixel_tensor=pt)
+                results.append(latent)
+            return results
+
+        elif task == 'decode':
+            if not isinstance(media, torch.Tensor):
+                raise TypeError("Para 'decode', 'media' deve ser um único tensor latente.")
+            return ltx_aduc_manager.submit_job(job_type='vae', job_func=_job_decode_latent_to_pixels, 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
+                ))
             return conditioning_items
-        finally:
-            self._cleanup_gpu()
-
-    @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"Decoded latents with shape {latent_tensor.shape} in {time.time() - t0:.2f}s.")
-            return pixels.cpu() # Retorna na CPU para liberar VRAM da GPU do VAE
-        finally:
-            self._cleanup_gpu()
-
-# --- Instância Singleton ---
-# A inicialização ocorre quando o módulo é importado pela primeira vez.
+        else:
+            raise ValueError(f"Tarefa desconhecida: '{task}'. Opções: 'encode', 'decode', 'create_conditioning_items'.")
+
+# --- INSTÂNCIA SINGLETON DO CLIENTE ---
 try:
-    vae_ltx_aduc_pipeline = VaeLtxAducPipeline()
+    vae_aduc_pipeline = VaeAducPipeline()
 except Exception as e:
-    logging.critical("CRITICAL: Failed to initialize VaeLtxAducPipeline singleton.", exc_info=True)
-    vae_ltx_aduc_pipeline = None
+    logging.critical("CRITICAL: Failed to initialize the VaeAducPipeline client.", exc_info=True)
+    vae_aduc_pipeline = None