| import torch |
| from torch import nn |
| from torch.nn import functional |
| from torch.nn.attention import SDPBackend, sdpa_kernel |
| from torch.utils.checkpoint import checkpoint |
|
|
| CAUSAL_CLOCK_POLICY = 'paired_av_clocks_v1' |
| CAUSAL_BATCH_NUMERICS_POLICY = 'chronological_kv_v1' |
|
|
| def embed_plan_timesteps(model, plan): |
| if plan.bidirectional and plan.clock_groups: |
| parts = [model.time_embedder(model.time_proj(clocks).to(model.time_embedder.linear_1.weight.dtype)) for clocks in plan.clock_groups] |
| clocks = torch.cat(plan.clock_groups) |
| values = torch.cat(parts) |
| owners = torch.stack([torch.where(clocks == time_input)[0][0] for time_input in plan.timesteps]) |
| return values.index_select(0, owners) |
| if plan.bidirectional: |
| return model.time_embedder(model.time_proj(plan.timesteps).to(model.time_embedder.linear_1.weight.dtype)) |
| parts = [model.time_embedder(model.time_proj(clocks).to(model.time_embedder.linear_1.weight.dtype)) for clocks in plan.clock_groups] |
| clocks = torch.cat(plan.clock_groups) |
| order = clocks.argsort() |
| if not torch.equal(clocks[order], plan.timesteps): |
| raise ValueError('Native paired clocks no longer cover the fixed inventory') |
| return torch.cat(parts).index_select(0, order) |
|
|
| def shared_kv_gradient_source(value): |
| if torch.is_grad_enabled() and value.requires_grad and (value.dtype in (torch.bfloat16, torch.float16)): |
| return value.float() |
| return value |
|
|
| def _attention(attn, hidden, rotary, plan): |
| from diffusers.models.transformers.transformer_minimax_h3 import _apply_rotary_emb |
| if attn.fused_projections: |
| raise ValueError('Chunked runtime requires unfused base projections and a shared LoRA owner') |
| query, key, value = [getattr(attn, name)(hidden).unflatten(-1, (attn.heads, -1)) for name in ('to_q', 'to_k', 'to_v')] |
| query, key = (attn.norm_q(query), attn.norm_k(key)) |
| query = torch.cat([_apply_rotary_emb(query[batch_index:batch_index + 1], *rotary[batch_index]) for batch_index in range(query.shape[0])]) |
| key = torch.cat([_apply_rotary_emb(key[batch_index:batch_index + 1], *rotary[batch_index]) for batch_index in range(key.shape[0])]) |
| query, key, value = (values.transpose(1, 2) for values in (query, key, value)) |
| if query.shape[2] > 4096 and query.device.type != 'cuda': |
| raise ValueError('Full-resolution attention requires GPU execution') |
| backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION] |
| if query.shape[2] <= 4096: |
| backends.append(SDPBackend.MATH) |
| outputs = [] |
| for batch_index, groups in enumerate(plan.groups): |
| out = torch.zeros_like(query[batch_index:batch_index + 1]) |
| key_source = shared_kv_gradient_source(key[batch_index:batch_index + 1]) |
| value_source = shared_kv_gradient_source(value[batch_index:batch_index + 1]) |
| for queries, keys in groups: |
| with sdpa_kernel(backends): |
| value = functional.scaled_dot_product_attention(query[batch_index:batch_index + 1].index_select(2, queries), key_source.index_select(2, keys).to(key.dtype), value_source.index_select(2, keys).to(value.dtype), dropout_p=0.0) |
| out = out.index_copy(2, queries, value) |
| outputs.append(out) |
| merged = torch.cat(outputs).transpose(1, 2).flatten(2, 3).type_as(query) |
| return attn.to_out[1](attn.to_out[0](merged)) |
|
|
| def _block_forward(block, hidden, temb, indices, rotary, plan): |
| shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = block.adaln_proj(temb) |
| normalized = block.norm1(hidden) * (1 + scale_a[indices]) + shift_a[indices] |
| hidden = hidden + gate_a[indices] * _attention(block.attn, normalized, rotary, plan) |
| normalized = block.norm2(hidden) * (1 + scale_f[indices]) + shift_f[indices] |
| return hidden + gate_f[indices] * block.ff(normalized) |
|
|
| class ChunkedH3Runtime(nn.Module): |
|
|
| def __init__(self, model, *, checkpoint_blocks=True): |
| super().__init__() |
| self.model = model |
| self.checkpoint_blocks = checkpoint_blocks |
| self.attention_contract = 'explicit_causal_rows_stage_text_native_clocks_kv_order_v4' |
| self.causal_clock_policy = CAUSAL_CLOCK_POLICY |
| self.causal_batch_numerics_policy = CAUSAL_BATCH_NUMERICS_POLICY |
|
|
| def forward(self, video_rows, audio_rows, prompt_embeds, plan, *, return_hidden=False): |
| from diffusers.models.transformers.transformer_minimax_h3 import MiniMaxH3TransformerOutput |
| model = self.model |
| batch, row_count = plan.routes.valid_rows.shape |
| if video_rows.shape[:2] != (batch, plan.video_indices.numel()): |
| raise ValueError('Video input/plan shape differs') |
| if audio_rows.shape[:2] != (batch, plan.audio_indices.numel()): |
| raise ValueError('Audio input/plan shape differs') |
| if prompt_embeds.shape[:2] != (batch, plan.max_text_length): |
| raise ValueError('Text input/plan shape differs') |
| if len({video_rows.device, audio_rows.device, prompt_embeds.device, plan.routes.token_tags.device}) != 1: |
| raise ValueError('Inputs and plan must share a device') |
| video = model.proj_in(video_rows.to(model.proj_in.weight.dtype)) |
| audio = model.audio_proj_in(audio_rows.to(model.audio_proj_in.weight.dtype)) |
| refined = [] |
| for batch_index, length in enumerate(plan.text_lengths): |
| one = model.context_embedder(prompt_embeds[batch_index:batch_index + 1, :length].to(model.context_embedder.weight.dtype)) |
| one = model.token_refiner(one) |
| refined.append(functional.pad(one, (0, 0, 0, plan.max_text_length - length))) |
| text = torch.cat(refined).repeat(1, plan.copies, 1) |
| hidden = text.new_zeros((batch, row_count, text.shape[-1])) |
| hidden = hidden.index_copy(1, plan.text_indices, text) |
| hidden = hidden.index_copy(1, plan.video_indices, video.to(text.dtype)) |
| hidden = hidden.index_copy(1, plan.audio_indices, audio.to(text.dtype)) |
| temb = embed_plan_timesteps(model, plan) |
| indices = plan.timestep_indices * 3 + plan.routes.token_tags[None].clamp(min=0) |
| rotary = tuple((model.rope(pos) for pos in plan.position_ids)) |
| for block in model.transformer_blocks: |
|
|
| def call(hidden_input, time_input, block=block): |
| return _block_forward(block, hidden_input, time_input, indices, rotary, plan) |
| if self.checkpoint_blocks and torch.is_grad_enabled(): |
| hidden = checkpoint(call, hidden, temb, use_reentrant=False) |
| else: |
| hidden = call(hidden, temb) |
| shift, scale = model.norm_out.linear(functional.silu(temb).to(model.norm_out.linear.weight.dtype)).chunk(2, dim=-1) |
| hidden = model.norm_out.norm(hidden) * (1 + scale[plan.timestep_indices]) + shift[plan.timestep_indices] |
| video_hidden = hidden.index_select(1, plan.video_indices).to(model.proj_out.weight.dtype) |
| audio_hidden = hidden.index_select(1, plan.audio_indices).to(model.audio_proj_out.weight.dtype) |
| video = torch.cat([model.proj_out(video_hidden[batch_index:batch_index + 1]) for batch_index in range(batch)]) |
| audio = torch.cat([model.audio_proj_out(audio_hidden[batch_index:batch_index + 1]) for batch_index in range(batch)]) |
| output = MiniMaxH3TransformerOutput(sample=video, audio_sample=audio) |
| return (output, hidden) if return_hidden else output |
|
|
