Update api/ltx/vae_aduc_pipeline.py

api/ltx/vae_aduc_pipeline.py

@@ -87,12 +87,14 @@ class VaeAducPipeline:
             self._initialized = True
             logging.info(f"✅ VaeServer ready. VAE model is 'hot' on {self.device} with dtype {self.dtype}. Startup time: {time.time() - t0:.2f}s")
 
+    @log_function_io
     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()
 
+    @log_function_io
     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 para encodar."""
         if isinstance(item, Image.Image):
@@ -112,7 +114,7 @@ class VaeAducPipeline:
         tensor_5d = tensor.unsqueeze(0).unsqueeze(2)
         return (tensor_5d * 2.0) - 1.0
 
-    @torch.no_grad()
+    @log_function_io
     def generate_conditioning_items(
         self,
         media_items: List[Union[Image.Image, torch.Tensor]],
@@ -143,7 +145,7 @@ class VaeAducPipeline:
         finally:
             self._cleanup_gpu()
 
-    @torch.no_grad()
+    @log_function_io
     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()