Update ltx_manager_helpers.py

ltx_manager_helpers.py

@@ -111,4 +111,90 @@ class LtxPoolManager:
         motion_prompt: str, conditioning_items_data: list,
         width: int, height: int, seed: int, cfg: float, video_total_frames: int,
         video_fps: int, num_inference_steps: int, use_attention_slicing: bool,
-        decode_timestep: float, im
+        decode_timestep: float, image_cond_noise_scale: float,
+        current_fragment_index: int, output_path: str, progress
+    ):
+        worker_to_use = None
+        try:
+            with self.lock:
+                if self.last_cleanup_thread and self.last_cleanup_thread.is_alive():
+                    print("LTX POOL MANAGER: Aguardando limpeza da GPU anterior...")
+                    self.last_cleanup_thread.join()
+                    print("LTX POOL MANAGER: Limpeza anterior concluída.")
+
+                worker_to_use = self.workers[self.current_worker_index]
+                previous_worker_index = (self.current_worker_index - 1 + len(self.workers)) % len(self.workers)
+                worker_to_cleanup = self.workers[previous_worker_index]
+
+                cleanup_thread = threading.Thread(target=self._cleanup_worker, args=(worker_to_cleanup,))
+                cleanup_thread.start()
+                self.last_cleanup_thread = cleanup_thread
+                
+                worker_to_use.to_gpu()
+                
+                self.current_worker_index = (self.current_worker_index + 1) % len(self.workers)
+            
+            target_device = worker_to_use.device
+            
+            if use_attention_slicing:
+                worker_to_use.pipeline.enable_attention_slicing()
+
+            media_paths = [item[0] for item in conditioning_items_data]
+            start_frames = [item[1] for item in conditioning_items_data]
+            strengths = [item[2] for item in conditioning_items_data]
+
+            padded_h, padded_w = ((height - 1) // 32 + 1) * 32, ((width - 1) // 32 + 1) * 32
+            padding_vals = calculate_padding(height, width, padded_h, padded_w)
+
+            conditioning_items = prepare_conditioning(
+                conditioning_media_paths=media_paths, conditioning_strengths=strengths,
+                conditioning_start_frames=start_frames, height=height, width=width,
+                num_frames=video_total_frames, padding=padding_vals, pipeline=worker_to_use.pipeline,
+            )
+            
+            for item in conditioning_items:
+                item.media_item = item.media_item.to(target_device)
+
+            first_pass_config = worker_to_use.config.get("first_pass", {}).copy()
+            first_pass_config['num_inference_steps'] = int(num_inference_steps)
+
+            kwargs = {
+                "prompt": motion_prompt, "negative_prompt": "blurry, distorted, bad quality, artifacts",
+                "height": padded_h, "width": padded_w, "num_frames": video_total_frames,
+                "frame_rate": video_fps,
+                "generator": torch.Generator(device=target_device).manual_seed(int(seed) + current_fragment_index),
+                "output_type": "pt", "guidance_scale": float(cfg),
+                "timesteps": first_pass_config.get("timesteps"),
+                "conditioning_items": conditioning_items,
+                "decode_timestep": decode_timestep,
+                "decode_noise_scale": worker_to_use.config.get("decode_noise_scale"),
+                "image_cond_noise_scale": image_cond_noise_scale,
+                "stochastic_sampling": worker_to_use.config.get("stochastic_sampling"),
+                "is_video": True, "vae_per_channel_normalize": True,
+                "mixed_precision": (worker_to_use.config.get("precision") == "mixed_precision"),
+                "enhance_prompt": False, "decode_every": 4, "num_inference_steps": int(num_inference_steps)
+            }
+            
+            # --- Configura os parâmetros dinâmicos do TeaCache antes da geração ---
+            if hasattr(worker_to_use.pipeline.transformer, 'enable_teacache') and worker_to_use.pipeline.transformer.enable_teacache:
+                print(f"LTX POOL MANAGER em {worker_to_use.device}: Configurando TeaCache com num_steps={int(num_inference_steps)}.")
+                worker_to_use.pipeline.transformer.num_steps = int(num_inference_steps)
+                worker_to_use.pipeline.transformer.cnt = 0
+            
+            progress(0.1, desc=f"[Câmera LTX em {worker_to_use.device}] Filmando Cena {current_fragment_index}...")
+            result_tensor = worker_to_use.generate_video_fragment_internal(**kwargs).images
+            
+            pad_l, pad_r, pad_t, pad_b = map(int, padding_vals); slice_h = -pad_b if pad_b > 0 else None; slice_w = -pad_r if pad_r > 0 else None
+            cropped_tensor = result_tensor[:, :, :video_total_frames, pad_t:slice_h, pad_l:slice_w]
+            video_np = (cropped_tensor[0].permute(1, 2, 3, 0).cpu().float().numpy() * 255).astype(np.uint8)
+            
+            with imageio.get_writer(output_path, fps=video_fps, codec='libx264', quality=8) as writer:
+                for frame in video_np: writer.append_data(frame)
+            
+            return output_path, video_total_frames
+
+        finally:
+            if use_attention_slicing and worker_to_use and worker_to_use.pipeline:
+                worker_to_use.pipeline.disable_attention_slicing()
+
+ltx_manager_singleton = LtxPoolManager(device_ids=['cuda:0', 'cuda:1', 'cuda:2', 'cuda:3'])