Update api/ltx_server.py

api/ltx_server.py

@@ -579,85 +579,94 @@ class VideoService:
         latents = None
         multi_scale_pipeline = None
 
+
         try:
             if improve_texture:
                 if not self.latent_upsampler:
                     raise ValueError("Upscaler espacial não carregado.")
-                print("[DEBUG] Multi-escala: construindo pipeline...")
-                multi_scale_pipeline = LTXMultiScalePipeline(self.pipeline, self.latent_upsampler)
+                
+                print("[DEBUG] Multi-escala: Iniciando Passo 1 (geração de latentes base).")
+                
+                # 1. Configurar e executar o primeiro passo
                 first_pass_args = self.config.get("first_pass", {}).copy()
-                first_pass_args["guidance_scale"] = float(guidance_scale)
-                second_pass_args = self.config.get("second_pass", {}).copy()
-                second_pass_args["guidance_scale"] = float(guidance_scale)
-
-                #multi_scale_call_kwargs = call_kwargs.copy()
-                #multi_scale_call_kwargs.update(
-                #    {
-                #        "downscale_factor": self.config["downscale_factor"],
-                #        "first_pass": first_pass_args,
-                #        "second_pass": second_pass_args,
-                #    }
-                #)
-
-                # 
                 first_pass_kwargs = call_kwargs.copy()
-                first_pass_kwargs.update(first_pass_args)
-                second_pass_kwargs = call_kwargs.copy()
-                second_pass_kwargs.update(second_pass_args)
+                first_pass_kwargs.update({
+                    "guidance_scale": float(guidance_scale),
+                    "stg_scale": first_pass_args.get("stg_scale"),
+                    "rescaling_scale": first_pass_args.get("rescaling_scale"),
+                    "skip_block_list": first_pass_args.get("skip_block_list"),
+                })
+                # Opcional: ajustar timesteps se especificado no config
+                schedule = first_pass_args.get("timesteps") or first_pass_args.get("guidance_timesteps")
+                if schedule:
+                    first_pass_kwargs["timesteps"] = schedule
+                    first_pass_kwargs["guidance_timesteps"] = schedule
                 
-                print("[DEBUG] Chamando multi_scale_pipeline...")
-                t_ms = time.perf_counter()
-                ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
+                # Reduzir dimensões para o primeiro passo
+                downscale_factor = self.config.get("downscale_factor", 2)
+                original_height = first_pass_kwargs["height"]
+                original_width = first_pass_kwargs["width"]
+                first_pass_kwargs["height"] = original_height // downscale_factor
+                first_pass_kwargs["width"] = original_width // downscale_factor
                 
-                print("[DEBUG] Executando FIRST PASS (pipeline base)...")
+                print(f"[DEBUG] Passo 1: Dimensões reduzidas para {first_pass_kwargs['height']}x{first_pass_kwargs['width']}")
+                
+                t_p1 = time.perf_counter()
+                ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
                 with ctx:
-                     result_first = self.pipeline(**first_pass_kwargs)
-
-
-                if hasattr(result_first, "latents"):
-                    latents = result_first.latents
-                elif hasattr(result_first, "images") and isinstance(result_first.images, torch.Tensor):
-                    latents = result_first.images
-                else:
-                    latents = result_first
-                print(f"[DEBUG] Latentes (single-pass): shape={tuple(latents.shape)}")
-                latents_first = latents 
-
+                    # Executa a pipeline principal para o primeiro passo
+                    first_pass_result = self.pipeline(**first_pass_kwargs)
+                    
+                # Extrai os latentes do resultado
+                latents_low_res = first_pass_result.latents if hasattr(first_pass_result, "latents") else first_pass_result
+                print(f"[DEBUG] Passo 1 concluído em {time.perf_counter()-t_p1:.3f}s. Shape dos latentes de baixa resolução: {tuple(latents_low_res.shape)}")
+                log_tensor_info(latents_low_res, "Latentes (Passo 1)")
+                
+                # Limpeza de memória entre os passos
+                del first_pass_result, first_pass_kwargs
+                gc.collect()
+                if self.device == "cuda": torch.cuda.empty_cache()
 
-                ctx = torch.autocast(device_type="cuda", dtype=self.runtime_autocast_dtype) if self.device == "cuda" else contextlib.nullcontext()
+                # 2. Upscale dos latentes
+                print("[DEBUG] Multi-escala: Fazendo upscale dos latentes com latent_upsampler.")
+                with ctx:
+                    latents_high_res = self.latent_upsampler(
+                        latents=latents_low_res,
+                        output_height=original_height,
+                        output_width=original_width,
+                    )
+                log_tensor_info(latents_high_res, "Latentes (Pós-Upscale)")
+                del latents_low_res
+                gc.collect()
+                if self.device == "cuda": torch.cuda.empty_cache()
                 
+                # 3. Configurar e executar o segundo passo
+                print("[DEBUG] Multi-escala: Iniciando Passo 2 (refinamento em alta resolução).")
+                second_pass_args = self.config.get("second_pass", {}).copy()
+                second_pass_kwargs = call_kwargs.copy()
+                second_pass_kwargs.update({
+                    "guidance_scale": float(guidance_scale),
+                    "stg_scale": second_pass_args.get("stg_scale"),
+                    "rescaling_scale": second_pass_args.get("rescaling_scale"),
+                    "skip_block_list": second_pass_args.get("skip_block_list"),
+                })
+                # O segundo passo geralmente usa uma fração dos timesteps totais (ex: 70%)
+                schedule_p2 = second_pass_args.get("timesteps") or second_pass_args.get("guidance_timesteps")
+                if schedule_p2:
+                    second_pass_kwargs["timesteps"] = schedule_p2
+                    second_pass_kwargs["guidance_timesteps"] = schedule_p2
+                    
+                # Adiciona os latentes do upscale como 'latents' iniciais para o segundo passo
+                second_pass_kwargs["latents"] = latents_high_res
                 
-                # --- SECOND PASS ---
-                print("[DEBUG] Executando SECOND PASS (latent_upsampler)...")
+                t_p2 = time.perf_counter()
                 with ctx:
-                    result = self.latent_upsampler(
-                        #**second_pass_kwargs,
-                        #conditions=conditioning_items,
-                        prompt=prompt,
-                        negative_prompt=negative_prompt,
-                        width=width_padded,
-                        height=height_padded,
-                        num_frames=actual_num_frames,
-                        latents=latents_first,
-                        denoise_strength=0.4,  
-                        num_inference_steps=10,
-                        decode_timestep=0.05,
-                        image_cond_noise_scale=0.025,
-                        generator=generator,
-                        output_type="latent",
-                    )
+                    # Executa a pipeline principal para o segundo passo
+                    second_pass_result = self.pipeline(**second_pass_kwargs)
+                    
+                latents = second_pass_result.latents if hasattr(second_pass_result, "latents") else second_pass_result
+                print(f"[DEBUG] Passo 2 concluído em {time.perf_counter()-t_p2:.3f}s. Shape dos latentes finais: {tuple(latents.shape)}")
                 
-                #with ctx:
-                #    result = multi_scale_pipeline(**multi_scale_call_kwargs)
-                #print(f"[DEBUG] multi_scale_pipeline tempo={time.perf_counter()-t_ms:.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
-                print(f"[DEBUG] Latentes (multi-escala): shape={tuple(latents.shape)}")
             else:
                 single_pass_kwargs = call_kwargs.copy()
                 first_pass_config = self.config.get("first_pass", {})