Update upscaler_specialist.py

upscaler_specialist.py

@@ -1,31 +1,32 @@
 # upscaler_specialist.py
-# Copyright (C) 2025 Carlos Rodrigues dos Santos
+# Copyright (C) 2025 Carlos Rodrigues
 # Especialista ADUC para upscaling espacial de tensores latentes.
 
 import torch
 import logging
 from diffusers import LTXLatentUpsamplePipeline
-
 from ltx_manager_helpers import ltx_manager_singleton
 
 logger = logging.getLogger(__name__)
 
-
 class UpscalerSpecialist:
-    def __init__(self, device="cuda"):
+    """
+    Especialista responsável por aumentar a resolução espacial de tensores latentes
+    usando o LTX Video Spatial Upscaler.
+    """
+    def __init__(self):
+        # Força uso de CUDA se disponível
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
-        self.cpu_device = torch.device("cpu")
-        self.workspace_dir = workspace_dir
-        self.pipe_upsample = None
         self.base_vae = None
+        self.pipe_upsample = None
 
     def _lazy_init(self):
-        """Inicializa o VAE e o pipeline somente quando for chamado."""
+        """Inicializa VAE e pipeline apenas quando necessário."""
         if self.base_vae is None:
             try:
-                from ltx_manager_helpers import ltx_manager_singleton
                 if ltx_manager_singleton.workers:
                     self.base_vae = ltx_manager_singleton.workers[0].pipeline.vae
+                    logger.info("[Upscaler] VAE base obtido com sucesso.")
                 else:
                     logger.warning("[Upscaler] Nenhum worker disponível no ltx_manager_singleton.")
             except Exception as e:
@@ -34,7 +35,6 @@ class UpscalerSpecialist:
 
         if self.pipe_upsample is None and self.base_vae is not None:
             try:
-                from ltx_video.pipelines.latent_upscale import LTXLatentUpsamplePipeline
                 self.pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained(
                     "linoyts/LTX-Video-spatial-upscaler-0.9.8",
                     vae=self.base_vae,
@@ -44,25 +44,25 @@ class UpscalerSpecialist:
             except Exception as e:
                 logger.error(f"[Upscaler] Falha ao carregar pipeline: {e}")
 
+    @torch.no_grad()
     def upscale(self, latents: torch.Tensor) -> torch.Tensor:
+        """Aplica o upscaling 2x nos tensores latentes fornecidos."""
         self._lazy_init()
         if self.pipe_upsample is None:
             logger.warning("[Upscaler] Pipeline indisponível. Retornando latentes originais.")
             return latents
+
         try:
-            with torch.no_grad():
-                result = self.pipe_upsample(latents=latents, output_type="latent")
+            logger.info(f"[Upscaler] Recebido shape {latents.shape}. Executando upscale em {self.device}...")
+            result = self.pipe_upsample(latents=latents, output_type="latent")
+            logger.info(f"[Upscaler] Upscale concluído. Novo shape: {result.latents.shape}")
             return result.latents
         except Exception as e:
-            logger.error(f"[Upscaler] Erro durante upscale: {e}")
+            logger.error(f"[Upscaler] Erro durante upscale: {e}", exc_info=True)
             return latents
-            
 
-# Instanciação Singleton
-# Depende do VAE do ltx_manager, então o obtemos de lá.
-try:
-    base_vae_for_upscaler = ltx_manager_singleton.workers[0].pipeline.vae
-    upscaler_specialist_singleton = UpscalerSpecialist(base_vae=base_vae_for_upscaler)
-except Exception as e:
-    logger.error(f"Não foi possível inicializar o UpscalerSpecialist Singleton: {e}")
-    upscaler_specialist_singleton = None
+
+# ---------------------------
+# Singleton global
+# ---------------------------
+upscaler_specialist_singleton = UpscalerSpecialist()