Update aduc_framework/utils/callbacks.py

aduc_framework/utils/callbacks.py

@@ -1,38 +1,56 @@
 # aduc_framework/utils/callbacks.py
 
+import imageio
 import torch
 import numpy as np
+import tempfile
+from diffusers.utils.export_utils import export_to_video
 
 class DenoiseStepLogger:
     """
-    Callback simplificado que apenas decodifica latentes para um tensor de vídeo em pixels (como um array NumPy).
-    A lógica de salvar em arquivo foi movida para o manager para melhor controle do fluxo com 'yield'.
+    Callback que, em cada passo do denoising, decodifica a sequência de vídeo
+    inteira e a salva como um clipe MP4 individual.
     """
-    def __init__(self, pipe):
+    def __init__(self, pipe, fps=8):
         self.pipe = pipe
-        # Mantém os tensores na CPU por padrão. Eles serão movidos para o dispositivo
-        # correto no momento do uso para garantir compatibilidade com `device_map="auto"`.
+        self.fps = fps
+        # Armazena os caminhos para os vídeos gerados em cada passo
+        self.step_video_paths = []
         self.latents_mean = torch.tensor(pipe.vae.config.latents_mean).view(1, pipe.vae.config.z_dim, 1, 1, 1)
         self.latents_std = 1.0 / torch.tensor(pipe.vae.config.latents_std).view(1, pipe.vae.config.z_dim, 1, 1, 1)
 
-    def decode_latents_to_video_tensor(self, latents: torch.Tensor) -> np.ndarray:
+    def decode_latents_to_video_path(self, latents: torch.Tensor, step: int) -> str:
         """
-        Decodifica um tensor latente 5D para um array numpy de frames de vídeo.
-        
-        Args:
-            latents (torch.Tensor): O tensor 5D vindo do processo de denoising.
-
-        Returns:
-            np.ndarray: Um array NumPy representando o lote de frames de vídeo (B, F, H, W, C).
+        Decodifica um tensor latente 5D, salva o vídeo resultante em um arquivo
+        temporário e retorna o caminho para esse arquivo.
         """
-        # Garante que as operações aconteçam no mesmo dispositivo que os latentes
         correct_device = latents.device
         latents_unscaled = latents / self.latents_std.to(correct_device) + self.latents_mean.to(correct_device)
         
         latents_unscaled = latents_unscaled.to(self.pipe.vae.dtype)
         video_tensor = self.pipe.vae.decode(latents_unscaled, return_dict=False)[0]
         
-        # O post-processador da pipeline converte o tensor de vídeo para um array NumPy
-        video_np = self.pipe.video_processor.postprocess_video(video=video_tensor, output_type="np")
+        # O resultado já é um lote de frames de vídeo
+        frames = self.pipe.video_processor.postprocess_video(video=video_tensor, output_type="np")
+        
+        with tempfile.NamedTemporaryFile(suffix=f"_step_{step+1}.mp4", delete=False) as tmp:
+            video_path = tmp.name
+        
+        export_to_video(frames[0], video_path, fps=self.fps)
+        return video_path
+
+    def __call__(self, pipe, step: int, timestep: int, callback_kwargs: dict):
+        """
+        Chamado pela pipeline a cada passo.
+        """
+        print(f"  -> Callback: Decodificando vídeo completo do passo de denoising {step+1}...")
+        latents = callback_kwargs["latents"]
+        
+        # Gera o vídeo para o passo atual e armazena seu caminho
+        video_path = self.decode_latents_to_video_path(latents, step)
+        self.step_video_paths.append(video_path)
+
+        yield None, None, self.denoising_step_videos
+        
         
-        return video_np
+        return callback_kwargs