Update managers/ltx_manager.py

managers/ltx_manager.py

@@ -139,7 +139,7 @@ class LtxPoolManager:
     def _prepare_pipeline_params(self, worker: 'LtxWorker', **kwargs) -> dict:
         pipeline_params = {
             "height": kwargs['height'], "width": kwargs['width'], "num_frames": kwargs['video_total_frames'],
-            "callback" : kwargs['callback']
+            "callback_on_step_end" : kwargs['callback_on_step_end']
             "frame_rate": kwargs.get('video_fps', 24),
             "generator": torch.Generator(device=worker.device).manual_seed(int(time.time()) + kwargs.get('current_fragment_index', 0)),
             "is_video": True, "vae_per_channel_normalize": True,
@@ -172,7 +172,7 @@ class LtxPoolManager:
             height, width = kwargs['height'], kwargs['width']
             padded_h, padded_w = ((height - 1) // 32 + 1) * 32, ((width - 1) // 32 + 1) * 32
             padding_vals = calculate_padding(height, width, padded_h, padded_w)
-            kwargs['callback'],
+            kwargs['callback_on_step_end'],
             kwargs['height'], kwargs['width'] = padded_h, padded_w
             pipeline_params = self._prepare_pipeline_params(worker_to_use, **kwargs)
             logger.info(f"Initiating GENERATION on {worker_to_use.device} with shape {padded_w}x{padded_h}")