Update upscaler_specialist.py

upscaler_specialist.py

@@ -10,53 +10,49 @@ from ltx_manager_helpers import ltx_manager_singleton
 
 logger = logging.getLogger(__name__)
 
+
 class UpscalerSpecialist:
-    """
-    Especialista responsável por aumentar a resolução espacial de tensores latentes
-    usando o LTX Video Spatial Upscaler.
-    """
-    def __init__(self, base_vae):
-        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    def __init__(self, device="cuda"):
+        self.device = device if torch.cuda.is_available() else "cpu"
         self.pipe_upsample = None
-        if base_vae is not None:
-            logger.info("Inicializando Especialista de Upscale Latente...")
+        self.base_vae = None
+
+    def _lazy_init(self):
+        """Inicializa o VAE e o pipeline somente quando for chamado."""
+        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
+                else:
+                    logger.warning("[Upscaler] Nenhum worker disponível no ltx_manager_singleton.")
+            except Exception as e:
+                logger.error(f"[Upscaler] Falha ao inicializar VAE: {e}")
+                return
+
+        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=base_vae,
-                    torch_dtype=torch.bfloat16,
+                    vae=self.base_vae,
+                    torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
                 ).to(self.device)
-                logger.info("Especialista de Upscale Latente pronto.")
+                logger.info("[Upscaler] Pipeline carregado com sucesso.")
             except Exception as e:
-                logger.error(f"Falha ao carregar o modelo de upscale: {e}", exc_info=True)
-        else:
-            logger.warning("VAE base não fornecido. Especialista de Upscale desativado.")
+                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 não está disponível. Retornando latentes originais.")
+            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")
+            return result.latents
+        except Exception as e:
+            logger.error(f"[Upscaler] Erro durante upscale: {e}")
             return latents
-        
-        logger.info(f"Upscaler: Recebeu latentes com shape {latents.shape}. Aplicando upscale 2x...")
-        
-        # O upscaler opera em um batch de latentes.
-        upscaled_latents = self.pipe_upsample(
-            latents=latents,
-            output_type="latent"
-        ).frames
-        
-        logger.info(f"Upscaler: Latentes redimensionados para {upscaled_latents.shape}.")
-        return upscaled_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
+            
+