Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -451,8 +451,130 @@ class VideoService:
         print(f"\n[SUCCESS] ETAPA 1 Concluída. Vídeo guia: {final_video_path}, Latentes guia: {final_latents_path}")
         return final_video_path, final_latents_path, used_seed
 
+
+    # Em ltx_server_refactored.py, dentro da classe VideoService
+
+def apply_secondary_refinement(
+    self,
+    latents_path: str,
+    prompt: str,
+    negative_prompt: str,
+    guidance_scale: float,
+    seed: int,
+    num_decode_chunks_per_worker: int = 2
+) -> Tuple[str, str]:
+    """
+    Orquestra o pipeline de Pós-Produção: Upscale, Denoise em Chunks, e Costura final.
+    Gerencia o estado da GPU explicitamente para cada etapa.
+    """
+    print("\n======================================================================")
+    print("====== [INFO] Iniciando Pós-Produção: Refinamento em Chunks e Costura Final ======")
+    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...")
+    
+    # --- Gerenciamento de VRAM para a Etapa de Upscale ---
+    # Descarrega o Transformer (se a otimização estiver ativa) para dar espaço.
+    if ENABLE_MEMORY_OPTIMIZATION:
+        if next(self.pipeline.transformer.parameters()).is_cuda:
+            self.pipeline.transformer.to('cpu')
+            print("  [VRAM Manager] Transformer movido para a CPU para a etapa de Upscale.")
+    
+    # Garante que os modelos necessários para o upscale (VAE e Upsampler) estejam na GPU.
+    if not next(self.pipeline.vae.parameters()).is_cuda:
+        self.pipeline.vae.to(self.device)
+        print("  [VRAM Manager] VAE movido para a GPU para a etapa de Upscale.")
+    if self.latent_upsampler and not next(self.latent_upsampler.parameters()).is_cuda:
+         self.latent_upsampler.to(self.device)
+    torch.cuda.empty_cache()
+    # --- Fim do Gerenciamento de VRAM ---
+    
+    low_res_latents = torch.load(latents_path)
+    with torch.no_grad(), torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
+        # Agora o VAE e o latents estarão ambos na GPU
+        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 latents HD: {latents_hd_upscaled_cpu.shape}")
+    del latents_hd_upscaled, low_res_latents
+    torch.cuda.empty_cache()
+
+    # --- ETAPA 2: Fragmentar, Processar em Chunks e Coletar Tensores ---
+    print("\n[LOG] Etapa 2: Fragmentando e processando cada chunk para obter tensores de pixels e latentes...")
+    total_latents = latents_hd_upscaled_cpu.shape[2]
+    sobreposicao = 2
+    num_divisoes = 3
+    tamanho_base = (total_latents - 1) // num_divisoes
+    
+    if tamanho_base < sobreposicao:
+        raise ValueError("Latentes muito curtos para dividir em 3. Use o refinamento direto ou reduza o número de chunks.")
+    
+    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 = []
+    refined_latent_chunks_list = []
+
+    for i, chunk in enumerate(hd_chunks_to_process):
+        print(f"\n  --> Processando 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
+
+        # A chamada a `refine_texture_only` internamente chamará `_set_generation_environment`,
+        # trazendo o Transformer de volta para a GPU e (se otimizado) movendo o VAE para a CPU.
+        pixel_chunk, refined_latent_chunk = self.refine_texture_only(
+            latents_path=chunk_path,
+            prompt=prompt, negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale, seed=chunk_seed,
+            num_decode_chunks=num_decode_chunks_per_worker
+        )
+        pixel_chunks_list.append(pixel_chunk)
+        refined_latent_chunks_list.append(refined_latent_chunk)
+        print(f"  --> Tensores do Chunk {i+1} recebidos.")
+
+    # --- ETAPA 3: Costurar os Tensores (Pixels e Latentes) ---
+    print("\n[LOG] Etapa 3: Costurando os tensores de pixels e latentes...")
+    
+    final_pixel_tensor = self._stitch_chunks(
+        chunks_list=pixel_chunks_list, num_divisoes=num_divisoes,
+        total_original_items=total_latents, sobreposicao_items=sobreposicao,
+        item_multiplier=8, domain_name="Pixels (Final)"
+    )
+
+    final_latents_hd = self._stitch_chunks(
+        chunks_list=refined_latent_chunks_list, num_divisoes=num_divisoes,
+        total_original_items=total_latents, sobreposicao_items=sobreposicao,
+        item_multiplier=1, domain_name="Latentes (Final)"
+    )
+
+    # --- ETAPA 4: Salvar os Artefatos Finais ---
+    print("\n[LOG] Etapa 4: Salvando o vídeo MP4 e o tensor de latentes finais...")
+    
+    final_video_path = self._save_video_from_tensor(
+        final_pixel_tensor, f"final_progressive_stitched_{seed}", 
+        seed, temp_dir_orchestrator 
+    )
+    final_latents_path = self._save_latents_to_disk(
+        final_latents_hd, f"final_hd_stitched_{seed}", 
+        seed
+    )
+
+    # Limpeza final
+    self._finalize()
+
+    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_refinement(
+    def apply_secondary_refinement2(
         self,
         latents_path: str,
         prompt: str,