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 3a585cd9..92fa8b71 100644 --- a/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py +++ b/extensions_built_in/diffusion_models/minimax_h3/minimax_h3.py @@ -68,6 +68,7 @@ from .src.packing import ( build_packed_sequence, pack_audio_latents, pad_layouts_to_batch, + unpack_audio_tokens, patchify_video_latents, remap_sigma, unpatchify_video_tokens, @@ -169,6 +170,11 @@ class MinimaxH3Model(BaseModel): self.processor = None # Qwen3VLProcessor self._warned_frame_trim = False self.latent_space_version = "minimax_h3_v1" + # caption token cap (vision blocks are never truncated); the released + # stack has no limit — set 0 to disable + self.max_text_length = int( + self.model_config.model_kwargs.get("max_text_length", 512) + ) @staticmethod def get_train_scheduler(): @@ -558,6 +564,7 @@ class MinimaxH3Model(BaseModel): keyframes=keyframes, device=self.device_torch, dtype=self.torch_dtype, + max_length=self.max_text_length, ) embeds_list.append(embeds) tags_list.append(tags) @@ -639,6 +646,23 @@ class MinimaxH3Model(BaseModel): self.vae.to(self.vae_device_torch) return self.audio_vae.decode(latents.to(self.audio_vae.device, torch.float32)) + @property + def audio_sample_rate(self) -> int: + return packing.AUDIO_SAMPLE_RATE + + def decode_packed_audio_rows(self, rows: torch.Tensor) -> torch.Tensor: + # differentiable: audio perceptual losses backprop through the audio VAE + """Packed channel-major audio rows (B, 2*T, 32) -> stereo waveform + (B, 2, T*800) at 32 kHz. Each stereo channel decodes as its own batch + item through the mono audio VAE.""" + a_lat = rows.shape[1] // packing.AUDIO_CHANNELS + latents = unpack_audio_tokens(rows, a_lat) # (B, 2, 32, T) + b = latents.shape[0] + waveform = self.decode_audio_latents( + latents.reshape(b * packing.AUDIO_CHANNELS, latents.shape[2], a_lat) + ) # (B*2, 1, samples) + return waveform.reshape(b, packing.AUDIO_CHANNELS, -1) + @torch.no_grad() def encode_audio(self, audio_data_list): """[{"waveform": (C, L), "sample_rate": int}, ...] -> packed audio @@ -759,6 +783,11 @@ class MinimaxH3Model(BaseModel): # the stored target follows ai-toolkit's noise - clean batch.audio_target = (audio_noise - raw_audio).detach() audio_rows = (1.0 - sa) * raw_audio + sa * audio_noise + # expose what audio perceptual losses need to rebuild the + # clean estimate (x0 = noisy - sigma_a * pred) and its target + batch.audio_latents = raw_audio + batch.audio_noisy = audio_rows + batch.audio_sigma = sigma_a else: # no soundtrack: silence (zeros) noised at the audio sigma # rides along without contributing to the loss 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 0685042d..58613c45 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 @@ -33,6 +33,7 @@ def encode_minimax_h3_prompt( keyframes: Optional[List] = None, # PIL images already on the target canvas device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, + max_length: Optional[int] = None, # cap on PROMPT tokens (vision blocks are never cut) ): """Encode ONE prompt (with optional keyframes) into MiniMax-H3 conditioning. @@ -72,6 +73,10 @@ def encode_minimax_h3_prompt( token_tags += [TEXT_TAG] * len(label_ids) + [VIDEO_TAG] * len(vision_ids) prompt_ids = tokenizer(prompt, add_special_tokens=False)["input_ids"] + if max_length is not None and max_length > 0: + # the cap applies to the caption only; a keyframe's vision block is + # structural conditioning and cannot be truncated without corrupting it + prompt_ids = prompt_ids[:max_length] token_ids += prompt_ids token_tags += [TEXT_TAG] * len(prompt_ids) if len(token_ids) == 0: