| from __future__ import annotations |
|
|
| import gc |
| import hashlib |
| import importlib.metadata |
| import json |
| import logging |
| import os |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| from transformers import AutoTokenizer |
|
|
| from .ops import attention, rms_norm, rotate, silu |
| from .weights import Linear, WeightStore |
|
|
| log = logging.getLogger(__name__) |
|
|
|
|
| class TextEncoder: |
| """Qwen3-VL's first fifty language layers, with no final RMS normalization. |
| |
| Layers stream one at a time. Conditioning is cached by weights, exact token |
| presentation and precision policy; it is independent of diffusion seed. |
| """ |
|
|
| def __init__( |
| self, path, *, cache_dir="cache", bits=None, group_size=64, query_chunk=0 |
| ): |
| self.path = Path(path) |
| self.config = json.loads((self.path / "config.json").read_text()) |
| self.text_config = self.config["text_config"] |
| self.store = WeightStore(self.path) |
| self.tokenizer = AutoTokenizer.from_pretrained( |
| str(self.path), local_files_only=True |
| ) |
| self.cache_dir = Path(cache_dir) |
| self.bits, self.group_size, self.query_chunk = bits, group_size, query_chunk |
|
|
| def encode(self, prompt, *, images=(), on_layer=None): |
| self.cache_hit = False |
| if not prompt.strip(): |
| raise ValueError("Prompt must be nonempty") |
| if images: |
| from .vision import image_presentation |
|
|
| ids, tags, positions, vision = image_presentation(self, prompt, images) |
| else: |
| raw = self.tokenizer.encode(prompt, add_special_tokens=False)[:4096] |
| ids = raw + [0] * (-len(raw) % 8) |
| tags = mx.ones(len(raw), mx.int32) |
| positions = None |
| vision = None |
| identity = { |
| "version": 5, |
| "weights": self.store.fingerprint(), |
| "config": self.config, |
| "versions": { |
| p: importlib.metadata.version(p) |
| for p in ("mlx", "numpy", "Pillow", "transformers") |
| }, |
| "ids": ids, |
| "bits": self.bits, |
| "group_size": self.group_size, |
| "query_chunk": self.query_chunk, |
| "tap": 50, |
| "tf32": os.environ.get("MLX_ENABLE_TF32"), |
| "images": [ |
| hashlib.sha256(im.tobytes()).hexdigest() + str(im.size) for im in images |
| ], |
| } |
| if images: |
| identity["vision_layer_norm_revision"] = 2 |
| key = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest()[ |
| :24 |
| ] |
| target = self.cache_dir / f"conditioning-{key}.safetensors" |
| if target.exists(): |
| data = mx.load(str(target)) |
| mx.eval(data) |
| self.cache_hit = True |
| log.info( |
| "Conditioning cache hit: %d effective tokens", data["hidden"].shape[0] |
| ) |
| return data["hidden"], data["tags"] |
| prefix = "model.language_model" |
| x = self.store.get(prefix + ".embed_tokens.weight", rows=ids)[None] |
| mx.eval(x) |
| if vision is not None: |
| from .vision import encode_vision |
|
|
| features, deep = encode_vision( |
| self.store, self.config["vision_config"], vision |
| ) |
| image_indices = mx.array( |
| [i for i, t in enumerate(ids) if t == self.config["image_token_id"]] |
| ) |
| x[0, image_indices] = features |
| cfg = self.text_config |
| dim, nq, nk = ( |
| cfg["head_dim"], |
| cfg["num_attention_heads"], |
| cfg["num_key_value_heads"], |
| ) |
| if positions is None: |
| positions = mx.broadcast_to( |
| mx.arange(len(ids), dtype=mx.float32)[None], (3, len(ids)) |
| ) |
| else: |
| positions = mx.array(positions, mx.float32) |
| inv = 1 / (cfg["rope_theta"] ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)) |
| angles = positions[:, :, None] * inv[None, None] |
| |
| sections = cfg["rope_scaling"]["mrope_section"] |
| freq = angles[0] |
| for axis, offset in ((1, 1), (2, 2)): |
| stop = sections[axis] * 3 |
| freq[:, offset:stop:3] = angles[axis, :, offset:stop:3] |
| freq = mx.concatenate((freq, freq), axis=-1) |
| rope = ( |
| mx.cos(freq).astype(mx.bfloat16)[None, :, None], |
| mx.sin(freq).astype(mx.bfloat16)[None, :, None], |
| ) |
| mx.eval(rope) |
| for i in range(50): |
| layer = f"{prefix}.layers.{i}" |
|
|
| def linear(suffix, layer=layer): |
| return Linear.load( |
| self.store, |
| layer + "." + suffix, |
| bits=self.bits, |
| group_size=self.group_size, |
| ) |
|
|
| qproj, kproj, vproj = ( |
| linear("self_attn." + p + "_proj") for p in ("q", "k", "v") |
| ) |
| oproj = linear("self_attn.o_proj") |
| h = rms_norm( |
| x, |
| self.store.get(layer + ".input_layernorm.weight"), |
| cfg["rms_norm_eps"], |
| rounded=True, |
| ) |
| q = qproj(h).reshape(1, len(ids), nq, dim) |
| k = kproj(h).reshape(1, len(ids), nk, dim) |
| v = vproj(h).reshape(1, len(ids), nk, dim) |
| q = rms_norm( |
| q, |
| self.store.get(layer + ".self_attn.q_norm.weight"), |
| cfg["rms_norm_eps"], |
| rounded=True, |
| ) |
| k = rms_norm( |
| k, |
| self.store.get(layer + ".self_attn.k_norm.weight"), |
| cfg["rms_norm_eps"], |
| rounded=True, |
| ) |
| q, k = rotate(q, *rope), rotate(k, *rope) |
| h = attention(q, k, v, causal=True, query_chunk=self.query_chunk).reshape( |
| 1, len(ids), -1 |
| ) |
| x = x + oproj(h) |
| mx.eval(x) |
| del qproj, kproj, vproj, oproj, q, k, v, h |
| gate, up, down = ( |
| linear("mlp." + p + "_proj") for p in ("gate", "up", "down") |
| ) |
| h = rms_norm( |
| x, |
| self.store.get(layer + ".post_attention_layernorm.weight"), |
| cfg["rms_norm_eps"], |
| rounded=True, |
| ) |
| x = x + down(silu(gate(h)) * up(h)) |
| if vision is not None and i < len(deep): |
| x = x.at[0, image_indices].add(deep[i]) |
| mx.eval(x) |
| del gate, up, down, h |
| if on_layer: |
| on_layer(i, x) |
| if (i + 1) % 10 == 0: |
| log.info("Text encoder %d/50", i + 1) |
| mx.clear_cache() |
| hidden = x[0, : len(tags)].astype(mx.bfloat16) |
| mx.eval(hidden) |
| self.cache_dir.mkdir(parents=True, exist_ok=True) |
| tmp = target.with_suffix(".tmp.safetensors") |
| mx.save_safetensors( |
| str(tmp), |
| {"hidden": hidden, "tags": tags}, |
| metadata={"identity": json.dumps(identity)}, |
| ) |
| tmp.replace(target) |
| gc.collect() |
| mx.clear_cache() |
| return hidden, tags |
|
|
