Add mask_drop_step for single-pass harmonize

minimax_h3_inpaint_blocks.py

@@ -655,6 +655,11 @@ class MiniMaxH3InpaintSetTimestepsStep(MiniMaxH3SetTimestepsStep):
             InputParam(
                 name="inpaint_audio_row_mask", type_hint=torch.Tensor, required=True, description="Per audio row."
             ),
+            InputParam(
+                name="mask_drop_step", type_hint=int,
+                description="If set, the video mask is dropped (all rows become targets) from this step on, so the "
+                "preserved subject free-refines into the scene over the final steps. None keeps the mask all the way.",
+            ),
         ]
 
     @staticmethod
@@ -701,6 +706,8 @@ class MiniMaxH3InpaintSetTimestepsStep(MiniMaxH3SetTimestepsStep):
         block_state.timesteps = components.scheduler.timesteps
         block_state.audio_timesteps = components.audio_scheduler.timesteps
 
+        drop = block_state.mask_drop_step
+        ones_video = torch.ones_like(block_state.inpaint_row_mask)
         block_state.row_timestep_plan = [
             tuple(
                 tensor.to(device)
@@ -714,11 +721,12 @@ class MiniMaxH3InpaintSetTimestepsStep(MiniMaxH3SetTimestepsStep):
                     float(audio_timestep),
                     max(float(timestep), components.keyframe_noise_aug),
                     AUDIO_COND_TIMESTEP,
-                    block_state.inpaint_row_mask,
+                    # after the drop step every video row is a free target; audio stays preserved either way
+                    (ones_video if (drop is not None and i >= drop) else block_state.inpaint_row_mask),
                     block_state.inpaint_audio_row_mask,
                 )
             )
-            for timestep, audio_timestep in zip(block_state.timesteps, block_state.audio_timesteps)
+            for i, (timestep, audio_timestep) in enumerate(zip(block_state.timesteps, block_state.audio_timesteps))
         ]
 
         self.set_block_state(state, block_state)
@@ -756,6 +764,7 @@ class MiniMaxH3InpaintLoopSchedulerStep(MiniMaxH3LoopSchedulerStep):
             InputParam(
                 name="inpaint_audio_row_mask", type_hint=torch.Tensor, required=True, description="Per audio row."
             ),
+            InputParam(name="mask_drop_step", type_hint=int, description="See the set-timesteps step."),
         ]
 
     @torch.no_grad()
@@ -765,17 +774,27 @@ class MiniMaxH3InpaintLoopSchedulerStep(MiniMaxH3LoopSchedulerStep):
         # On the last step the schedule has reached the clean end, and so does the source: the preserved rows are
         # handed to the decoder as the footage itself rather than as the anchor's 0.999.
         last = i + 1 >= block_state.timesteps.numel()
-        video_level = 1.0 if last else VISUAL_COND_TIMESTEP
+        drop = block_state.mask_drop_step
 
         num_condition_video_rows = block_state.num_condition_video_rows
-        block_state.latents[num_condition_video_rows:] = impose_source(
-            block_state.latents[num_condition_video_rows:],
-            block_state.source_rows,
-            block_state.inpaint_noise_rows,
-            block_state.inpaint_row_mask,
-            video_level,
-            components.scheduler,
-        )
+
+        # Video write-back. With a mask-drop, stop imposing the source from step `drop` on (the subject then
+        # free-refines into the scene), and on the handoff step `drop - 1` re-noise the preserved rows to step
+        # `drop`'s level so the next forward reads them on distribution. Without a drop, the normal inpaint write-back.
+        if drop is None or i < drop:
+            if drop is not None and i == drop - 1:
+                video_level = float(block_state.timesteps[drop])
+            else:
+                video_level = 1.0 if last else VISUAL_COND_TIMESTEP
+            block_state.latents[num_condition_video_rows:] = impose_source(
+                block_state.latents[num_condition_video_rows:],
+                block_state.source_rows,
+                block_state.inpaint_noise_rows,
+                block_state.inpaint_row_mask,
+                video_level,
+                components.scheduler,
+            )
+        # Audio is always preserved: we harmonize the picture, not the soundtrack.
         if block_state.source_audio_rows is not None:
             num_condition_audio_rows = block_state.num_condition_audio_rows
             block_state.audio_latents[num_condition_audio_rows:] = impose_source(