Update api/ltx_server.py

api/ltx_server.py

@@ -473,6 +473,7 @@ class VideoService:
                 pass  
 
                 
+
     def generate(
         self,
         prompt,
@@ -575,11 +576,10 @@ class VideoService:
             print(f"[DEBUG] media_items shape={tuple(media.shape)}")
 
         latents = None
-        multi_scale_pipeline = None
-
-        # ltx_server.py (dentro da função generate)
-
+        
         try:
+            ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
+
             if improve_texture:
                 if not self.latent_upsampler:
                     raise ValueError("Upscaler espacial não carregado.")
@@ -617,7 +617,6 @@ class VideoService:
                 
                 print(f"[DEBUG] Passo 1: Dimensões reduzidas e ajustadas para {height_p1}x{width_p1}")
                 
-                ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
                 with ctx:
                     first_pass_result = self.pipeline(**first_pass_kwargs)
                     
@@ -656,14 +655,24 @@ class VideoService:
                     "skip_block_list": second_pass_args.get("skip_block_list"),
                 })
                 
-                strength = second_pass_args.get("strength", second_pass_args.get("denoising_strength", 0.7))
-                second_pass_kwargs["strength"] = strength
-                print(f"[DEBUG] Passo 2: Denoising strength definido para {strength}")
+                schedule_p2 = second_pass_args.get("timesteps") or second_pass_args.get("guidance_timesteps")
+                if schedule_p2:
+                    timesteps_para_refinamento = schedule_p2
+                    print(f"[DEBUG] Passo 2: Usando {len(timesteps_para_refinamento)} timesteps pré-definidos do config para refinamento.")
+                else:
+                    strength_p2 = second_pass_args.get("strength", second_pass_args.get("denoising_strength", 0.4))
+                    num_steps_passo2_total = second_pass_args.get("num_inference_steps", 20)
+                    
+                    self.pipeline.scheduler.set_timesteps(num_steps_passo2_total, device=self.device)
+                    todos_os_timesteps_p2 = self.pipeline.scheduler.timesteps
+                    
+                    ponto_de_corte = int(len(todos_os_timesteps_p2) * (1.0 - strength_p2))
+                    timesteps_para_refinamento = todos_os_timesteps_p2[ponto_de_corte:]
+                    print(f"[DEBUG] Passo 2: Calculando {len(timesteps_para_refinamento)} timesteps manuais (strength ≈ {strength_p2})")
 
-                if "timesteps" in second_pass_kwargs:
-                    del second_pass_kwargs["timesteps"]
-                if "guidance_timesteps" in second_pass_kwargs:
-                    del second_pass_kwargs["guidance_timesteps"]
+                second_pass_kwargs["timesteps"] = timesteps_para_refinamento
+                
+                if "strength" in second_pass_kwargs: del second_pass_kwargs["strength"]
 
                 second_pass_kwargs["latents"] = latents_high_res
                 
@@ -672,8 +681,9 @@ class VideoService:
                     
                 latents = second_pass_result.images
                 log_tensor_info(latents, "Latentes Finais (Passo 2)")
-                
+
             else:
+                # --- PASSO ÚNICO (SINGLE-PASS) ---
                 single_pass_kwargs = call_kwargs.copy()
                 first_pass_config = self.config.get("first_pass", {})
                 single_pass_kwargs.update(
@@ -693,30 +703,22 @@ class VideoService:
                 print(f"[DEBUG] Single-pass: timesteps_len={len(schedule) if schedule else 0}")
 
                 print("\n[INFO] Executando pipeline de etapa única...")
-                t_sp = time.perf_counter()
-                ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
                 with ctx:
                     result = self.pipeline(**single_pass_kwargs)
-                print(f"[DEBUG] single-pass tempo={time.perf_counter()-t_sp:.3f}s")
-
-                if hasattr(result, "latents"):
-                    latents = result.latents
-                elif hasattr(result, "images") and isinstance(result.images, torch.Tensor):
-                    latents = result.images
-                else:
-                    latents = result
+                
+                latents = result.images
                 print(f"[DEBUG] Latentes (single-pass): shape={tuple(latents.shape)}")
 
-            # Staging e escrita MP4 (simples: VAE → pixels → MP4)
+            # --- DECODIFICAÇÃO E CODIFICAÇÃO DE VÍDEO FINAL ---
             
             latents_cpu = latents.detach().to("cpu", non_blocking=True)
-            torch.cuda.empty_cache()
-            try:
-                torch.cuda.ipc_collect()
-            except Exception:
-                pass
+            if self.device == "cuda":
+                torch.cuda.empty_cache()
+                try:
+                    torch.cuda.ipc_collect()
+                except Exception:
+                    pass
                 
-            # 2) Divide em duas partes
             lat_a, lat_b = self._dividir_latentes(latents_cpu)
             lat_a1, lat_a2 = self._dividir_latentes(lat_a)
             lat_b1, lat_b2 = self._dividir_latentes(lat_b)
@@ -729,22 +731,19 @@ class VideoService:
             partes_mp4 = []
             par = 0
             
-            for latents in latents_parts:
-                print(f"[DEBUG] Partição {par}: {tuple(latents.shape)}")
-       
+            for part in latents_parts:
+                print(f"[DEBUG] Partição {par}: {tuple(part.shape)}")
                 par = par + 1
                 output_video_path = os.path.join(temp_dir, f"output_{used_seed}_{par}.mp4")
-                final_output_path = None
     
                 print("[DEBUG] Decodificando bloco de latentes com VAE → tensor de pixels...")
-                # Usar manager com timestep por item; previne target_shape e rota NoneType.decode
                 pixel_tensor = vae_manager_singleton.decode(
-                    latents.to(self.device, non_blocking=True),
+                    part.to(self.device, non_blocking=True),
                     decode_timestep=float(self.config.get("decode_timestep", 0.05))
                 )
                 log_tensor_info(pixel_tensor, "Pixel tensor (VAE saída)")
     
-                print("[DEBUG] Codificando MP4 a partir do tensor de pixels (bloco inteiro)...")
+                print("[DEBUG] Codificando MP4 a partir do tensor de pixels...")
                 video_encode_tool_singleton.save_video_from_tensor(
                     pixel_tensor,
                     output_video_path,
@@ -758,9 +757,7 @@ class VideoService:
                     final_output_path = candidate
                     print(f"[DEBUG] MP4 parte {par} movido para {final_output_path}")
                     partes_mp4.append(final_output_path)
-                    
                 except Exception as e:
-                    final_output_path = output_video_path
                     print(f"[DEBUG] Falha no move; usando tmp como final: {e}")
 
             final_concat = os.path.join(results_dir, f"concat_fim_{used_seed}.mp4")
@@ -769,31 +766,23 @@ class VideoService:
             self._log_gpu_memory("Fim da Geração")
             return final_concat, used_seed
             
-
         except Exception as e:
             print("[DEBUG] EXCEÇÃO NA GERAÇÃO:")
             print("".join(traceback.format_exception(type(e), e, e.__traceback__)))
             raise
         finally:
             try:
-                del latents
-            except Exception:
-                pass
-            try:
-                del multi_scale_pipeline
+                del latents, latents_low_res, latents_high_res, second_pass_result
             except Exception:
                 pass
 
             gc.collect()
-            try:
-                if self.device == "cuda":
+            if self.device == "cuda":
+                try:
                     torch.cuda.empty_cache()
-                    try:
-                        torch.cuda.ipc_collect()
-                    except Exception:
-                        pass
-            except Exception as e:
-                print(f"[DEBUG] Limpeza GPU no finally falhou: {e}")
+                    torch.cuda.ipc_collect()
+                except Exception as e:
+                    print(f"[DEBUG] Limpeza GPU no finally falhou: {e}")
 
             try:
                 self.finalize(keep_paths=[])
@@ -801,4 +790,4 @@ class VideoService:
                 print(f"[DEBUG] finalize() no finally falhou: {e}")
 
 print("Criando instância do VideoService. O carregamento do modelo começará agora...")
-video_generation_service = VideoService()
+video_generation_service = VideoService()