Upload ltx_server.py

api/ltx_server.py

@@ -575,100 +575,41 @@ class VideoService:
             print(f"[DEBUG] media_items shape={tuple(media.shape)}")
 
         latents = None
-        
+        multi_scale_pipeline = None
+
         try:
             if improve_texture:
                 if not self.latent_upsampler:
                     raise ValueError("Upscaler espacial não carregado.")
-                
-                # --- INÍCIO DA SEPARAÇÃO DOS PASSOS ---
-                
-                print("[DEBUG] Multi-escala: Iniciando Passo 1 (geração de latentes base).")
-                
-                # 1. Configurar e executar o primeiro passo
+                print("[DEBUG] Multi-escala: construindo pipeline...")
+                multi_scale_pipeline = LTXMultiScalePipeline(self.pipeline, self.latent_upsampler)
                 first_pass_args = self.config.get("first_pass", {}).copy()
-                first_pass_kwargs = call_kwargs.copy()
-                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"),
-                })
-                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
-                
-                # Reduzir dimensões para o primeiro passo, garantindo divisibilidade por 24
-                downscale_factor = self.config.get("downscale_factor", 2)
-                original_height = first_pass_kwargs["height"]
-                original_width = first_pass_kwargs["width"]
-                divisor = 24
-
-                target_height_p1 = original_height // downscale_factor
-                first_pass_kwargs["height"] = round(target_height_p1 / divisor) * divisor
+                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)
 
-                target_width_p1 = original_width // downscale_factor
-                first_pass_kwargs["width"] = round(target_width_p1 / divisor) * divisor
-                
-                if first_pass_kwargs["height"] == 0: first_pass_kwargs["height"] = divisor
-                if first_pass_kwargs["width"] == 0: first_pass_kwargs["width"] = divisor
-                
-                print(f"[DEBUG] Passo 1: Dimensões reduzidas e ajustadas para {first_pass_kwargs['height']}x{first_pass_kwargs['width']}")
-                
-                t_p1 = time.perf_counter()
+                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,
+                    }
+                )
+                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()
                 with ctx:
-                    first_pass_result = self.pipeline(**first_pass_kwargs)
-                    
-                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)")
-                
-                del first_pass_result
-                gc.collect()
-                if self.device == "cuda": torch.cuda.empty_cache()
-
-                # 2. Upscale dos latentes
-                print("[DEBUG] Multi-escala: Fazendo upscale dos latentes com latent_upsampler.")
-                with ctx:
-                    # Chamada corrigida: posicional, sem argumentos de palavra-chave extras
-                    latents_high_res = self.latent_upsampler(latents_low_res)
+                    result = multi_scale_pipeline(**multi_scale_call_kwargs)
+                print(f"[DEBUG] multi_scale_pipeline tempo={time.perf_counter()-t_ms:.3f}s")
 
-                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["height"] = original_height
-                second_pass_kwargs["width"] = original_width
-                
-                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"),
-                })
-                
-                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
-                    
-                second_pass_kwargs["latents"] = latents_high_res
-                
-                t_p2 = time.perf_counter()
-                with ctx:
-                    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)}")
-                
+                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", {})
@@ -775,7 +716,11 @@ class VideoService:
                 del latents
             except Exception:
                 pass
-            
+            try:
+                del multi_scale_pipeline
+            except Exception:
+                pass
+
             gc.collect()
             try:
                 if self.device == "cuda":
@@ -793,4 +738,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()