Update h3_dpmpp_2s_ancestral.py

h3_dpmpp_2s_ancestral.py

@@ -377,7 +377,10 @@ class MiniMaxH3DPMppSdeStep(ModularPipelineBlocks):
     def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
         ncv = block_state.num_condition_video_rows
         nca = block_state.num_condition_audio_rows
-        is_last_step = i == len(block_state.row_timestep_plan) - 1
+        # Not the loop's own last index — for a Staged Denoising slice that's this *stage's* last step, not
+        # necessarily the trajectory's. `components.scheduler.sigmas[i + 1] == 0.0` is the real end regardless
+        # of how many steps this particular call's schedule holds.
+        is_last_step = float(components.scheduler.sigmas[i + 1]) == 0.0
         is_first_step = i == 0
 
         unique_timesteps, timestep_indices = block_state.row_timestep_plan[i]
@@ -497,7 +500,10 @@ class MiniMaxH3Seeds2Step(ModularPipelineBlocks):
     def __call__(self, components: MiniMaxH3ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
         ncv = block_state.num_condition_video_rows
         nca = block_state.num_condition_audio_rows
-        is_last_step = i == len(block_state.row_timestep_plan) - 1
+        # Not the loop's own last index — for a Staged Denoising slice that's this *stage's* last step, not
+        # necessarily the trajectory's. `components.scheduler.sigmas[i + 1] == 0.0` is the real end regardless
+        # of how many steps this particular call's schedule holds.
+        is_last_step = float(components.scheduler.sigmas[i + 1]) == 0.0
         is_first_step = i == 0
 
         unique_timesteps, timestep_indices = block_state.row_timestep_plan[i]