CthomasdesignandClaude Opus 4.8
1db2179abe
Groundwork for image-grounded conditioning (rung B). Vendors mlx-vlm's pure-MLX Qwen3-VL vision tower into krea2_engine/vision/ (self-contained: inlined SDPA helper, local VisionConfig) and adds torch-free image preprocessing (smart-resize + CLIP normalize + patch flatten). Both validated against the PyTorch reference (tools/test_vision_tower.py): - tower vs transformers Qwen3VLVisionModel: merged embeds + all 3 deepstack features match to max|Δ|~2e-4, cosine 1.0 (bf16 rounding). - preprocessing vs transformers image processor: max|Δ|~2e-7, grid exact. Loads the visual.* weights already present in the Krea-2-Turbo text encoder (315 tensors, 0 missing). Not yet wired into any node — the grounded-encode integration (M-RoPE, deepstack injection, 12-layer tap) is the next step. mlx-vlm is MIT; attribution added to NOTICE. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
461 lines
16 KiB
Python
461 lines
16 KiB
Python
"""Qwen3-VL vision tower — vendored from Blaizzy/mlx-vlm (MIT). See ../../NOTICE for attribution.
|
|
|
|
Unmodified except: `ensure_fused_sdpa` inlined (was mlx_vlm.models.base) and the config import
|
|
pointed at the local self-contained VisionConfig. Loads the `visual.*` weights already present in
|
|
krea/Krea-2-Turbo/text_encoder; returns (merged_image_embeds, [deepstack features]).
|
|
"""
|
|
from itertools import accumulate
|
|
|
|
import mlx.core as mx
|
|
import mlx.nn as nn
|
|
|
|
from .config import VisionConfig
|
|
|
|
|
|
def ensure_fused_sdpa(q, k, v, scale, mask=None):
|
|
fused_dims = (64, 80, 128) # supported by MLX's fused SDPA kernel
|
|
d = q.shape[-1]
|
|
target = next((t for t in fused_dims if d <= t), d)
|
|
if target != d:
|
|
pad = [(0, 0)] * (q.ndim - 1) + [(0, target - d)]
|
|
q, k, v = mx.pad(q, pad), mx.pad(k, pad), mx.pad(v, pad)
|
|
return mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=mask)[..., :d]
|
|
|
|
|
|
def check_array_shape(arr):
|
|
shape = arr.shape
|
|
|
|
# Check if the shape has 4 or 5 dimensions
|
|
if len(shape) not in [4, 5]:
|
|
return False
|
|
|
|
B, out_channels, kH, KW, t = shape
|
|
|
|
if t == 3:
|
|
return True
|
|
|
|
# Check if out_channels is the largest, and kH and KW are the same
|
|
if (out_channels >= kH) and (out_channels >= KW) and (kH == KW):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
|
|
def rotate_half(x):
|
|
"""Rotates half the hidden dims of the input."""
|
|
x1 = x[..., : x.shape[-1] // 2]
|
|
x2 = x[..., x.shape[-1] // 2 :]
|
|
return mx.concatenate([-x2, x1], axis=-1)
|
|
|
|
|
|
def apply_rotary_pos_emb_vision(tensor, freqs) -> mx.array:
|
|
orig_dtype = tensor.dtype
|
|
|
|
cos = mx.cos(freqs)
|
|
sin = mx.sin(freqs)
|
|
|
|
cos = mx.expand_dims(cos, axis=1)
|
|
cos = mx.tile(cos, (1, 1, 2))
|
|
cos = mx.expand_dims(cos, axis=0)
|
|
|
|
sin = mx.expand_dims(sin, axis=1)
|
|
sin = mx.tile(sin, (1, 1, 2))
|
|
sin = mx.expand_dims(sin, axis=0)
|
|
|
|
output = (tensor * cos) + (rotate_half(tensor) * sin)
|
|
return output.astype(orig_dtype)
|
|
|
|
|
|
class VisionRotaryEmbedding(nn.Module):
|
|
def __init__(self, dim: int, theta: float = 10000.0) -> None:
|
|
super().__init__()
|
|
self.dim = dim
|
|
self.theta = theta
|
|
|
|
def __call__(self, seqlen: int) -> mx.array:
|
|
inv_freq = 1.0 / (
|
|
self.theta ** (mx.arange(0, self.dim, 2, dtype=mx.float32) / self.dim)
|
|
)
|
|
seq = mx.arange(seqlen, dtype=inv_freq.dtype)
|
|
freqs = mx.outer(seq, inv_freq)
|
|
return freqs
|
|
|
|
|
|
class PatchEmbed(nn.Module):
|
|
def __init__(
|
|
self,
|
|
patch_size: int = 14,
|
|
temporal_patch_size: int = 2,
|
|
in_channels: int = 3,
|
|
hidden_size: int = 1152,
|
|
) -> None:
|
|
super().__init__()
|
|
self.patch_size = patch_size
|
|
self.temporal_patch_size = temporal_patch_size
|
|
self.in_channels = in_channels
|
|
self.hidden_size = hidden_size
|
|
|
|
kernel_size = [temporal_patch_size, patch_size, patch_size]
|
|
self.proj = nn.Conv3d(
|
|
in_channels,
|
|
hidden_size,
|
|
kernel_size=kernel_size,
|
|
stride=kernel_size,
|
|
bias=True,
|
|
)
|
|
|
|
def __call__(self, hidden_states: mx.array) -> mx.array:
|
|
hidden_states = hidden_states.reshape(
|
|
-1,
|
|
self.in_channels,
|
|
self.temporal_patch_size,
|
|
self.patch_size,
|
|
self.patch_size,
|
|
).moveaxis(1, 4)
|
|
|
|
hidden_states = self.proj(hidden_states)
|
|
hidden_states = hidden_states.reshape(-1, self.hidden_size)
|
|
return hidden_states
|
|
|
|
|
|
class PatchMerger(nn.Module):
|
|
def __init__(self, config: VisionConfig, use_postshuffle_norm=False) -> None:
|
|
super().__init__()
|
|
self.hidden_size = config.hidden_size * (config.spatial_merge_size**2)
|
|
self.use_postshuffle_norm = use_postshuffle_norm
|
|
self.norm = nn.LayerNorm(
|
|
self.hidden_size if use_postshuffle_norm else config.hidden_size, eps=1e-6
|
|
)
|
|
self.linear_fc1 = nn.Linear(self.hidden_size, self.hidden_size)
|
|
self.act_fn = nn.GELU()
|
|
self.linear_fc2 = nn.Linear(self.hidden_size, config.out_hidden_size)
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
x = self.norm(
|
|
x.reshape(-1, self.hidden_size) if self.use_postshuffle_norm else x
|
|
).reshape(-1, self.hidden_size)
|
|
x = self.linear_fc2(self.act_fn(self.linear_fc1(x)))
|
|
return x
|
|
|
|
|
|
class Attention(nn.Module):
|
|
def __init__(self, dim: int, num_heads: int = 16) -> None:
|
|
super().__init__()
|
|
self.num_heads = num_heads
|
|
self.head_dim = dim // num_heads
|
|
self.scale = self.head_dim**-0.5
|
|
self.qkv = nn.Linear(dim, dim * 3, bias=True)
|
|
self.proj = nn.Linear(dim, dim)
|
|
|
|
def __call__(
|
|
self, x: mx.array, cu_seqlens: mx.array, rotary_pos_emb: mx.array = None
|
|
) -> mx.array:
|
|
seq_length = x.shape[0]
|
|
qkv = (
|
|
self.qkv(x).reshape(seq_length, 3, self.num_heads, -1).transpose(1, 0, 2, 3)
|
|
)
|
|
q, k, v = mx.split(qkv, 3)
|
|
|
|
q = apply_rotary_pos_emb_vision(mx.expand_dims(q, 0), rotary_pos_emb)[0]
|
|
k = apply_rotary_pos_emb_vision(mx.expand_dims(k, 0), rotary_pos_emb)[0]
|
|
|
|
q = q.transpose(0, 2, 1, 3)
|
|
k = k.transpose(0, 2, 1, 3)
|
|
v = v.transpose(0, 2, 1, 3)
|
|
|
|
splits = [
|
|
mx.split(tensor, cu_seqlens[1:-1].tolist(), axis=2) for tensor in (q, k, v)
|
|
]
|
|
|
|
attn_outputs = []
|
|
for q, k, v in zip(*splits):
|
|
attn_outputs.append(ensure_fused_sdpa(q, k, v, self.scale))
|
|
|
|
output = mx.concatenate(attn_outputs, axis=2)
|
|
output = output.transpose(0, 2, 1, 3).reshape(seq_length, -1)
|
|
return self.proj(output)
|
|
|
|
|
|
class MLP(nn.Module):
|
|
def __init__(self, dim, hidden_dim):
|
|
super().__init__()
|
|
self.linear_fc1 = nn.Linear(dim, hidden_dim, bias=True)
|
|
self.linear_fc2 = nn.Linear(hidden_dim, dim, bias=True)
|
|
self.act_fn = nn.GELU(approx="tanh")
|
|
|
|
def __call__(self, x: mx.array) -> mx.array:
|
|
return self.linear_fc2(self.act_fn(self.linear_fc1(x)))
|
|
|
|
|
|
class Qwen3VLMoEVisionBlock(nn.Module):
|
|
def __init__(self, config: VisionConfig) -> None:
|
|
super().__init__()
|
|
self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-6)
|
|
self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-6)
|
|
|
|
self.attn = Attention(dim=config.hidden_size, num_heads=config.num_heads)
|
|
self.mlp = MLP(dim=config.hidden_size, hidden_dim=config.intermediate_size)
|
|
|
|
def __call__(self, hidden_states, cu_seqlens, rotary_pos_emb) -> mx.array:
|
|
hidden_states = hidden_states + self.attn(
|
|
self.norm1(hidden_states),
|
|
cu_seqlens=cu_seqlens,
|
|
rotary_pos_emb=rotary_pos_emb,
|
|
)
|
|
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
|
|
return hidden_states
|
|
|
|
|
|
class VisionModel(nn.Module):
|
|
def __init__(self, config: VisionConfig) -> None:
|
|
super().__init__()
|
|
self.config = config
|
|
self.model_type = config.model_type
|
|
|
|
if self.model_type not in [
|
|
"qwen3_vl",
|
|
"qwen3_5",
|
|
"qwen3_5_moe",
|
|
"qwen3_5_vision",
|
|
"qwen3_5_moe_vision",
|
|
]:
|
|
raise ValueError(f"Unsupported model type: {self.model_type}")
|
|
|
|
self.spatial_merge_size = config.spatial_merge_size
|
|
|
|
self.patch_embed = PatchEmbed(
|
|
patch_size=config.patch_size,
|
|
temporal_patch_size=config.temporal_patch_size,
|
|
in_channels=config.in_channels,
|
|
hidden_size=config.hidden_size,
|
|
)
|
|
|
|
head_dim = config.hidden_size // config.num_heads
|
|
self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)
|
|
|
|
self.pos_embed = nn.Embedding(
|
|
config.num_position_embeddings, config.hidden_size
|
|
)
|
|
self.num_grid_per_side = int(config.num_position_embeddings**0.5)
|
|
|
|
self.blocks = [Qwen3VLMoEVisionBlock(config) for _ in range(config.depth)]
|
|
self.merger = PatchMerger(config=config, use_postshuffle_norm=False)
|
|
|
|
self.deepstack_visual_indexes = config.deepstack_visual_indexes
|
|
self.deepstack_merger_list = [
|
|
PatchMerger(
|
|
config=config,
|
|
use_postshuffle_norm=True,
|
|
)
|
|
for _ in range(len(config.deepstack_visual_indexes))
|
|
]
|
|
|
|
def rot_pos_emb(self, grid_thw: mx.array) -> mx.array:
|
|
merge_size = self.spatial_merge_size
|
|
|
|
# Get max grid size for frequency table
|
|
max_hw = int(mx.max(grid_thw[:, 1:]).item())
|
|
freq_table = self.rotary_pos_emb(max_hw) # Shape: (max_hw, dim // 2)
|
|
|
|
pos_ids = []
|
|
|
|
for num_frames, height, width in grid_thw.tolist():
|
|
num_frames, height, width = int(num_frames), int(height), int(width)
|
|
merged_h, merged_w = height // merge_size, width // merge_size
|
|
|
|
# Create block indices
|
|
block_rows = mx.arange(merged_h)
|
|
block_cols = mx.arange(merged_w)
|
|
|
|
# Create intra-block indices
|
|
intra_row = mx.arange(merge_size)
|
|
intra_col = mx.arange(merge_size)
|
|
|
|
# Compute full-resolution positions
|
|
row_idx = (
|
|
block_rows[:, None, None, None] * merge_size
|
|
+ intra_row[None, None, :, None]
|
|
)
|
|
col_idx = (
|
|
block_cols[None, :, None, None] * merge_size
|
|
+ intra_col[None, None, None, :]
|
|
)
|
|
|
|
# Broadcast and flatten
|
|
row_idx = mx.broadcast_to(
|
|
row_idx, (merged_h, merged_w, merge_size, merge_size)
|
|
).reshape(-1)
|
|
col_idx = mx.broadcast_to(
|
|
col_idx, (merged_h, merged_w, merge_size, merge_size)
|
|
).reshape(-1)
|
|
|
|
# Stack into coordinate pairs
|
|
coords = mx.stack([row_idx, col_idx], axis=-1)
|
|
|
|
# Repeat for temporal dimension
|
|
if num_frames > 1:
|
|
coords = mx.tile(coords, (num_frames, 1))
|
|
|
|
pos_ids.append(coords)
|
|
|
|
# Concatenate all position IDs - shape: (total_tokens, 2)
|
|
pos_ids = mx.concatenate(pos_ids, axis=0)
|
|
|
|
# Lookup embeddings: freq_table[h_pos] and freq_table[w_pos]
|
|
# pos_ids[:, 0] = height positions, pos_ids[:, 1] = width positions
|
|
h_embeddings = freq_table[pos_ids[:, 0]] # (total_tokens, dim // 2)
|
|
w_embeddings = freq_table[pos_ids[:, 1]] # (total_tokens, dim // 2)
|
|
|
|
# Concatenate height and width embeddings
|
|
embeddings = mx.concatenate([h_embeddings, w_embeddings], axis=-1)
|
|
|
|
return embeddings
|
|
|
|
def fast_pos_embed_interpolate(self, grid_thw):
|
|
grid_thw_list = grid_thw.tolist()
|
|
idx_list = [[] for _ in range(4)]
|
|
weight_list = [[] for _ in range(4)]
|
|
|
|
for t, h, w in grid_thw_list:
|
|
h = int(h)
|
|
w = int(w)
|
|
t = int(t)
|
|
|
|
h_idxs = mx.linspace(0, self.num_grid_per_side - 1, h)
|
|
w_idxs = mx.linspace(0, self.num_grid_per_side - 1, w)
|
|
|
|
h_idxs_floor = h_idxs.astype(mx.int32)
|
|
w_idxs_floor = w_idxs.astype(mx.int32)
|
|
h_idxs_ceil = mx.minimum(h_idxs_floor + 1, self.num_grid_per_side - 1)
|
|
w_idxs_ceil = mx.minimum(w_idxs_floor + 1, self.num_grid_per_side - 1)
|
|
|
|
dh = h_idxs - h_idxs_floor.astype(mx.float32)
|
|
dw = w_idxs - w_idxs_floor.astype(mx.float32)
|
|
|
|
base_h = h_idxs_floor * self.num_grid_per_side
|
|
base_h_ceil = h_idxs_ceil * self.num_grid_per_side
|
|
|
|
indices = [
|
|
(base_h[:, None] + w_idxs_floor[None, :]).flatten(),
|
|
(base_h[:, None] + w_idxs_ceil[None, :]).flatten(),
|
|
(base_h_ceil[:, None] + w_idxs_floor[None, :]).flatten(),
|
|
(base_h_ceil[:, None] + w_idxs_ceil[None, :]).flatten(),
|
|
]
|
|
|
|
weights = [
|
|
((1 - dh)[:, None] * (1 - dw)[None, :]).flatten(),
|
|
((1 - dh)[:, None] * dw[None, :]).flatten(),
|
|
(dh[:, None] * (1 - dw)[None, :]).flatten(),
|
|
(dh[:, None] * dw[None, :]).flatten(),
|
|
]
|
|
|
|
for i in range(4):
|
|
idx_list[i].extend(indices[i].tolist())
|
|
weight_list[i].extend(weights[i].tolist())
|
|
|
|
idx_tensor = mx.array(idx_list, dtype=mx.int32)
|
|
weight_tensor = mx.array(weight_list, dtype=self.pos_embed.weight.dtype)
|
|
|
|
pos_embeds = self.pos_embed(idx_tensor) * weight_tensor[:, :, None]
|
|
patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
|
|
|
|
split_sizes = [int(h * w) for t, h, w in grid_thw_list]
|
|
if len(split_sizes) > 1:
|
|
split_indices = list(accumulate(split_sizes[:-1]))
|
|
patch_pos_embeds_split = mx.split(patch_pos_embeds, split_indices, axis=0)
|
|
else:
|
|
patch_pos_embeds_split = [patch_pos_embeds]
|
|
|
|
patch_pos_embeds_permute = []
|
|
merge_size = self.config.spatial_merge_size
|
|
|
|
for pos_embed, (t, h, w) in zip(patch_pos_embeds_split, grid_thw_list):
|
|
t, h, w = int(t), int(h), int(w)
|
|
feature_dim = pos_embed.shape[-1]
|
|
pos_embed = mx.tile(pos_embed, (t, 1))
|
|
pos_embed = pos_embed.reshape(t, h, w, feature_dim)
|
|
pos_embed = (
|
|
pos_embed.reshape(
|
|
t,
|
|
h // merge_size,
|
|
merge_size,
|
|
w // merge_size,
|
|
merge_size,
|
|
feature_dim,
|
|
)
|
|
.transpose(0, 1, 3, 2, 4, 5)
|
|
.reshape(-1, feature_dim)
|
|
)
|
|
patch_pos_embeds_permute.append(pos_embed)
|
|
|
|
patch_pos_embeds = mx.concatenate(patch_pos_embeds_permute)
|
|
return patch_pos_embeds
|
|
|
|
def __call__(
|
|
self,
|
|
hidden_states: mx.array,
|
|
grid_thw: mx.array,
|
|
**kwargs,
|
|
) -> mx.array:
|
|
|
|
hidden_states = self.patch_embed(hidden_states)
|
|
pos_embeds = self.fast_pos_embed_interpolate(grid_thw)
|
|
hidden_states = hidden_states + pos_embeds
|
|
rotary_pos_emb = self.rot_pos_emb(grid_thw)
|
|
|
|
seq_len = hidden_states.shape[0]
|
|
hidden_states = hidden_states.reshape(seq_len, -1)
|
|
rotary_pos_emb = rotary_pos_emb.reshape(seq_len, -1)
|
|
|
|
# Assuming grid_thw has shape (batch_size, 3)
|
|
batch_size = grid_thw.shape[0]
|
|
|
|
# Calculate cu_seqlens for each item in the batch
|
|
cu_seqlens = []
|
|
for i in range(batch_size):
|
|
seq_len = grid_thw[i, 1] * grid_thw[i, 2]
|
|
cu_seqlens.append(mx.repeat(seq_len, grid_thw[i, 0]))
|
|
|
|
# Concatenate the cu_seqlens for all items in the batch
|
|
cu_seqlens = mx.concatenate(cu_seqlens)
|
|
|
|
cu_seqlens = mx.cumsum(cu_seqlens.astype(mx.int32), axis=0)
|
|
cu_seqlens = mx.pad(cu_seqlens, (1, 0), mode="constant", constant_values=0)
|
|
|
|
deepstack_feature_lists = []
|
|
for layer_num, blk in enumerate(self.blocks):
|
|
hidden_states = blk(
|
|
hidden_states,
|
|
cu_seqlens=cu_seqlens,
|
|
rotary_pos_emb=rotary_pos_emb,
|
|
)
|
|
if layer_num in self.deepstack_visual_indexes:
|
|
deepstack_feature = self.deepstack_merger_list[
|
|
self.deepstack_visual_indexes.index(layer_num)
|
|
](hidden_states)
|
|
deepstack_feature_lists.append(deepstack_feature)
|
|
|
|
hidden_states = self.merger(hidden_states)
|
|
|
|
return hidden_states, deepstack_feature_lists
|
|
|
|
def sanitize(self, weights):
|
|
sanitized_weights = {}
|
|
for k, v in weights.items():
|
|
if "position_ids" in k:
|
|
# Remove unused position_ids
|
|
continue
|
|
elif "patch_embed.proj.weight" in k:
|
|
# PyTorch conv2d weight tensors have shape:
|
|
# [out_channels, in_channels, kH, KW]
|
|
# MLX conv2d expects the weight be of shape:
|
|
# [out_channels, kH, KW, in_channels]
|
|
if check_array_shape(v):
|
|
sanitized_weights[k] = v
|
|
else:
|
|
sanitized_weights[k] = v.transpose(0, 2, 3, 4, 1)
|
|
else:
|
|
sanitized_weights[k] = v
|
|
|
|
return sanitized_weights
|