Update upscaler_specialist.py

upscaler_specialist.py

@@ -20,30 +20,51 @@ class UpscalerSpecialist:
         self.base_vae = None
         self.pipe_upsample = None
 
+
     def _lazy_init(self):
-        """Inicializa VAE e pipeline apenas quando necessário."""
-        if self.base_vae is None:
-            try:
-                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:
-                logger.error(f"[Upscaler] Falha ao inicializar VAE: {e}")
-                return
+    if self.pipe_upsample is not None:
+        return
 
-        if self.pipe_upsample is None and self.base_vae is not None:
-            try:
-                self.pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained(
+    try:
+        # Tenta usar o VAE do ltx_manager
+        if ltx_manager_singleton.workers:
+            candidate_vae = ltx_manager_singleton.workers[0].pipeline.vae
+            if candidate_vae.__class__.__name__ == "AutoencoderKLLTXVideo":
+                self.base_vae = candidate_vae
+                logger.info("[Upscaler] Usando VAE do ltx_manager (AutoencoderKLLTXVideo).")
+            else:
+                logger.warning(f"[Upscaler] VAE incompatível: {type(candidate_vae)}. "
+                               "Carregando AutoencoderKLLTXVideo manualmente...")
+                from diffusers.models.autoencoders import AutoencoderKLLTXVideo
+                self.base_vae = AutoencoderKLLTXVideo.from_pretrained(
                     "linoyts/LTX-Video-spatial-upscaler-0.9.8",
-                    vae=self.base_vae,
+                    subfolder="vae",
                     torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
                 ).to(self.device)
-                logger.info("[Upscaler] Pipeline carregado com sucesso.")
-            except Exception as e:
-                logger.error(f"[Upscaler] Falha ao carregar pipeline: {e}")
+        else:
+            logger.warning("[Upscaler] Nenhum worker disponível, carregando VAE manualmente...")
+            from diffusers.models.autoencoders import AutoencoderKLLTXVideo
+            self.base_vae = AutoencoderKLLTXVideo.from_pretrained(
+                "linoyts/LTX-Video-spatial-upscaler-0.9.8",
+                subfolder="vae",
+                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
+            ).to(self.device)
+
+        # Carregar pipeline
+        self.pipe_upsample = LTXLatentUpsamplePipeline.from_pretrained(
+            "linoyts/LTX-Video-spatial-upscaler-0.9.8",
+            vae=self.base_vae,
+            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32
+        ).to(self.device)
+
+        logger.info("[Upscaler] Pipeline carregado com sucesso.")
+
+    except Exception as e:
+        logger.error(f"[Upscaler] Falha ao carregar pipeline: {e}")
+        self.pipe_upsample = None
+        
 
+    
     @torch.no_grad()
     def upscale(self, latents: torch.Tensor) -> torch.Tensor:
         """Aplica o upscaling 2x nos tensores latentes fornecidos."""