On minimax h3, trim cached text embeddings to max tokens if they are longer than specified

extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py

@@ -83,6 +83,7 @@ from .src.text_encoder import (
     VideoRef,
     encode_minimax_h3_prompt,
     load_video_ref,
+    trim_caption_tokens,
 )
 from .src.transformer import MiniMaxH3Transformer, MiniMaxH3TransformerParams
 from .src.vae import MiniMaxH3VideoVAE
@@ -1003,12 +1004,23 @@ class MinimaxH3Model(BaseModel):
                     dtype=torch.float32,
                 )
 
+            # embeds cached with a longer max_text_length: cap the caption
+            # tail (vision blocks are never touched)
+            trimmed = [
+                trim_caption_tokens(e, t, self.max_text_length)
+                for e, t in zip(
+                    text_embeddings.text_embeds, text_embeddings.text_token_tags
+                )
+            ]
+            text_embed_list = [e for e, _ in trimmed]
+            text_tag_list = [t for _, t in trimmed]
+
             # --- packed layout (per item: text lengths differ) --------------
             layouts = []
             for i in range(batch_size):
                 layouts.append(
                     build_packed_sequence(
-                        text_token_tags=text_embeddings.text_token_tags[i].to("cpu"),
+                        text_token_tags=text_tag_list[i].to("cpu"),
                         num_latent_frames=t_lat,
                         latent_height=h_lat,
                         latent_width=w_lat,
@@ -1046,11 +1058,11 @@ class MinimaxH3Model(BaseModel):
             text_batch = torch.zeros(
                 batch_size,
                 max_text,
-                text_embeddings.text_embeds[0].shape[-1],
+                text_embed_list[0].shape[-1],
                 device=device,
                 dtype=dtype,
             )
-            for i, emb in enumerate(text_embeddings.text_embeds):
+            for i, emb in enumerate(text_embed_list):
                 text_batch[i, : emb.shape[0]] = emb.to(device, dtype)
 
             video_rows = patchify_video_latents(

extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py

@@ -26,6 +26,7 @@ from PIL import Image
 from diffusers.utils.torch_utils import randn_tensor
 
 from . import packing
+from .text_encoder import trim_caption_tokens
 from .packing import (
     AUDIO_CHANNELS,
     AUDIO_SIGMA_SHIFT,
@@ -100,8 +101,13 @@ class MiniMaxH3Pipeline:
         w_lat = width // 16
         a_lat = packing.audio_latent_num_frames(num_frames)
 
-        text_embeds = conditional_embeds.text_embeds[0].to(device, dtype)
-        token_tags = conditional_embeds.text_token_tags[0].to("cpu", torch.long)
+        text_embeds, token_tags = trim_caption_tokens(
+            conditional_embeds.text_embeds[0],
+            conditional_embeds.text_token_tags[0],
+            getattr(model, "max_text_length", None),
+        )
+        text_embeds = text_embeds.to(device, dtype)
+        token_tags = token_tags.to("cpu", torch.long)
 
         # --- packed layout -------------------------------------------------
         if ctrl_img is not None and ref_images:

extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py

@@ -89,6 +89,30 @@ def load_video_ref(path, max_frames: int = 0, has_audio=None) -> "VideoRef":
     return VideoRef(frames=frames, timestamps=times, has_audio=has_audio)
 
 
+def trim_caption_tokens(embeds: torch.Tensor, tags: torch.Tensor, max_length):
+    """Cap the CAPTION rows of an already-encoded prompt at ``max_length``.
+
+    The presentation is ``[labels + vision blocks ...] + caption``: everything
+    up to and including the last vision row (tag 0) is structural conditioning
+    and kept intact; only the trailing text-tagged caption tail is truncated.
+    Lets embeds cached with a longer max_text_length serve a shorter one
+    without re-encoding. Returns (embeds, tags) — same objects when nothing to
+    trim.
+    """
+    if max_length is None or max_length <= 0:
+        return embeds, tags
+    is_vision = tags.to("cpu") == VIDEO_TAG
+    if bool(is_vision.any()):
+        head = int(is_vision.nonzero()[-1].item()) + 1
+    else:
+        head = 0
+    caption_len = int(tags.shape[0]) - head
+    if caption_len <= max_length:
+        return embeds, tags
+    keep = head + max_length
+    return embeds[:keep], tags[:keep]
+
+
 @torch.no_grad()
 def encode_minimax_h3_prompt(
     text_encoder,  # transformers Qwen3VLForConditionalGeneration