Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -554,7 +554,7 @@ class VideoService:
             
             self._finalize()
 
-            final_video_path, final_latents_path = self.refine_texture_only(
+            final_video_path, final_latents_path, _ = self.refine_texture_only(
                 latents_path=tensor_path,
                 prompt=prompt,
                 negative_prompt=negative_prompt,
@@ -570,10 +570,123 @@ class VideoService:
         return final_video_path, final_latents_path, used_seed
 
 
-    # Em api/ltx_server_refactored.py, dentro da classe VideoService
-
-
+    
     def apply_secondary_refinement(
+        self,
+        low_res_latents_path: str,
+        prompt: str,
+        negative_prompt: str,
+        guidance_scale: float,
+        seed: int
+    ) -> Tuple[str, str]:
+        """
+        Orquestra o pipeline de Pós-Produção de forma elegante:
+        1. Upscale dos latentes de baixa resolução.
+        2. Processa cada chunk usando `refine_texture_only`, coletando os TENSORES DE PIXELS.
+        3. Costura os tensores de pixels na memória e salva o MP4 final.
+        """
+        print("\n======================================================================")
+        print("====== [INFO] Iniciando Pós-Produção: Costura de Tensores de Pixels ======")
+        print("======================================================================\n")
+    
+        temp_dir_orchestrator = tempfile.mkdtemp(prefix="refine_orchestrator_")
+        self._register_tmp_dir(temp_dir_orchestrator)
+    
+        # --- ETAPA 1: Upscale Espacial do Tensor Completo ---
+        print("[LOG] Etapa 1: Aplicando Upscale Espacial nos latentes de baixa resolução...")
+        self._set_generation_environment()
+        
+        low_res_latents = torch.load(low_res_latents_path)
+        with torch.no_grad(), torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
+            latents_hd_upscaled = self._upsample_and_filter_latents(low_res_latents.to(self.device))
+        
+        latents_hd_upscaled_cpu = latents_hd_upscaled.cpu()
+        print(f"  [LOG] Upscale concluído. Novo shape de latentes HD: {latents_hd_upscaled_cpu.shape}")
+        del latents_hd_upscaled, low_res_latents
+        torch.cuda.empty_cache()
+    
+        # --- ETAPA 2: Fragmentar e Processar cada Chunk para Obter Tensores de Pixels ---
+        print("\n[LOG] Etapa 2: Fragmentando e processando cada chunk para obter tensores de pixels...")
+        total_latents = latents_hd_upscaled_cpu.shape[2]
+        sobreposicao = 2
+        num_divisoes = 3
+        tamanho_base = (total_latents - 1) // num_divisoes
+        
+        pontos_de_corte = [
+            (0, tamanho_base + sobreposicao),
+            (tamanho_base - sobreposicao, (2 * tamanho_base) + sobreposicao),
+            ((2 * tamanho_base) - sobreposicao, total_latents)
+        ]
+        hd_chunks_to_process = [latents_hd_upscaled_cpu[:, :, s:e, :, :] for s, e in pontos_de_corte]
+        
+        pixel_chunks_list = []
+        # Também coletaremos os latentes refinados de cada chunk para a costura final de latentes
+        refined_latent_chunk_paths = []
+    
+        for i, chunk in enumerate(hd_chunks_to_process):
+            print(f"\n  --> Processando para Pixels o Chunk {i+1}/{len(hd_chunks_to_process)}...")
+            chunk_path = os.path.join(temp_dir_orchestrator, f"hd_chunk_to_process_{i}.pt")
+            torch.save(chunk, chunk_path)
+            chunk_seed = seed + i + 1
+    
+            # Chamamos a função worker e pegamos todos os seus retornos
+            _, refined_latent_path, pixel_chunk = self.refine_texture_only(
+                latents_path=chunk_path,
+                prompt=prompt, negative_prompt=negative_prompt,
+                guidance_scale=guidance_scale, seed=chunk_seed
+            )
+            pixel_chunks_list.append(pixel_chunk)
+            refined_latent_chunk_paths.append(refined_latent_path)
+            print(f"  --> Tensor de pixels do Chunk {i+1} recebido. Shape: {pixel_chunk.shape}")
+    
+        # --- ETAPA 3: Costurar os Tensores de Pixels ---
+        print("\n[LOG] Etapa 3: Costurando os tensores de pixels na CPU...")
+        
+        frames_p1 = tamanho_base * 8
+        parte1_pix = pixel_chunks_list[0][:, :, :frames_p1, :, :]
+        
+        descarte_inicio_p2 = sobreposicao * 8
+        frames_p2 = tamanho_base * 8
+        parte2_pix = pixel_chunks_list[1][:, :, descarte_inicio_p2 : descarte_inicio_p2 + frames_p2, :, :]
+        
+        descarte_inicio_p3 = sobreposicao * 8
+        parte3_pix = pixel_chunks_list[2][:, :, descarte_inicio_p3:, :, :]
+        
+        final_pixel_tensor = torch.cat([parte1_pix, parte2_pix, parte3_pix], dim=2)
+        print(f"  [LOG] Tensor de pixels final montado. Shape Final: {final_pixel_tensor.shape}")
+        del pixel_chunks_list, parte1_pix, parte2_pix, parte3_pix # Limpeza de memória
+    
+        # --- ETAPA 4: Costurar os Latentes (para consistência do estado final) ---
+        print("\n[LOG] Etapa 4: Costurando os latentes refinados para o estado final...")
+        refined_hd_chunks = [torch.load(path) for path in refined_latent_chunk_paths]
+        
+        parte1_lat = refined_hd_chunks[0][:, :, :tamanho_base, :, :]
+        parte2_lat = refined_hd_chunks[1][:, :, sobreposicao : sobreposicao + tamanho_base, :, :]
+        parte3_lat = refined_hd_chunks[2][:, :, sobreposicao:, :, :]
+        
+        final_latents_hd = torch.cat([parte1_lat, parte2_lat, parte3_lat], dim=2)
+        print(f"  [LOG] Latentes HD finais montados. Shape Final: {final_latents_hd.shape}")
+        del refined_hd_chunks, parte1_lat, parte2_lat, parte3_lat
+    
+        # --- ETAPA 5: Salvar os Artefatos Finais ---
+        print("\n[LOG] Etapa 5: Salvando o vídeo MP4 e o tensor de latentes finais...")
+        
+        # Salva o tensor de pixels costurado como o vídeo final
+        final_video_path = self._save_video_from_tensor(
+            final_pixel_tensor, 
+            f"final_progressive_stitched_{seed}", 
+            seed, 
+            temp_dir_orchestrator 
+        )
+        
+        # Salva o tensor de latentes costurado para o estado do app
+        final_latents_path = self._save_latents_to_disk(final_latents_hd, "final_hd_stitched", seed)
+    
+        print(f"\n[SUCCESS] Pós-Produção completa! Vídeo final: {final_video_path}")
+        return final_video_path, final_latents_path
+        
+    
+    def apply_secondary_refinement1(
         self,
         latents_path: str,
         prompt: str,
@@ -707,18 +820,19 @@ class VideoService:
         
         video_path_out = self._save_video_from_tensor(pixel_tensor, "refined_video_final", used_seed, temp_dir)
         latents_path_out = self._save_latents_to_disk(final_latents_cpu, "latents_refined_final", used_seed)
-        
+
+                
         print("  [LOG] Tarefa concluída. Restaurando ambiente de GERAÇÃO na GPU para a próxima execução...")
         self._set_generation_environment()
         
         print("  [LOG] Liberando tensores finais da memória da CPU.")
-        del pixel_tensor, final_latents_cpu
+        del final_latents_cpu
         self._finalize()
         
         print("\n======================================================================")
         print("============ [SUCCESS] ETAPA 2 Concluída com Sucesso =============")
         print("======================================================================\n")
-        return video_path_out, latents_path_out
+        return video_path_out, latents_path_out, pixel_tensor