Spaces:
Running on Zero
Running on Zero
| """Caches the keyframe/condition-row slice of MiniMax-H3's video patch projection across denoising steps. | |
| `MiniMaxH3Transformer3DModel.forward` projects the whole packed video latent through `proj_in` on every call, | |
| including the leading condition (keyframe) rows -- which the denoise loop never re-steps (see | |
| `MiniMaxH3LoopSchedulerStep`: `latents[num_condition_video_rows:] = scheduler.step(...)`, an in-place slice | |
| assignment that never touches the condition rows at all). Those rows carry the same value, in the same storage, | |
| for the whole denoise loop, so `proj_in`'s output for them is identical on every step -- recomputing it is pure | |
| waste. Because `proj_in` is a plain linear layer, splitting it over the condition/generated row boundary and | |
| concatenating is mathematically identical to running it over the whole tensor at once: this changes nothing about | |
| the output, only when the work happens. | |
| The cache key is the condition slice's own `data_ptr()` + shape + device, which is safe specifically because of the | |
| in-place assignment above: the condition rows live in the same storage, unchanged, for an entire `pipe(...)` call, | |
| and a *new* call (a new stage, a new chunk, an unrelated generation) always allocates a fresh latents tensor -- the | |
| cache invalidates itself automatically, with no explicit per-request reset needed. | |
| `patch()` wraps `forward` once, at load time, preserving the original's `@apply_lora_scale` behavior inline (a | |
| real decorator -- it scales and unscales active LoRA layers around the forward pass, not just bookkeeping -- so | |
| bypassing it would silently break LoRA strength). The wrapped method checks `transformer._proj_cache_enabled` | |
| fresh on every call, so `set_enabled()` can toggle it per request without needing to unpatch anything. | |
| """ | |
| from __future__ import annotations | |
| import torch | |
| _PATCHED: set[int] = set() | |
| def _condition_row_count(video_indices: torch.Tensor) -> int: | |
| """Number of leading, contiguous condition (keyframe) rows in `video_indices`' own order.""" | |
| host = video_indices.detach().cpu() | |
| if host.numel() < 2: | |
| return 0 | |
| discontinuities = (host[1:] - host[:-1] != 1).nonzero().flatten() | |
| return int(discontinuities[0]) + 1 if len(discontinuities) else 0 | |
| def patch(transformer) -> None: | |
| """Wrap `transformer.forward` once. Safe to call at startup; a no-op on a second call for the same instance.""" | |
| if id(transformer) in _PATCHED: | |
| return | |
| from diffusers.models.modeling_utils import get_parameter_dtype | |
| from diffusers.models.transformers.transformer_minimax_h3 import ( | |
| MINIMAX_H3_MODALITY_NUM, | |
| MiniMaxH3TransformerOutput, | |
| ) | |
| from diffusers.utils import USE_PEFT_BACKEND | |
| from diffusers.utils.peft_utils import scale_lora_layers, unscale_lora_layers | |
| original_forward = transformer.forward | |
| def forward( | |
| hidden_states, | |
| audio_hidden_states, | |
| encoder_hidden_states, | |
| timestep, | |
| timestep_indices, | |
| token_tags, | |
| position_ids, | |
| video_indices, | |
| audio_indices, | |
| text_indices, | |
| attention_kwargs=None, | |
| return_dict=True, | |
| ): | |
| self = transformer | |
| if not getattr(self, "_proj_cache_enabled", True): | |
| return original_forward( | |
| hidden_states, audio_hidden_states, encoder_hidden_states, timestep, timestep_indices, token_tags, | |
| position_ids, video_indices, audio_indices, text_indices, attention_kwargs, return_dict, | |
| ) | |
| # Verbatim from the stock `@apply_lora_scale("attention_kwargs")` decorator this replaces. | |
| lora_scale = 1.0 | |
| kwargs = attention_kwargs | |
| if kwargs is not None: | |
| kwargs = kwargs.copy() | |
| lora_scale = kwargs.pop("scale", 1.0) | |
| if USE_PEFT_BACKEND: | |
| scale_lora_layers(self, lora_scale) | |
| try: | |
| if position_ids.ndim != 2 or position_ids.shape[-1] != 3: | |
| raise ValueError(f"`position_ids` must be a `(seq_len, 3)` tensor, got {list(position_ids.shape)}.") | |
| sequence_length = position_ids.shape[0] | |
| if token_tags.shape != (sequence_length,) or timestep_indices.shape != (sequence_length,): | |
| raise ValueError( | |
| "`token_tags` and `timestep_indices` must both be `(seq_len,)` tensors matching `position_ids`, " | |
| f"got {list(token_tags.shape)} and {list(timestep_indices.shape)} for seq_len={sequence_length}." | |
| ) | |
| rotary_emb = self.rope(position_ids) | |
| # The only change from the stock forward: `proj_in` is split at the condition/generated row boundary | |
| # and the condition half is cached. `torch.cat([proj_in(a), proj_in(b)]) == proj_in(torch.cat([a, b]))` | |
| # for a plain linear layer, so this is bit-identical to the stock computation, never an approximation. | |
| condition_rows = _condition_row_count(video_indices) | |
| if condition_rows > 0: | |
| dtype = get_parameter_dtype(self.proj_in) | |
| condition_slice = hidden_states[:, :condition_rows] | |
| key = (condition_slice.data_ptr(), tuple(condition_slice.shape), condition_slice.device) | |
| cached = getattr(self, "_condition_embeds_cache", None) | |
| if cached is not None and cached[0] == key: | |
| condition_embeds = cached[1] | |
| else: | |
| condition_embeds = self.proj_in(condition_slice.to(dtype)) | |
| self._condition_embeds_cache = (key, condition_embeds) | |
| generated_slice = hidden_states[:, condition_rows:] | |
| video_embeds = ( | |
| torch.cat((condition_embeds, self.proj_in(generated_slice.to(dtype))), dim=1) | |
| if generated_slice.shape[1] > 0 | |
| else condition_embeds | |
| ) | |
| else: | |
| video_embeds = self.proj_in(hidden_states.to(get_parameter_dtype(self.proj_in))) | |
| audio_embeds = self.audio_proj_in(audio_hidden_states.to(get_parameter_dtype(self.audio_proj_in))) | |
| text_embeds = self.context_embedder(encoder_hidden_states.to(get_parameter_dtype(self.context_embedder))) | |
| text_embeds = self.token_refiner(text_embeds) | |
| hidden_states = text_embeds.new_zeros((text_embeds.shape[0], sequence_length, text_embeds.shape[-1])) | |
| hidden_states = hidden_states.index_copy(1, text_indices, text_embeds) | |
| hidden_states = hidden_states.index_copy(1, video_indices, video_embeds.to(text_embeds.dtype)) | |
| hidden_states = hidden_states.index_copy(1, audio_indices, audio_embeds.to(text_embeds.dtype)) | |
| temb = self.time_proj(timestep) | |
| temb = self.time_embedder(temb.to(get_parameter_dtype(self.time_embedder))) | |
| adaln_indices = timestep_indices * MINIMAX_H3_MODALITY_NUM + token_tags | |
| for block in self.transformer_blocks: | |
| if torch.is_grad_enabled() and self.gradient_checkpointing: | |
| hidden_states = self._gradient_checkpointing_func( | |
| block, hidden_states, temb, adaln_indices, rotary_emb | |
| ) | |
| else: | |
| hidden_states = block(hidden_states, temb, adaln_indices, rotary_emb) | |
| hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(get_parameter_dtype(self.proj_out)) | |
| video_output = self.proj_out(hidden_states).index_select(1, video_indices) | |
| audio_output = self.audio_proj_out(hidden_states).index_select(1, audio_indices) | |
| if not return_dict: | |
| return (video_output, audio_output) | |
| return MiniMaxH3TransformerOutput(sample=video_output, audio_sample=audio_output) | |
| finally: | |
| if USE_PEFT_BACKEND: | |
| unscale_lora_layers(self, lora_scale) | |
| transformer.forward = forward | |
| transformer._proj_cache_enabled = True | |
| _PATCHED.add(id(transformer)) | |
| def set_enabled(transformer, enabled: bool) -> None: | |
| """Toggle for one request. Call right before `pipe(...)`; `patch()` must have run first (a no-op otherwise, | |
| since an unpatched transformer just runs its own stock forward regardless of this attribute).""" | |
| transformer._proj_cache_enabled = bool(enabled) | |
| def status() -> str: | |
| return "keyframe-projection cache patched" if _PATCHED else "keyframe-projection cache not patched" | |