Update ltx_worker_upscaler.py

ltx_worker_upscaler.py

@@ -1,6 +1,6 @@
-#-- START OF MODIFIED FILE app_fluxContext_Ltx/ltx_worker_upscaler.py ---
+#--- START OF MODIFIED FILE app_fluxContext_Ltx/ltx_worker_upscaler.py ---
 # ltx_worker_upscaler.py
-# Worker para fazer upscale dos fragmentos de vídeo para alta resolução.
+# Worker para fazer upscale de latentes de vídeo para alta resolução.
 # Este arquivo é parte do projeto Euia-AducSdr e está sob a licença AGPL v3.
 # Copyright (C) 4 de Agosto de 2025  Carlos Rodrigues dos Santos
 
@@ -15,7 +15,7 @@ import huggingface_hub
 
 from inference import create_ltx_video_pipeline
 from ltx_video.models.autoencoders.latent_upsampler import LatentUpsampler
-from ltx_video.models.autoencoders.vae_encode import vae_encode, vae_decode
+from ltx_video.models.autoencoders.vae_encode import vae_decode
 
 class LtxUpscaler:
     def __init__(self, device_id='cuda:0'):
@@ -72,28 +72,22 @@ class LtxUpscaler:
                 torch.cuda.empty_cache()
     
     @torch.no_grad()
-    def upscale_video_fragment(self, video_path_low_res: str, output_path: str, video_fps: int):
-        print(f"UPSCALER ({self.device}): Processando {os.path.basename(video_path_low_res)}")
+    def upscale_latents_to_video(self, latent_path: str, output_path: str, video_fps: int):
+        print(f"UPSCALER ({self.device}): Processando latentes de {os.path.basename(latent_path)}")
         
-        with imageio.get_reader(video_path_low_res) as reader:
-            video_frames = [frame for frame in reader]
-        video_np = np.stack(video_frames)
-        
-        video_tensor = torch.from_numpy(video_np).permute(0, 3, 1, 2).float() / 255.0
-        video_tensor = (video_tensor * 2.0) - 1.0
-        video_tensor = video_tensor.unsqueeze(0).permute(0, 2, 1, 3, 4)
-        video_tensor = video_tensor.to(self.device, dtype=self.model_dtype)
+        # Carrega os latentes do disco e os envia para a GPU
+        latents = torch.load(latent_path).to(self.device, dtype=self.model_dtype)
 
-        latents = vae_encode(video_tensor, self.vae)
+        # PASSO 1: Upscale Espacial (não precisamos mais de vae_encode)
         upsampled_latents = self.latent_upsampler(latents)
         
-        # --- CORREÇÃO AQUI ---
-        # O VAE espera um tensor de timestep. Vamos criar um.
+        # (Opcional: PASSO 2 - Upscale Temporal seria inserido aqui no futuro)
+
+        # PASSO 3: Decodificação Final
         decode_timestep = torch.tensor([0.0] * upsampled_latents.shape[0], device=self.device)
         upsampled_video_tensor = vae_decode(
             upsampled_latents, self.vae, is_video=True, timestep=decode_timestep
         )
-        # --- FIM DA CORREÇÃO ---
         
         upsampled_video_tensor = (upsampled_video_tensor.clamp(-1, 1) + 1) / 2.0
         video_np_high_res = (upsampled_video_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
@@ -102,6 +96,21 @@ class LtxUpscaler:
             for frame in video_np_high_res:
                 writer.append_data(frame)
 
-        print(f"UPSCALER ({self.device}): Arquivo salvo em {os.path.basename(output_path)}")
+        print(f"UPSCALER ({self.device}): Arquivo de vídeo salvo em {os.path.basename(output_path)}")
         return output_path
+        
+    @torch.no_grad()
+    def decode_single_latent_frame(self, latent_frame_tensor: torch.Tensor) -> Image.Image:
+        """Decodifica um único frame latente para uma imagem PIL para o Gemini."""
+        latent_frame_tensor = latent_frame_tensor.to(self.device, dtype=self.model_dtype)
+        
+        decode_timestep = torch.tensor([0.0] * latent_frame_tensor.shape[0], device=self.device)
+        decoded_tensor = vae_decode(
+            latent_frame_tensor, self.vae, is_video=True, timestep=decode_timestep
+        )
+        
+        decoded_tensor = (decoded_tensor.clamp(-1, 1) + 1) / 2.0
+        # Shape: (B, C, F, H, W) -> (H, W, C)
+        numpy_image = (decoded_tensor[0].permute(2, 3, 1, 0).squeeze().cpu().float().numpy() * 255).astype(np.uint8)
+        return Image.fromarray(numpy_image)
 #--- END OF MODIFIED FILE app_fluxContext_Ltx/ltx_worker_upscaler.py ---