Update api/ltx_server.py

api/ltx_server.py

@@ -769,6 +769,8 @@ class VideoService:
             print("[DEBUG] EXCEÇÃO NA GERAÇÃO:")
             print("".join(traceback.format_exception(type(e), e, e.__traceback__)))
             raise
+
+
 # ltx_server.py
 
     def generate(
@@ -789,7 +791,7 @@ class VideoService:
         frames_to_use=9,
         seed=42,
         randomize_seed=True,
-        guidance_scale=3.0,
+        guidance_scale=3.0, # Valor de referência/fallback
         improve_texture=True,
         progress_callback=None,
         external_decode=True,
@@ -846,16 +848,16 @@ class VideoService:
             "output_type": "latent",
             "conditioning_items": conditioning_items if conditioning_items else None,
             "media_items": None,
-            "decode_timestep": self.config["decode_timestep"],
-            "decode_noise_scale": self.config["decode_noise_scale"],
-            "stochastic_sampling": self.config["stochastic_sampling"],
-            "image_cond_noise_scale": 0.01,
+            "decode_timestep": self.config.get("decode_timestep"),
+            "decode_noise_scale": self.config.get("decode_noise_scale"),
+            "stochastic_sampling": self.config.get("stochastic_sampling"),
+            "image_cond_noise_scale": self.config.get("image_cond_noise_scale", 0.01),
             "is_video": True,
-            "vae_per_channel_normalize": True,
-            "mixed_precision": (self.config["precision"] == "mixed_precision"),
+            "vae_per_channel_normalize": self.config.get("vae_per_channel_normalize", True),
+            "mixed_precision": (self.config.get("precision") == "mixed_precision"),
             "offload_to_cpu": False,
             "enhance_prompt": False,
-            "skip_layer_strategy": SkipLayerStrategy.AttentionValues,
+            "skip_layer_strategy": SkipLayerStrategy[self.config.get("stg_mode", "AttentionValues")],
         }
         print(f"[DEBUG] output_type={call_kwargs['output_type']} skip_layer_strategy={call_kwargs['skip_layer_strategy']}")
 
@@ -884,32 +886,29 @@ class VideoService:
                 
                 first_pass_args = self.config.get("first_pass", {}).copy()
                 first_pass_kwargs = call_kwargs.copy()
+
                 first_pass_kwargs.update({
-                    "guidance_scale": float(guidance_scale),
+                    "guidance_scale": first_pass_args.get("guidance_scale", 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"),
+                    "guidance_timesteps": first_pass_args.get("guidance_timesteps"),
+                    "timesteps": first_pass_args.get("timesteps")
                 })
-                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(f"[DEBUG] Passo 1: Parâmetros do config carregados: guidance_scale={first_pass_kwargs['guidance_scale']}, stg_scale={first_pass_kwargs['stg_scale']}")
                 
                 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
                 height_p1 = round(target_height_p1 / divisor) * divisor
                 if height_p1 == 0: height_p1 = divisor
                 first_pass_kwargs["height"] = height_p1
-
                 target_width_p1 = original_width // downscale_factor
                 width_p1 = round(target_width_p1 / divisor) * divisor
                 if width_p1 == 0: width_p1 = divisor
                 first_pass_kwargs["width"] = width_p1
-                
                 print(f"[DEBUG] Passo 1: Dimensões reduzidas e ajustadas para {height_p1}x{width_p1}")
                 
                 with ctx:
@@ -937,43 +936,23 @@ class VideoService:
                 second_pass_args = self.config.get("second_pass", {}).copy()
                 second_pass_kwargs = call_kwargs.copy()
 
-                height_p2 = height_p1 * 2
-                width_p2 = width_p1 * 2
-                second_pass_kwargs["height"] = height_p2
-                second_pass_kwargs["width"] = width_p2
-                print(f"[DEBUG] Passo 2: Dimensões definidas para {height_p2}x{width_p2} para corresponder ao upscale.")
-
                 second_pass_kwargs.update({
-                    "guidance_scale": float(guidance_scale),
+                    "guidance_scale": second_pass_args.get("guidance_scale", 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"),
+                    "guidance_timesteps": second_pass_args.get("guidance_timesteps"),
+                    "timesteps": second_pass_args.get("timesteps")
                 })
+                print(f"[DEBUG] Passo 2: Parâmetros do config carregados: guidance_scale={second_pass_kwargs['guidance_scale']}, stg_scale={second_pass_kwargs['stg_scale']}")
                 
-                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})")
-
-                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
+                height_p2 = height_p1 * 2
+                width_p2 = width_p1 * 2
+                second_pass_kwargs["height"] = height_p2
+                second_pass_kwargs["width"] = width_p2
+                print(f"[DEBUG] Passo 2: Dimensões definidas para {height_p2}x{width_p2}")
                 
-                num_timesteps_p2 = len(timesteps_para_refinamento)
-                if 'guidance_mapping' not in second_pass_kwargs:
-                     second_pass_kwargs['guidance_mapping'] = list(range(num_timesteps_p2))
-                print(f"[DEBUG] Passo 2: Injetando 'guidance_mapping' de identidade com {num_timesteps_p2} passos.")
+                second_pass_kwargs["latents"] = latents_high_res
 
                 with ctx:
                     second_pass_result = self.pipeline(**second_pass_kwargs)
@@ -984,20 +963,16 @@ class VideoService:
             else:
                 # --- PASSO ÚNICO (SINGLE-PASS) ---
                 single_pass_kwargs = call_kwargs.copy()
-                first_pass_config = self.config.get("first_pass", {})
+                
                 single_pass_kwargs.update({
-                    "guidance_scale": float(guidance_scale),
-                    "stg_scale": first_pass_config.get("stg_scale"),
-                    "rescaling_scale": first_pass_config.get("rescaling_scale"),
-                    "skip_block_list": first_pass_config.get("skip_block_list"),
+                    "guidance_scale": self.config.get("guidance_scale", guidance_scale),
+                    "stg_scale": self.config.get("stg_scale"),
+                    "rescaling_scale": self.config.get("rescaling_scale"),
+                    "skip_block_list": self.config.get("skip_block_list"),
+                    "guidance_timesteps": self.config.get("guidance_timesteps"),
+                    "timesteps": self.config.get("timesteps"),
+                    "num_inference_steps": self.config.get("num_inference_steps", 20)
                 })
-                schedule = first_pass_config.get("timesteps") or first_pass_config.get("guidance_timesteps")
-                if mode == "video-to-video":
-                    schedule = [0.7]; print("[INFO] Modo video-to-video (etapa única): timesteps=[0.7]")
-                if isinstance(schedule, (list, tuple)) and len(schedule) > 0:
-                    single_pass_kwargs["timesteps"] = schedule
-                    single_pass_kwargs["guidance_timesteps"] = schedule
-                print(f"[DEBUG] Single-pass: timesteps_len={len(schedule) if schedule else 0}")
 
                 print("\n[INFO] Executando pipeline de etapa única...")
                 with ctx:
@@ -1012,22 +987,28 @@ class VideoService:
                 torch.cuda.empty_cache()
                 try: torch.cuda.ipc_collect()
                 except Exception: pass
-                
-            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)
             
-            latents_parts = [lat_a1, lat_a2, lat_b1, lat_b2]
+            lat_a, lat_b = self._dividir_latentes(latents_cpu)
+            if lat_a is not None:
+                lat_a1, lat_a2 = self._dividir_latentes(lat_a)
+            else:
+                lat_a1, lat_a2 = None, None
+            if lat_b is not None:
+                lat_b1, lat_b2 = self._dividir_latentes(lat_b)
+            else:
+                lat_b1, lat_b2 = None, None
             
+            latents_parts = [p for p in [lat_a1, lat_a2, lat_b1, lat_b2] if p is not None]
+            if not latents_parts:
+                latents_parts = [latents_cpu]
+
             temp_dir = tempfile.mkdtemp(prefix="ltxv_"); self._register_tmp_dir(temp_dir)
             results_dir = "/app/output"; os.makedirs(results_dir, exist_ok=True)
-
             partes_mp4 = []
             par = 0
             
             for part in latents_parts:
                 par += 1
-                if part is None: continue
                 print(f"[DEBUG] Partição {par}: {tuple(part.shape)}")
                 output_video_path = os.path.join(temp_dir, f"output_{used_seed}_{par}.mp4")
     
@@ -1074,10 +1055,12 @@ class VideoService:
             print("".join(traceback.format_exception(type(e), e, e.__traceback__)))
             raise
         finally:
-            # Limpeza de variáveis para liberar memória
-            try: del latents, latents_low_res, latents_high_res, second_pass_result
-            except NameError: pass
-            except Exception as e: print(f"[DEBUG] Erro na limpeza de variáveis: {e}")
+            try:
+                del latents, latents_low_res, latents_high_res, second_pass_result, first_pass_result, result
+            except NameError:
+                pass
+            except Exception as e:
+                print(f"[DEBUG] Erro na limpeza de variáveis: {e}")
 
             gc.collect()
             if self.device == "cuda":
@@ -1092,5 +1075,6 @@ class VideoService:
             except Exception as e:
                 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()