Upload vae_manager (1).py

managers/vae_manager (1).py

@@ -0,0 +1,96 @@
+# FILE: managers/vae_manager.py (Versão Final com vae_decode corrigido)
+
+import torch
+import contextlib
+import logging
+import sys
+from pathlib import Path
+import os
+import io
+
+LTX_VIDEO_REPO_DIR = Path("/data/LTX-Video")
+
+def add_deps_to_path():
+    """
+    Adiciona o diretório do repositório LTX ao sys.path para garantir que suas
+    bibliotecas possam ser importadas.
+    """
+    repo_path = str(LTX_VIDEO_REPO_DIR.resolve())
+    if repo_path not in sys.path:
+        sys.path.insert(0, repo_path)
+        logging.info(f"[ltx_utils] LTX-Video repository added to sys.path: {repo_path}")
+
+# Executa a função imediatamente para configurar o ambiente antes de qualquer importação.
+add_deps_to_path()
+
+
+# --- IMPORTAÇÃO CRÍTICA ---
+# Importa a função helper oficial da biblioteca LTX para decodificação.
+try:
+    from ltx_video.models.autoencoders.vae_encode import vae_decode
+except ImportError:
+    raise ImportError("Could not import 'vae_decode' from LTX-Video library. Check sys.path and repo integrity.")
+
+
+class _SimpleVAEManager:
+    """
+    Manages VAE decoding, now using the official 'vae_decode' helper function
+    for maximum compatibility.
+    """
+    def __init__(self):
+        self.pipeline = None
+        self.device = torch.device("cpu")
+        self.autocast_dtype = torch.float32
+
+    def attach_pipeline(self, pipeline, device=None, autocast_dtype=None):
+        self.pipeline = pipeline
+        if device is not None:
+            self.device = torch.device(device)
+            logging.info(f"[VAEManager] VAE device successfully set to: {self.device}")
+        if autocast_dtype is not None:
+            self.autocast_dtype = autocast_dtype
+
+    @torch.no_grad()
+    def decode(self, latent_tensor: torch.Tensor, decode_timestep: float = 0.05) -> torch.Tensor:
+        """
+        Decodes a latent tensor into a pixel tensor using the 'vae_decode' helper.
+        """
+        if self.pipeline is None:
+            raise RuntimeError("VAEManager: No pipeline has been attached.")
+        
+        # Move os latentes para o dispositivo VAE dedicado.
+        latent_tensor_on_vae_device = latent_tensor.to(self.device)
+
+        # Prepara o tensor de timesteps no mesmo dispositivo.
+        num_items_in_batch = latent_tensor_on_vae_device.shape[0]
+        timestep_tensor = torch.tensor([decode_timestep] * num_items_in_batch, device=self.device)
+    
+        autocast_device_type = self.device.type
+        ctx = torch.autocast(
+            device_type=autocast_device_type,
+            dtype=self.autocast_dtype,
+            enabled=(autocast_device_type == 'cuda')
+        )
+        
+        with ctx:
+            logging.debug(f"[VAEManager] Decoding latents with shape {latent_tensor_on_vae_device.shape} on {self.device}.")
+            
+            # --- CORREÇÃO PRINCIPAL ---
+            # Usa a função helper `vae_decode` em vez de chamar `vae.decode` diretamente.
+            # Esta função sabe como lidar com o argumento 'timestep'.
+            pixels = vae_decode(
+                latents=latent_tensor_on_vae_device,
+                vae=self.pipeline.vae,
+                is_video=True,
+                timestep=timestep_tensor,
+                vae_per_channel_normalize=True, # Importante manter este parâmetro consistente
+            )
+    
+        # A função vae_decode já retorna no intervalo [0, 1], mas um clamp extra não faz mal.
+        pixels = pixels.clamp(0, 1)
+        
+        logging.debug("[VAEManager] Decoding complete. Moving pixel tensor to CPU.")
+        return pixels.cpu()
+
+# Singleton global
+vae_manager_singleton = _SimpleVAEManager()