Update ltx_video/pipelines/pipeline_ltx_video.py

ltx_video/pipelines/pipeline_ltx_video.py

@@ -25,6 +25,9 @@ from transformers import (
     AutoTokenizer,
 )
 
+
+import logging
+
 from ltx_video.models.autoencoders.causal_video_autoencoder import (
     CausalVideoAutoencoder,
 )
@@ -50,6 +53,12 @@ from ltx_video.models.autoencoders.vae_encode import (
 
 
 
+import logging # Garanta que a importação padrão esteja no topo do arquivo
+import torch   # Já deve estar lá
+
+# A linha abaixo já existe no seu arquivo e está correta
+logger = logging.get_logger(__name__) 
+
 def log_memory_snapshot(event_name: str):
     """
     Exibe um log do consumo de memória VRAM da GPU atual no exato momento da chamada.
@@ -62,8 +71,11 @@ def log_memory_snapshot(event_name: str):
     reserved = torch.cuda.memory_reserved(device) / (1024 ** 2)
     total = torch.cuda.get_device_properties(device).total_memory / (1024 ** 2)
     
-    # Usamos o emoji do cérebro para indicar um log de diagnóstico de memória
-    print(
+    # =========================================================================
+    # ========================== ALTERAÇÃO PRINCIPAL ==========================
+    # Substitua 'print' por 'logger.info' ou 'logger.debug'
+    # =========================================================================
+    logger.info(
         f"🧠 MEMORY LOG [{event_name.ljust(35)}]: "
         f"Alocada: {allocated:.2f} MB | "
         f"Reservada: {reserved:.2f} MB | "
@@ -71,7 +83,6 @@ def log_memory_snapshot(event_name: str):
         f"({(reserved/total)*100:.1f}%)"
     )
 
-
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name