From 2cbc2bb09715f87d0a8e74373842aeec7fbeb24a Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 16 Aug 2026 08:42:01 -0600 Subject: [PATCH] On minimax h3, trim cached text embeddings to max tokens if they are longer than specified --- .../diffusion_models/minimax_h3/minimax_h3.py | 18 +++++++++++--- .../minimax_h3/src/pipeline.py | 10 ++++++-- .../minimax_h3/src/text_encoder.py | 24 +++++++++++++++++++ 3 files changed, 47 insertions(+), 5 deletions(-) diff --git a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py index fbf3a54..d0cbf42 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/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( diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py b/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py index f1bec54..6e08cd3 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/pipeline.py +++ b/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: diff --git a/extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py b/extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py index 3475e1b..8a7f103 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/src/text_encoder.py +++ b/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