Update api/ltx_server_refactored.py

api/ltx_server_refactored.py

@@ -509,7 +509,7 @@ class VideoService:
         print("[SUCCESS] ETAPA 1 Concluída.")
         return final_video_path, final_latents_path, used_seed
 
-    
+
     def refine_texture_only(
         self,
         latents_path: str,
@@ -518,6 +518,150 @@ class VideoService:
         guidance_scale: float,
         seed: Optional[int] = None,
         conditioning_items: Optional[List[ConditioningItem]] = None
+    ) -> Tuple[str, str]:
+        """
+        ETAPA 2: Refina a textura dos latentes existentes.
+        Implementa uma decodificação VAE com janela deslizante para processar vídeos longos
+        sem estourar a memória da GPU, preservando a continuidade temporal.
+        """
+        print("\n======================================================================")
+        print("====== [INFO] Iniciando ETAPA 2: Refinamento de Textura Detalhado ======")
+        print("======================================================================\n")
+    
+        # --- FASE 1: Configuração e Preparação ---
+        print("[LOG] FASE 1: Configuração e Preparação")
+        used_seed = random.randint(0, 2**32 - 1) if seed is None else int(seed)
+        print(f"  [LOG] Seed a ser utilizado (consistente com Etapa 1): {used_seed}")
+    
+        temp_dir = tempfile.mkdtemp(prefix="ltxv_refine_")
+        self._register_tmp_dir(temp_dir)
+        print(f"  [LOG] Diretório temporário para artefatos criado em: {temp_dir}")
+        
+        print(f"  [LOG] Carregando latentes de baixa resolução de: {latents_path}")
+        latents_to_refine = torch.load(latents_path).to(self.device)
+        print(f"  [LOG] Latentes carregados para a GPU. Shape: {latents_to_refine.shape}, Dtype: {latents_to_refine.dtype}")
+        
+        if conditioning_items:
+            print(f"  [LOG] {len(conditioning_items)} item(ns) de condicionamento serão utilizados no refinamento.")
+    
+        # --- FASE 2: Execução do Refinamento (Transformer) ---
+        print("\n[LOG] FASE 2: Execução do Refinamento (Transformer Pass)")
+        with torch.autocast(device_type=self.device.split(':')[0], dtype=self.runtime_autocast_dtype, enabled=(self.device == 'cuda')):
+            
+            refine_height = latents_to_refine.shape[3] * self.pipeline.vae_scale_factor
+            refine_width = latents_to_refine.shape[4] * self.pipeline.vae_scale_factor
+            print(f"  [LOG] Dimensões alvo para o refinamento: {refine_height}x{refine_width}")
+            
+            second_pass_kwargs = {
+                "prompt": prompt, "negative_prompt": negative_prompt,
+                "height": refine_height, "width": refine_width,
+                "frame_rate": int(DEFAULT_FPS), "num_frames": latents_to_refine.shape[2],
+                "latents": latents_to_refine, "guidance_scale": float(guidance_scale),
+                "output_type": "latent",
+                "generator": torch.Generator(device=self.device).manual_seed(used_seed),
+                "conditioning_items": conditioning_items,
+                **(self.config.get("second_pass", {}))
+            }
+            
+            print("  [LOG] Enviando tensor completo para a pipeline de refinamento (Transformer)...")
+            final_latents = self.pipeline(**second_pass_kwargs).images
+            print(f"  [LOG] [SUCESSO] Latentes refinados pelo Transformer. Shape resultante: {final_latents.shape}")
+    
+        # --- FASE 3: Decodificação VAE com Janela Deslizante ---
+        print("\n[LOG] FASE 3: Decodificação VAE com Lógica de Janela Deslizante")
+        
+        print("  [LOG] Gerenciamento de Memória: Movendo tensor de latentes refinados para a CPU...")
+        final_latents_cpu = final_latents.cpu()
+        print("  [LOG] Gerenciamento de Memória: Liberando tensores da VRAM para maximizar espaço para o VAE.")
+        del final_latents, latents_to_refine
+        torch.cuda.empty_cache()
+    
+        total_latents = final_latents_cpu.shape[2]
+        
+        # Caso 1: Vídeo curto. Processo simplificado.
+        if total_latents <= 10:
+            print(f"  [LOG] Detecção: Vídeo curto ({total_latents} latentes <= 10). Usando decodificação direta.")
+            pixel_tensor = vae_manager_singleton.decode(
+                final_latents_cpu.to(self.device),
+                decode_timestep=float(self.config.get("decode_timestep", 0.05))
+            ).cpu()
+    
+        # Caso 2: Vídeo longo. Aplica sua lógica de 3 grupos.
+        else:
+            print(f"  [LOG] Detecção: Vídeo longo ({total_latents} latentes). Ativando modo de janela deslizante.")
+            # 3.1: Cálculo dos Pontos de Corte
+            print("    [LOG] 3.1: Calculando janelas de decodificação (grupos)...")
+            sobreposicao = 2
+            tamanho_base = (total_latents - 1) // 3
+            print(f"      [CALC] Total Latentes: {total_latents}, Sobreposição: {sobreposicao}, Tamanho Base por Grupo: {tamanho_base}")
+            
+            pontos_de_corte = [
+                (0, tamanho_base + sobreposicao),
+                (tamanho_base - sobreposicao, (2 * tamanho_base) + sobreposicao),
+                ((2 * tamanho_base) - sobreposicao, total_latents)
+            ]
+            print(f"      [CALC] Pontos de corte definidos: Grupo 1 ({pontos_de_corte[0][0]}-{pontos_de_corte[0][1]-1}), Grupo 2 ({pontos_de_corte[1][0]}-{pontos_de_corte[1][1]-1}), Grupo 3 ({pontos_de_corte[2][0]}-{pontos_de_corte[2][1]-1})")
+    
+            # 3.2: Decodificação dos Chunks em Loop
+            print("\n    [LOG] 3.2: Decodificando cada grupo em loop...")
+            pixel_chunks_list = []
+            for i, (start, end) in enumerate(pontos_de_corte):
+                latent_chunk = final_latents_cpu[:, :, start:end, :, :]
+                print(f"      [LOOP] Processando Grupo {i+1}/{len(pontos_de_corte)} (latentes {start} a {end-1}), shape: {latent_chunk.shape}")
+                
+                pixel_chunk = vae_manager_singleton.decode(
+                    latent_chunk.to(self.device),
+                    decode_timestep=float(self.config.get("decode_timestep", 0.05))
+                )
+                print(f"        [VAE] Grupo {i+1} decodificado. Shape de pixels resultante: {pixel_chunk.shape}")
+                pixel_chunks_list.append(pixel_chunk.cpu())
+                print(f"        [VRAM] Resultado movido para CPU. Limpando VRAM para o próximo grupo.")
+                torch.cuda.empty_cache()
+    
+            # 3.3: Costura Inteligente dos Vídeos
+            print("\n    [LOG] 3.3: Costurando os vídeos decodificados para montagem final...")
+            
+            frames_p1 = tamanho_base * 8
+            parte1 = pixel_chunks_list[0][:, :, :frames_p1, :, :]
+            print(f"      [STITCH] Parte 1 (do Grupo 1): Pegando os primeiros {frames_p1} frames. Shape: {parte1.shape}")
+    
+            descarte_inicio_p2 = sobreposicao * 8
+            frames_p2 = tamanho_base * 8
+            parte2 = pixel_chunks_list[1][:, :, descarte_inicio_p2 : descarte_inicio_p2 + frames_p2, :, :]
+            print(f"      [STITCH] Parte 2 (do Grupo 2): Descartando {descarte_inicio_p2} frames, pegando os próximos {frames_p2}. Shape: {parte2.shape}")
+    
+            descarte_inicio_p3 = sobreposicao * 8
+            parte3 = pixel_chunks_list[2][:, :, descarte_inicio_p3:, :, :]
+            print(f"      [STITCH] Parte 3 (do Grupo 3): Descartando {descarte_inicio_p3} frames, pegando todo o resto. Shape: {parte3.shape}")
+            
+            pixel_tensor = torch.cat([parte1, parte2, parte3], dim=2)
+    
+        print(f"\n[LOG] [SUCESSO] Tensor de pixels final montado com shape: {pixel_tensor.shape}")
+    
+        # --- FASE 4: Salvamento e Limpeza ---
+        print("\n[LOG] FASE 4: Salvamento dos Artefatos Finais e Limpeza")
+        video_path_out = self._save_video_from_tensor(pixel_tensor, "refined_video_final", used_seed, temp_dir)
+        print(f"  [LOG] Vídeo final salvo em: {video_path_out}")
+        latents_path_out = self._save_latents_to_disk(final_latents_cpu, "latents_refined_final", used_seed)
+        print(f"  [LOG] Latentes refinados finais (CPU) salvos em: {latents_path_out}")
+            
+        print("  [LOG] Liberando tensores finais da memória.")
+        del pixel_tensor, final_latents_cpu
+        self._finalize()
+        
+        print("\n======================================================================")
+        print("============ [SUCCESS] ETAPA 2 Concluída com Sucesso =============")
+        print("======================================================================\n")
+        return video_path_out, latents_path_out
+    
+    def refine_texture_only1(
+        self,
+        latents_path: str,
+        prompt: str,
+        negative_prompt: str,
+        guidance_scale: float,
+        seed: Optional[int] = None,
+        conditioning_items: Optional[List[ConditioningItem]] = None
     ) -> Tuple[str, str]:
         """
         ETAPA 2: Refina a textura dos latentes existentes SEM alterar sua resolução