Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -307,7 +307,6 @@ class VideoService:
     
     
 
-    # ADICIONE A FUNÇÃO ABAIXO
     @torch.no_grad()
     def _image_to_latents(self, image_input: Union[str, Image.Image], height: int, width: int) -> torch.Tensor:
         """
@@ -428,7 +427,7 @@ class VideoService:
             tensor_path = self._save_latents_to_disk(latents, "latents_low_res", used_seed)
             
             final_video_path = self._save_video_from_tensor(pixel_tensor, f"final_video_{seed}", seed, temp_dir, fps=DEFAULT_FPS)
-            return final_video_path
+            # A linha "return final_video_path" foi removida daqui!
         
         # --- Limpeza ---
         self._finalize()
@@ -437,6 +436,118 @@ class VideoService:
         return final_video_path, tensor_path, used_seed
 
 
+    # --------------------------------------------------------------------------
+    # --- Métodos Públicos (API do Serviço) ---
+    # --------------------------------------------------------------------------
+    
+    def generate_upscale_denoise(
+        self,
+        latents_path: str,
+        prompt: str,
+        negative_prompt: str,
+        guidance_scale: float,
+        seed: int,
+        # Estes deveriam ser parâmetros no mundo ideal, mas vamos recuperá-los do latente se necessário
+        # height: int,
+        # width: int,
+        # duration_secs: float,
+    ) -> Tuple[str, str]:
+        """
+        ETAPA 2: Aplica upscaling espacial e denoise na textura usando o pipeline principal
+        para refinar um tensor de latentes de baixa resolução.
+        """
+        print("[INFO] Iniciando ETAPA 2: Upscale e Refinamento de Textura (LTX)...")
+        
+        # --- 1. Carregar Latentes de Entrada ---
+        if not Path(latents_path).exists():
+            raise FileNotFoundError(f"Latentes não encontrados no caminho: {latents_path}")
+            
+        latents_low_res = torch.load(latents_path).to(self.device, dtype=self._get_precision_dtype())
+        log_tensor_info(latents_low_res, "Latentes Carregados (Baixa Resolução)")
+
+        # --- 2. Upsample Espacial e Filtro AdaIN ---
+        with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
+            upsampled_latents = self._upsample_and_filter_latents(latents_low_res)
+        log_tensor_info(upsampled_latents, "Latentes Upscaled (Antes do Denoise)")
+        
+        # --- 3. Denoise (Segundo Passo da Pipeline) ---
+        _, _, _, latent_height, latent_width = upsampled_latents.shape
+        # O upscaler dobra a resolução espacial
+        target_height = latent_height * self.pipeline.vae_scale_factor
+        target_width = latent_width * self.pipeline.vae_scale_factor
+
+        # O num_frames latente é preservado
+        latent_num_frames = upsampled_latents.shape[2]
+        actual_num_frames = (latent_num_frames - 1) * 8 + 1 # Reverte de latent frames para pixel frames (n*8+1)
+        duration_secs = actual_num_frames / DEFAULT_FPS
+        
+        print(f"  - Resolução de Saída Estimada: {target_height}x{target_width}")
+        print(f"  - Frames Estimados: {actual_num_frames}")
+
+        temp_dir = tempfile.mkdtemp(prefix="ltxv_high_")
+        self._register_tmp_dir(temp_dir)
+        
+        second_pass_kwargs = {
+            "prompt": prompt,
+            "negative_prompt": negative_prompt,
+            "height": target_height,
+            "width": target_width,
+            "num_frames": actual_num_frames,
+            "frame_rate": int(DEFAULT_FPS),
+            "generator": torch.Generator(device=self.device).manual_seed(seed),
+            "output_type": "latent",
+            "vae_per_channel_normalize": True,
+            "is_video": True,
+            "latents": upsampled_latents, # Passa os latentes upscaled como ponto de partida
+            "guidance_scale": float(guidance_scale),
+            **(self.config.get("second_pass", {}))
+        }
+
+        print("  - Enviando para a pipeline LTX (Refinamento)...")
+        with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
+            refined_latents = self.pipeline(**second_pass_kwargs).images
+        log_tensor_info(refined_latents, "Latentes Refinados (Saída do Denoise)")
+
+        # --- 4. Decodificação e Saída ---
+        pixel_tensor = vae_manager_singleton.decode(refined_latents, decode_timestep=float(self.config.get("decode_timestep", 0.00)))
+        
+        tensor_path = self._save_latents_to_disk(refined_latents, "latents_refined", seed)
+        video_path = self._save_video_from_tensor(pixel_tensor, f"refined_video_{seed}", seed, temp_dir, fps=DEFAULT_FPS)
+        
+        self._finalize()
+        
+        print("[SUCCESS] ETAPA 2 Concluída.")
+        return video_path, tensor_path
+        
+    def move_to_cpu(self):
+        """Move todos os modelos para a CPU para liberar VRAM."""
+        print("[LTX/SWAP] Movendo modelos LTX para a CPU...")
+        self.pipeline.to("cpu")
+        if self.latent_upsampler:
+            self.latent_upsampler.to("cpu")
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        print("[LTX/SWAP] Modelos LTX na CPU.")
+        
+    def move_to_device(self, device: torch.device):
+        """Move todos os modelos para o dispositivo especificado (normalmente GPU)."""
+        if str(device) == "cpu":
+            return self.move_to_cpu()
+            
+        print(f"[LTX/SWAP] Movendo modelos LTX para {device}...")
+        self.pipeline.to(device)
+        if self.latent_upsampler:
+            self.latent_upsampler.to(device)
+        self.device = str(device)
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        print(f"[LTX/SWAP] Modelos LTX em {device}.")
+
+
     # --------------------------------------------------------------------------
     # --- Métodos Internos e Auxiliares ---
     # --------------------------------------------------------------------------