Update aduc_framework/utils/callbacks.py

aduc_framework/utils/callbacks.py

@@ -2,103 +2,52 @@
 
 import imageio
 import torch
-from PIL import Image, ImageDraw, ImageFont
 import numpy as np
-import math
-import os
+import tempfile
+from diffusers.utils.export_utils import export_to_video
 
 class DenoiseStepLogger:
     """
-    Uma classe de callback que "espiona" o processo de denoising.
-    Ela captura frames intermediários, pode salvá-los como um vídeo de processo
-    e também pode criar uma grade de comparação visual de todas as etapas.
+    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
-        self.intermediate_frames = []
-        # Mantém os tensores na CPU por padrão durante a inicialização.
-        # Eles serão movidos para o dispositivo correto no momento do uso.
+        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)
 
-    ### INÍCIO DA SEÇÃO CORRIGIDA ###
-    def decode_latents_to_pil(self, latents: torch.Tensor) -> Image.Image:
-        """Decodifica um tensor de latents para uma única imagem PIL."""
-        
-        # Pega o dispositivo correto do tensor de entrada `latents`
+    def decode_latents_to_video_path(self, latents: torch.Tensor, step: int) -> str:
+        """
+        Decodifica um tensor latente 5D, salva o vídeo resultante em um arquivo
+        temporário e retorna o caminho para esse arquivo.
+        """
         correct_device = latents.device
-        
-        # Move os tensores de média e desvio padrão para o mesmo dispositivo dos latents
-        # antes de realizar a operação. Isso evita o erro de "device meta".
         latents_unscaled = latents / self.latents_std.to(correct_device) + self.latents_mean.to(correct_device)
         
         latents_unscaled = latents_unscaled.to(self.pipe.vae.dtype)
-        decoded_video_tensor = self.pipe.vae.decode(latents_unscaled, return_dict=False)[0]
+        video_tensor = self.pipe.vae.decode(latents_unscaled, return_dict=False)[0]
+        
+        # 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
         
-        frame_tensor = decoded_video_tensor[0, :, 0, :, :]
-        frame_tensor = (frame_tensor / 2 + 0.5).clamp(0, 1)
-        frame_np = frame_tensor.cpu().permute(1, 2, 0).float().numpy()
-        pil_image = Image.fromarray((frame_np * 255).astype(np.uint8))
-        return pil_image
-    ### FIM DA SEÇÃO CORRIGIDA ###
+        export_to_video(frames[0], video_path, fps=self.fps)
+        return video_path
 
     def __call__(self, pipe, step: int, timestep: int, callback_kwargs: dict):
         """
-        Esta função é chamada pela pipeline da diffusers em cada passo de denoising.
+        Chamado pela pipeline a cada passo.
         """
-        print(f"  -> Callback: Capturando frame do passo de denoising {step+1}...")
+        print(f"  -> Callback: Decodificando vídeo completo do passo de denoising {step+1}...")
         latents = callback_kwargs["latents"]
-        pil_frame = self.decode_latents_to_pil(latents)
-        self.intermediate_frames.append(pil_frame)
-        return callback_kwargs
-
-    def save_as_video(self, output_path: str, fps: int = 5):
-        """Salva os frames intermediários capturados como um vídeo MP4."""
-        if not self.intermediate_frames:
-            print("  -> Callback: Nenhum frame intermediário para salvar como vídeo.")
-            return
-        print(f"  -> Callback: Codificando {len(self.intermediate_frames)} frames em vídeo em '{output_path}'...")
-        writer = imageio.get_writer(output_path, fps=fps, codec='libx264', quality=8, pixelformat='yuv420p')
-        for frame in self.intermediate_frames:
-            writer.append_data(np.array(frame))
-        writer.close()
-        print("  -> Callback: Vídeo de depuração salvo com sucesso.")
-
-    def create_steps_grid(self) -> Image.Image:
-        """
-        Organiza todos os frames intermediários capturados em uma única imagem de grade para comparação.
-        """
-        if not self.intermediate_frames:
-            print("  -> Callback: Nenhum frame intermediário para criar a grade.")
-            return None
-        print(f"  -> Callback: Criando grade de comparação com {len(self.intermediate_frames)} etapas...")
-        num_images = len(self.intermediate_frames)
-        cols = math.ceil(math.sqrt(num_images))
-        rows = math.ceil(num_images / cols)
-        frame_w, frame_h = self.intermediate_frames[0].size
-        grid_w, grid_h = frame_w * cols, frame_h * rows
-        grid_image = Image.new('RGB', (grid_w, grid_h), (20, 20, 20))
-        draw = ImageDraw.Draw(grid_image)
-        try:
-            font_path = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
-            if not os.path.exists(font_path):
-                font_path = "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf"
-            font = ImageFont.truetype(font_path, size=32)
-        except IOError:
-            print("  -> Callback WARNING: Fonte não encontrada. Usando fonte padrão.")
-            font = ImageFont.load_default()
-        for i, frame in enumerate(self.intermediate_frames):
-            x, y = (i % cols) * frame_w, (i // cols) * frame_h
-            grid_image.paste(frame, (x, y))
-            text = f"Passo {i+1}"
-            text_origin = (x + 10, y + 10)
-            try:
-                text_bbox = draw.textbbox(text_origin, text, font=font)
-            except AttributeError:
-                text_w, text_h = draw.textsize(text, font=font)
-                text_bbox = (text_origin[0], text_origin[1], text_origin[0] + text_w, text_origin[1] + text_h)
-            rect_coords = (text_bbox[0] - 5, text_bbox[1] - 5, text_bbox[2] + 5, text_bbox[3] + 5)
-            draw.rectangle(rect_coords, fill=(0, 0, 0, 180))
-            draw.text(text_origin, text, font=font, fill=(255, 255, 255))
-        print("  -> Callback: Grade de comparação criada com sucesso.")
-        return grid_image
+        
+        # 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)
+        
+        return callback_kwargs