Merge ae6551d70d into 37ac9ff44f

comfy/ldm/minimax/model.py

@@ -554,9 +554,9 @@ class MiniMaxH3Model(nn.Module):
         # model_base passes model_sampling.timestep(sigma) = sigma * 1000
         shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
         shift_a = float(transformer_options.get("minimax_h3_sigma_shift_audio", self.sigma_shift_audio))
-        sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6)
-        t_v = float(1.0 - sigma_v)
-        t_a = float(1.0 - time_shift_sigma(sigma_v, shift_v, shift_a))
+        sigma_v = float((timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6))
+        t_v = 1.0 - sigma_v
+        t_a = 1.0 - time_shift_sigma(sigma_v, shift_v, shift_a)
 
         # distinct timesteps are known analytically: text/pad follow video, cond rows pin near 1
         vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP))
@@ -572,12 +572,16 @@ class MiniMaxH3Model(nn.Module):
         seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "cond_audio": 2, "ref_audio": 2}
 
         text_tags = payload.get("text_token_tags")
+        tags = payload.get("_text_token_tags_list")
+        if text_tags is not None and tags is None:
+            # pull the tag list to host once per sampling run instead of once per step
+            tags = text_tags.view(-1).tolist()
+            payload["_text_token_tags_list"] = tags
         mod_segments = []
         for a, b, kind in layout.segments:
             row_base = t_row[seg_t[kind]] * 3
-            if kind == "text" and text_tags is not None:
+            if kind == "text" and tags is not None:
                 # the presentation text span mixes tags (vision pads carry the video modality) split into tag runs
-                tags = text_tags.view(-1).tolist()
                 run_start = 0
                 for i in range(1, b - a + 1):
                     if i == b - a or tags[i] != tags[run_start]:

tests-unit/comfy_test/test_minimax_h3_forward.py

@@ -0,0 +1,67 @@
+"""MiniMax H3 _forward should pull text-token tags to host once per sampling run, not once per step."""
+
+import torch
+
+from comfy.cli_args import args as cli_args
+
+if not torch.cuda.is_available():
+    cli_args.cpu = True
+
+import comfy.ops as comfy_ops  # noqa: E402
+from comfy.ldm.minimax.model import MiniMaxH3Model  # noqa: E402
+
+
+class _CountingTags:
+    """Stands in for the real tensor payload["text_token_tags"] and counts host pulls."""
+
+    def __init__(self, values):
+        self._values = values
+        self.calls = 0
+
+    def view(self, *shape):
+        self.calls += 1
+        return self
+
+    def tolist(self):
+        return list(self._values)
+
+
+def _make_model():
+    return MiniMaxH3Model(
+        hidden_size=8,
+        num_layers=0,
+        token_refiner_num_layers=1,
+        num_attention_heads=1,
+        attention_head_dim=8,
+        ffn_hidden_size=8,
+        latents_dim=2,
+        audio_latents_dim=2,
+        text_dim=8,  # == hidden_size so _forward skips the token refiner entirely
+        timestep_input_dim=16,
+        time_embed_hidden_size=8,
+        time_embed_dim=8,
+        rope_inv_freq_len=4,
+        operations=comfy_ops.disable_weight_init,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+    )
+
+
+def _forward_inputs(model, text_len=4, latent_t=1, lat_h=2, lat_w=2, audio_t=2):
+    video_x = torch.zeros(1, model.latents_dim, latent_t, lat_h, lat_w)
+    audio_x = torch.zeros(1, model.audio_latents_dim, 2, audio_t)
+    context = torch.zeros(1, text_len, model.hidden_size)
+    timestep = torch.tensor([500.0])
+    return video_x, audio_x, context, timestep
+
+
+def test_text_token_tags_cached_across_forward_calls():
+    model = _make_model()
+    tags = _CountingTags([0, 0, 1, 1])
+    payload = {"text_token_tags": tags}
+    video_x, audio_x, context, timestep = _forward_inputs(model)
+
+    model._forward([video_x, audio_x], timestep, context, minimax_payload=payload)
+    model._forward([video_x, audio_x], timestep, context, minimax_payload=payload)
+
+    assert tags.calls == 1