Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -452,128 +452,129 @@ class VideoService:
         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.")
+    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")
     
-    # 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 ---
+        temp_dir_orchestrator = tempfile.mkdtemp(prefix="refine_orchestrator_")
+        self._register_tmp_dir(temp_dir_orchestrator)
     
-    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))
+        # --- 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()
     
-    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
+        # --- 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 = []
     
-    if tamanho_base < sobreposicao:
-        raise ValueError("Latentes muito curtos para dividir em 3. Use o refinamento direto ou reduza o número de chunks.")
+        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
     
-    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]
+            # 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.")
     
-    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
+        # --- 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)"
         )
-        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_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)"
+        )
     
-    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
-    )
+        # --- 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    
 
-    # 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_refinement2(
         self,
         latents_path: str,