Fix MiniMax H3 denoise masks (#15988)

comfy/ldm/minimax/model.py

@@ -588,6 +588,12 @@ class MiniMaxH3Model(nn.Module):
         else:
             out = graph_out
 
+        # Masked rows predict at mask * sigma; scale their velocity to match the outer x0 conversion.
+        if denoise_mask is not None:
+            out[0] = out[0] * denoise_mask
+        if audio_denoise_mask is not None:
+            out[1] = out[1] * audio_denoise_mask
+
         if scale != 1.0:
             # d/d(sigma_v) of the carried variable
             out[1] = ((1.0 - scale) * (audio_src * carry)

tests-unit/comfy_test/test_minimax_h3_model.py

@@ -0,0 +1,61 @@
+import torch
+from torch import nn
+
+from comfy.ldm.minimax.model import MiniMaxH3Model, time_shift_sigma
+from comfy.model_sampling import CONST
+
+
+def make_model(video_output, audio_output):
+    model = MiniMaxH3Model.__new__(MiniMaxH3Model)
+    nn.Module.__init__(model)
+    model.sigma_shift_video = 12.0
+    model.sigma_shift_audio = 3.0
+    model._forward = lambda *args, **kwargs: [video_output.clone(), audio_output.clone()]
+    return model
+
+
+def test_forward_scales_velocity_to_mask_timestep():
+    video_output = torch.full((1, 2, 1, 2, 2), 2.0)
+    audio_output = torch.full((1, 2, 2, 3), 3.0)
+    video_mask = torch.tensor([[[[[1.0, 0.75], [0.5, 0.25]]]]])
+    audio_mask = torch.tensor([[[[1.0, 0.5, 0.25], [0.75, 0.5, 0.0]]]])
+    sigma = torch.tensor([0.5])
+    clean = torch.arange(video_output.numel(), dtype=torch.float32).reshape_as(video_output)
+    model_input = clean + sigma.reshape(1, 1, 1, 1, 1) * video_mask * video_output
+    model = make_model(video_output, audio_output)
+
+    out = model(
+        [model_input, torch.zeros_like(audio_output)],
+        sigma * 1000.0,
+        torch.empty(1, 1, 1),
+        minimax_payload={"audio_scale": 1.0},
+        denoise_mask=video_mask,
+        audio_denoise_mask=audio_mask,
+    )
+
+    torch.testing.assert_close(out[0], video_output * video_mask)
+    torch.testing.assert_close(out[1], audio_output * audio_mask)
+    denoised = CONST.calculate_denoised(None, sigma, out[0], model_input)
+    torch.testing.assert_close(denoised, clean)
+
+
+def test_forward_scales_audio_velocity_before_carry_conversion():
+    video_output = torch.ones((1, 1, 1, 1, 1))
+    audio_output = torch.full((1, 1, 2, 2), 3.0)
+    audio_src = torch.full_like(audio_output, 2.0)
+    audio_mask = torch.tensor([[[[0.75, 0.5], [0.25, 0.0]]]])
+    model = make_model(video_output, audio_output)
+    sigma_v = torch.tensor(0.5)
+    sigma_a = time_shift_sigma(sigma_v, 12.0, 3.0)
+    carry = sigma_a / sigma_v
+
+    out = model(
+        [torch.zeros_like(video_output), audio_src],
+        sigma_v.reshape(1) * 1000.0,
+        torch.empty(1, 1, 1),
+        minimax_payload={"audio_scale": 4.0},
+        audio_denoise_mask=audio_mask,
+    )
+
+    expected = -3.0 * audio_src * carry + (1.0 + 3.0 * sigma_a) * audio_output * audio_mask
+    torch.testing.assert_close(out[1], expected)