Update managers/vae_manager.py

managers/vae_manager.py

@@ -22,34 +22,35 @@ class _SimpleVAEManager:
         if autocast_dtype is not None:
             self.autocast_dtype = autocast_dtype
 
-    @torch.no_grad()
-    def decode(self, latents_5d: torch.Tensor) -> torch.Tensor:
-        """
-        Decodifica todo o bloco 5D de uma vez, replicando o fluxo simples do deformes4D.
-        Retorna tensor de pixels 5D em [0,1] com shape (B,C,T,H',W').
-        """
-        if self.pipeline is None:
-            raise RuntimeError("VAE Manager sem pipeline. Chame attach_pipeline primeiro.")
 
-        # Garante device correto
-        latents_5d = latents_5d.to(self.device, non_blocking=True)
 
+    @torch.no_grad()
+    def decode(self, latent_tensor: torch.Tensor, decode_timestep: float = 0.05) -> torch.Tensor:
+        
+        # Garante device e dtype conforme runtime
+        latent_tensor_gpu = latent_tensor.to(self.device, dtype=self.autocast_dtype if self.device == "cuda" else latent_tensor.dtype)
+        
+        # Constrói o vetor de timesteps (um por item no batch B)
+        num_items_in_batch = latent_tensor_gpu.shape[0]
+        timestep_tensor = torch.tensor([decode_timestep] * num_items_in_batch, device=self.device, dtype=latent_tensor_gpu.dtype)
+    
         ctx = torch.autocast(device_type="cuda", dtype=self.autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
         with ctx:
-            if hasattr(self.pipeline, "decode_latents"):
-                pixels_5d = self.pipeline.decode_latents(latents_5d)
-            elif hasattr(self.pipeline, "vae") and hasattr(self.pipeline.vae, "decode"):
-                pixels_5d = self.pipeline.vae.decode(latents_5d)
-            else:
-                raise RuntimeError("Pipeline não expõe decode_latents nem vae.decode.")
-
+            pixels = vae_decode(
+                latent_tensor_gpu,
+                self.pipeline.vae if hasattr(self.pipeline, "vae") else self.pipeline,  # compat
+                is_video=True,
+                timestep=timestep_tensor,
+                vae_per_channel_normalize=True,
+            )
+    
         # Normaliza para [0,1] se vier em [-1,1]
-        if pixels_5d.min() < 0:
-            pixels_5d = (pixels_5d.clamp(-1, 1) + 1.0) / 2.0
+        if pixels.min() < 0:
+            pixels = (pixels.clamp(-1, 1) + 1.0) / 2.0
         else:
-            pixels_5d = pixels_5d.clamp(0, 1)
-
-        return pixels_5d
+            pixels = pixels.clamp(0, 1)
+        return pixels
+    
 
 # Singleton global de uso simples
 vae_manager_singleton = _SimpleVAEManager()