File size: 3,786 Bytes
a066584 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | # Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Mk1 wrappers around the borrowed DiffusionGemma trunk layers.
The trunk implementation is imported as an internal dependency. The public
model type remains ``modilify_mk1``. Two Mk1-specific forwards are installed
on the constructed trunk:
* Router: softmax over all experts, then top-k and renormalize
* Canvas merge: RMS-capped latent residual, not previous-logit soft embeds
"""
from __future__ import annotations
import mlx.core as mx
import mlx.nn as nn
from mlx_vlm.models.diffusion_gemma.language import (
DiffusionGemma4Backbone,
Router,
geglu,
)
LATENT_RESIDUAL_RMS_RATIO_CAP = 0.5
class Mk1Router(Router):
"""Official router weights with Mk1 log-softmax / top-k routing."""
def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]:
x = mx.fast.rms_norm(x, None, self.eps)
x = x * self.scale * self._root_size
scores = self.proj(x)
probabilities = mx.softmax(scores, axis=-1, precise=True)
top_k = self.config.top_k_experts
indices = mx.argpartition(probabilities, kth=-top_k, axis=-1)[..., -top_k:]
weights = mx.take_along_axis(probabilities, indices, axis=-1)
weights = weights / mx.sum(weights, axis=-1, keepdims=True)
weights = weights * self.per_expert_scale[indices]
return indices, weights
def merge_latent_context(
mapper: nn.Module,
token_embeddings: mx.array,
latent_context: mx.array | None,
*,
rms_ratio_cap: float = LATENT_RESIDUAL_RMS_RATIO_CAP,
) -> mx.array:
"""Apply the native self-conditioning bridge with Mk1 RMS capping."""
context = (
mx.zeros_like(token_embeddings)
if latent_context is None
else latent_context.astype(token_embeddings.dtype)
)
if context.shape != token_embeddings.shape:
raise ValueError("Latent context must match the canvas embedding shape.")
normalized = mapper.pre_norm(context)
mapped = mapper.down_proj(
geglu(mapper.gate_proj(normalized), mapper.up_proj(normalized))
)
mapped_rms = mx.sqrt(
mx.mean(mx.square(mapped.astype(mx.float32)), axis=-1, keepdims=True)
)
token_rms = mx.sqrt(
mx.mean(
mx.square(token_embeddings.astype(mx.float32)), axis=-1, keepdims=True
)
)
cap = rms_ratio_cap * token_rms
scale = cap / mx.sqrt(mx.square(mapped_rms) + mx.square(cap) + 1.0e-12)
mapped = mapped * scale.astype(mapped.dtype)
return mapper.post_norm(token_embeddings + mapped)
def _install_mk1_embed_canvas(decoder: nn.Module) -> None:
def _embed_canvas(
canvas_ids,
self_conditioning_logits=None,
self_conditioning_embeddings=None,
):
if self_conditioning_logits is not None:
raise ValueError(
"Modilify Mk1 uses latent embeddings, not logits self-conditioning."
)
token_embeddings = decoder.embed_tokens(canvas_ids) * decoder.embed_scale
return merge_latent_context(
decoder.self_conditioning,
token_embeddings,
self_conditioning_embeddings,
)
decoder._embed_canvas = _embed_canvas
def _install_mk1_routers(backbone: DiffusionGemma4Backbone) -> None:
for layer in backbone.decoder.layers:
replacement = Mk1Router(layer.router.config)
replacement.update(layer.router.parameters())
layer.router = replacement
def build_mk1_backbone(trunk_config) -> DiffusionGemma4Backbone:
"""Construct the borrowed trunk and install Mk1 forwards."""
backbone = DiffusionGemma4Backbone(trunk_config)
_install_mk1_routers(backbone)
_install_mk1_embed_canvas(backbone.decoder)
return backbone
|