streamembed 0.1.0__py3-none-any.whl

streamembed/engine/llm.py

@@ -0,0 +1,170 @@
+"""LLM part: token sequence, exact prefix reuse (with KV Reuse (ViT)) and KV Reuse (LLM).
+
+Sequence = [prompt head (first image token at `first`)] + [one visual token per kept block] + [text tail incl. the final token].
+The embedding is the normalised final hidden state of the last token.
+
+Exact prefix reuse (on with KV Reuse): the LLM is causal, so positions before the first changed position P are bit-identical to the
+previous frame. Crop the previous frame's KV cache to P and prefill only [P, L). If nothing changed, return the previous embedding.
+
+KV Reuse (LLM) (approximate): recompute only R = visual tokens of changed blocks + every text token after the first change + the final token
+(+ any slot without cached K/V). Unchanged visual tokens after the first change reuse their cached per-layer K/V although their causal
+context changed: this is the approximation. Text after the change is always recomputed because it summarises the image.
+"""
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from transformers.cache_utils import DynamicCache
+from transformers.models.qwen3_vl import modeling_qwen3_vl as MQ
+
+from .model import Grid, Model
+from .temporal import FramePlan, to_device
+
+
+def full_grid_slots(g: Grid, kept: np.ndarray) -> np.ndarray:
+    """Full-grid slot of each position of the (pruned) sequence."""
+    return np.concatenate([np.arange(g.first), g.first + kept, np.arange(g.last + 1, g.length)])
+
+
+def first_changed_position(g: Grid, plan: FramePlan) -> int:
+    """P: length of the prefix identical to the previous frame's sequence (0 without history)."""
+    st = plan.state
+    if st is None or st.prev_kept is None:
+        return 0
+    kept, prev = plan.kept, st.prev_kept
+    recomputed = np.zeros(g.n_blocks, bool)
+    recomputed[plan.recompute] = True
+    n = min(len(kept), len(prev))
+    same = (kept[:n] == prev[:n]) & ~recomputed[kept[:n]]   # same block in the same slot and not recomputed
+    j = int(np.argmin(same)) if not same.all() else n
+    if j == n and len(kept) == len(prev):
+        return g.llm_len(len(kept))                         # nothing changed
+    return g.first + j
+
+
+def plan_llm(g: Grid, plan: FramePlan, approx_llm: bool):
+    """Fill plan.prefix, plan.skip_llm and (KV Reuse (LLM)) plan.llm_rows, on the host."""
+    st = plan.state
+    L = g.llm_len(len(plan.kept))
+    plan.prefix = first_changed_position(g, plan)
+    plan.skip_llm = st is not None and plan.prefix >= L and st.prev_emb is not None
+    if approx_llm and not plan.skip_llm:
+        plan.llm_rows = recompute_rows(g, plan)
+
+
+def recompute_rows(g: Grid, plan: FramePlan) -> np.ndarray:
+    """KV Reuse (LLM): R, the sorted positions to run through the LLM (always contains the final position L-1)."""
+    f, kept, P = g.first, plan.kept, plan.prefix
+    nk = len(kept)
+    L = g.llm_len(nk)
+    r = np.zeros(L, bool)
+    if plan.refresh or P == 0:
+        r[P:] = True
+    else:
+        changed = np.zeros(g.n_blocks, bool)
+        changed[plan.recompute] = True
+        r[f:f + nk] = changed[kept]                         # visual tokens of recomputed blocks
+        r[max(P, f + nk):] = True                           # text after the first change
+    r |= ~plan.state.llm_valid[full_grid_slots(g, kept)]   # slots never computed in this stream
+    r[L - 1] = True
+    return np.flatnonzero(r)
+
+
+def llm_inputs(g: Grid, kept: np.ndarray, kept_t: torch.Tensor, visual_tokens: torch.Tensor, m: Model):
+    """Input embeddings and mRoPE positions of the (pruned) sequence.
+    changed: dropped blocks are removed from the sequence; the rest keep their full-grid mRoPE positions."""
+    seq = torch.cat([torch.arange(g.first, device=m.device), g.first + kept_t, torch.arange(g.last + 1, g.length, device=m.device)])
+    x = g.token_emb[seq].clone()
+    x[g.first:g.first + len(kept)] = visual_tokens.to(x.dtype)
+    return x, g.llm_pos[:, seq]
+
+
+@torch.inference_mode()
+def run_exact(m: Model, g: Grid, plan: FramePlan, visual_tokens: torch.Tensor, deepstack: list) -> torch.Tensor:
+    """Stock HF LLM forward; with KV Reuse (ViT) it continues from the previous frame's KV cache cropped to P (exact)."""
+    st = plan.state
+    if plan.skip_llm:
+        return st.prev_emb
+    x, pos = llm_inputs(g, plan.kept, plan.kept_t, visual_tokens, m)
+    L, f, nk = x.shape[0], g.first, len(plan.kept)
+    P = plan.prefix
+    if st is None or st.exact_kv is None or P == 0:
+        kv, P = DynamicCache(config=m.tcfg), 0
+    else:
+        kv = st.exact_kv                                    # changed: reuse the previous frame's cache, cropped to P
+        cur = kv.get_seq_length()
+        if cur > P:
+            kv.crop(-(cur - P))
+    visual = torch.zeros(L, dtype=torch.bool, device=m.device)
+    visual[f:f + nk] = True
+    n_vis_in_prefix = min(max(P - f, 0), nk)
+    hs = m.llm(input_ids=None, inputs_embeds=x[None, P:], position_ids=pos[:, None, P:], past_key_values=kv, use_cache=True,
+               visual_pos_masks=visual[None, P:],
+               deepstack_visual_embeds=[d[n_vis_in_prefix:].to(x.dtype) for d in deepstack]).last_hidden_state
+    if st is not None:
+        st.exact_kv = kv
+    return F.normalize(hs[0, -1].float(), p=2, dim=-1)
+
+
+@torch.inference_mode()
+def run_approx(m: Model, g: Grid, plan: FramePlan, visual_tokens: torch.Tensor, deepstack: list) -> torch.Tensor:
+    """Run only rows R through every LLM layer. Per layer, write their fresh K/V into the full-grid slot cache and attend causally
+    (by sequence index) over the cached K/V of all current positions.
+    changed: re-implements the Qwen3 decoder layer (same modules) with a row subset and a slot-indexed K/V cache."""
+    st = plan.state
+    if plan.skip_llm:
+        return st.prev_emb
+    tc = m.tcfg
+    nh, nkv, hd = tc.num_attention_heads, tc.num_key_value_heads, tc.head_dim
+    f, nk = g.first, len(plan.kept)
+    x, pos = llm_inputs(g, plan.kept, plan.kept_t, visual_tokens, m)
+    rows_np = plan.llm_rows
+    slots_np = full_grid_slots(g, plan.kept)
+    is_vis = (rows_np >= f) & (rows_np < f + nk)
+    rows, slots = to_device(rows_np, m), to_device(slots_np, m)
+    row_slots = to_device(slots_np[rows_np], m)
+    vis_rows = to_device(np.flatnonzero(is_vis), m)                         # rows (within R) that are visual tokens
+    ds_rows = to_device(rows_np[is_vis] - f, m)                             # their index among the kept blocks
+    h = x[rows]
+    rpos = pos[:, rows][:, None]
+    if rpos.shape[0] == 4:                                                  # (text, t, h, w) layout -> the 3 mRoPE rows
+        rpos = rpos[1:]
+    ds = [d[ds_rows].to(h.dtype) for d in deepstack]
+    cos, sin = m.llm.rotary_emb(h, rpos)
+    cos, sin = cos[0], sin[0]
+    mask = (torch.arange(len(slots_np), device=m.device)[None, :] <= rows[:, None])[None, None]   # causal over sequence index
+    for layer_id, layer in enumerate(m.llm.layers):
+        at = layer.self_attn
+        q, k, v = KERNELS["llm_qkv"](layer, h, cos, sin, nh, nkv, hd)
+        st.llm_k[layer_id][row_slots] = k
+        st.llm_v[layer_id][row_slots] = v
+        kf, vf = st.llm_k[layer_id][slots], st.llm_v[layer_id][slots]
+        o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], kf.transpose(0, 1)[None], vf.transpose(0, 1)[None],
+                                           attn_mask=mask, dropout_p=0.0, scale=at.scaling, enable_gqa=True)
+        h = KERNELS["llm_out"](layer, h, o, len(rows_np))
+        if layer_id < len(ds) and len(vis_rows):                             # DeepStack: add ViT features after early layers
+            h[vis_rows] = h[vis_rows] + ds[layer_id]
+    st.llm_valid[slots_np[rows_np]] = True
+    return F.normalize(m.llm.norm(h[-1]).float(), p=2, dim=-1)              # R is sorted and ends with the final position
+
+
+def _llm_qkv(layer, h, cos, sin, nh: int, nkv: int, hd: int):
+    """input RMSNorm -> q/k/v projections -> q/k RMSNorm -> mRoPE (the dense part before attention)."""
+    at = layer.self_attn
+    a = layer.input_layernorm(h)
+    q = at.q_norm(at.q_proj(a).view(-1, nh, hd))
+    k = at.k_norm(at.k_proj(a).view(-1, nkv, hd))
+    v = at.v_proj(a).view(-1, nkv, hd)
+    q, k = MQ.apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1)
+    return q, k, v
+
+
+def _llm_out(layer, h, o, n: int):
+    """o_proj + residual, then post-attention RMSNorm -> MLP + residual (the dense part after attention)."""
+    h = h + layer.self_attn.o_proj(o[0].transpose(0, 1).reshape(n, -1))
+    return h + layer.mlp(layer.post_attention_layernorm(h))
+
+
+# changed: the dense per-layer parts of KV Reuse (LLM) are looked up here so cpu_compile.enable() can swap in torch.compile'd versions
+KERNELS = {"llm_qkv": _llm_qkv, "llm_out": _llm_out}

streamembed/engine/macs.py

@@ -0,0 +1,46 @@
+"""Multiply-accumulate (MAC) count of the work actually done per frame (weights and attention; norms / activations ignored).
+
+All results are Python ints (no overflow). MAC fraction = (vit_macs + llm_macs) / full_frame_macs(m, grid).
+"""
+from __future__ import annotations
+
+import numpy as np
+
+from .model import Grid, Model
+
+
+def vit_macs(m: Model, n_query: int, n_kv: int, n_merge: int) -> int:
+    """ViT with n_query query patches attending over n_kv patches in every layer; n_merge blocks through merger + DeepStack mergers."""
+    cv = m.vcfg
+    d, inter = cv.hidden_size, cv.intermediate_size
+    mh = d * cv.spatial_merge_size ** 2
+    per_layer = n_query * (4 * d * d + 2 * d * inter) + 2 * n_query * n_kv * d
+    patch_embed = n_query * cv.in_channels * cv.temporal_patch_size * cv.patch_size ** 2 * d
+    mergers = n_merge * (len(m.deepstack_layers) + 1) * (mh * mh + mh * cv.out_hidden_size)
+    return len(m.vit.blocks) * per_layer + patch_embed + mergers
+
+
+def _llm_linear_per_token(m: Model) -> int:
+    tc = m.tcfg
+    H, nh, nk, hd, inter = tc.hidden_size, tc.num_attention_heads, tc.num_key_value_heads, tc.head_dim, tc.intermediate_size
+    return 2 * H * nh * hd + 2 * H * nk * hd + 3 * H * inter       # q,o + k,v + gate,up,down
+
+
+def llm_macs_suffix(m: Model, P: int, L: int) -> int:
+    """Exact path: prefill positions [P, L); position i attends over i+1 positions."""
+    tc = m.tcfg
+    att = 2 * tc.num_attention_heads * tc.head_dim * (L * (L + 1) - P * (P + 1)) // 2
+    return tc.num_hidden_layers * ((L - P) * _llm_linear_per_token(m) + att)
+
+
+def llm_macs_rows(m: Model, rows: np.ndarray) -> int:
+    """KV Reuse (LLM) path: only positions `rows` (the host indices that were executed); row i attends over i+1 positions."""
+    tc = m.tcfg
+    rows = np.asarray(rows, dtype=np.int64)
+    att = 2 * tc.num_attention_heads * tc.head_dim * int((rows + 1).sum())
+    return tc.num_hidden_layers * (len(rows) * _llm_linear_per_token(m) + att)
+
+
+def full_frame_macs(m: Model, g: Grid) -> int:
+    """FULL config on this resolution (3.084e12 for 1280x704 / 880 blocks)."""
+    return vit_macs(m, g.n_patches, g.n_patches, g.n_blocks) + llm_macs_suffix(m, 0, g.length)

streamembed/engine/model.py

@@ -0,0 +1,254 @@
+"""Model loading (+ LoRA, optional bf16 / bf16-llm-fp32 weight storage), the linear patch embedding and per-resolution constants.
+
+Qwen3-VL-Embedding-2B = ViT (16x16-px patches) + LLM; the embedding is the normalised final hidden state of the last token.
+A 1280x704 screenshot gives 80x44 patches; the ViT merger fuses each 2x2 group of patches (= one 32x32-px "block") into one LLM token,
+so 880 blocks / visual tokens. The processor already orders patches block by block (4 consecutive rows = one block).
+"""
+from __future__ import annotations
+
+import warnings
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from transformers.models.qwen3_vl import modeling_qwen3_vl as MQ
+
+from .config import MODEL_ID
+from .qwen_embedder import Qwen3VLEmbedder
+
+
+@dataclass
+class Model:
+    embedder: Qwen3VLEmbedder           # processor + HF model; also the text embedder of the OCR route
+    device: torch.device
+    dtype: torch.dtype
+    grids: dict = field(default_factory=dict)   # (grid_thw, prompt length) -> Grid
+
+    def __post_init__(self):
+        base = self.embedder.model
+        self.core = base.model                      # HF Qwen3VLModel
+        self.vit = self.core.visual
+        self.llm = self.core.language_model
+        self.vcfg, self.tcfg = self.vit.config, self.llm.config
+        self.image_token_id = base.config.image_token_id
+        self.vit_heads = self.vcfg.num_heads
+        self.vit_head_dim = self.vcfg.hidden_size // self.vcfg.num_heads
+        self.deepstack_layers = list(self.vit.deepstack_visual_indexes)   # ViT layers whose output is added to early LLM layers
+        pe = self.vit.patch_embed
+        # changed: Conv3d (kernel == stride == patch) flattened once into a matrix for F.linear (see embed_patches)
+        self.patch_w = pe.proj.weight.detach().reshape(pe.embed_dim, -1)
+        self.patch_b = pe.proj.bias.detach()
+
+    def sync(self):
+        """Wait for the device so wall-clock timings cover the device work."""
+        if self.device.type == "mps":
+            torch.mps.synchronize()
+        elif self.device.type == "cuda":
+            torch.cuda.synchronize(self.device)
+
+
+def pick_device(name: str = "auto") -> str:
+    if name != "auto":
+        return name
+    if torch.cuda.is_available():
+        return "cuda"
+    return "mps" if torch.backends.mps.is_available() else "cpu"
+
+
+def set_mps_memory_cap(cap_gb: float) -> float:
+    """Cap this process's MPS allocations (useful on a shared 16 GB Mac). Returns the fraction set."""
+    frac = min(1.0, cap_gb * 2 ** 30 / torch.mps.recommended_max_memory())
+    torch.mps.set_per_process_memory_fraction(frac)
+    return frac
+
+
+WEIGHT_STORAGE = ("native", "bf16", "bf16-llm-fp32")
+
+
+def load_model(device: str = "auto", dtype: torch.dtype = torch.float16, lora: Optional[str] = None, weight_storage: str = "native",
+               model_id: str = MODEL_ID, lora_merge: bool = True) -> Model:
+    """Load Qwen3-VL-Embedding-2B on the CPU, add an optional LoRA, then move it to `device`.
+
+    lora: PEFT LoRA state dict (.pt); None = zero-shot.
+    lora_merge=True: W += B @ A at load, so the compute graph equals the base model (paper default, fp16 / fp32).
+    lora_merge=False: keep the low-rank path separate, x W^T + (x A^T) B^T, as PEFT computes it (fp32 adapter path). Use this in bf16:
+    rounding W + B @ A to bf16 (8-bit mantissa) loses much of the low-rank update.
+    weight_storage="bf16" (CPU + float32 only): keep Linear weights in the checkpoint's bf16, compute in fp32 (store_linear_weights_bf16).
+    weight_storage="bf16-llm-fp32": as "bf16", then the LLM decoder-layer weights are upcast to fp32 once (ViT and the input-embedding
+    table stay bf16): no per-call upcast in the LLM, bitwise identical outputs, ~7 GB instead of ~4.3 GB.
+    """
+    assert weight_storage in WEIGHT_STORAGE, weight_storage
+    device = pick_device(device)
+    bf16_store = weight_storage in ("bf16", "bf16-llm-fp32")
+    if bf16_store:
+        assert device == "cpu" and dtype == torch.float32, "bf16 weight storage is the CPU fp32 memory saver"
+    embedder = Qwen3VLEmbedder(model_id, dtype=torch.bfloat16 if bf16_store else dtype)
+    if lora and lora_merge and embedder.model.dtype == torch.bfloat16:
+        warnings.warn("merging the LoRA into bf16 weights rounds away part of the update; consider lora_merge=False (--no-lora-merge)")
+    if lora:
+        (merge_lora if lora_merge else attach_lora)(embedder.model, lora)
+    embedder.model.to(device).eval()
+    if bf16_store:
+        store_linear_weights_bf16(embedder.model)
+    if weight_storage == "bf16-llm-fp32":                  # changed: LLM weights upcast once at load (no per-call bf16 -> fp32 copy)
+        for p in embedder.model.model.language_model.layers.parameters():
+            p.data = p.data.float()
+    return Model(embedder=embedder, device=torch.device(device), dtype=dtype)
+
+
+def merge_lora(model: torch.nn.Module, path: str, scale: float = 1.0) -> int:
+    """Merge a PEFT LoRA state dict (alpha / r = 1 for the paper's r 32 / alpha 32) into the base weights: W += scale * B @ A in fp32.
+
+    Keys look like base_model.model.<module path>.lora_A.weight; the module path is relative to Qwen3VLForEmbedding.
+    """
+    pairs = _lora_pairs(path)
+    for module, a, b in pairs:
+        w = model.get_submodule(module).weight
+        w.data = (w.data.float() + scale * (b.float() @ a.float())).to(w.dtype)
+    return len(pairs)
+
+
+def _lora_pairs(path: str):
+    """(module path, A, B) per LoRA module of a PEFT state dict (keys base_model.model.<module path>.lora_A.weight), .safetensors or .pt."""
+    if str(path).endswith(".safetensors"):
+        from safetensors.torch import load_file
+        state = load_file(str(path), device="cpu")
+    else:
+        state = torch.load(path, map_location="cpu", weights_only=True)
+    pairs = [(key[len("base_model.model."):-len(".lora_A.weight")], a, state[key.replace(".lora_A.", ".lora_B.")])
+             for key, a in state.items() if key.endswith(".lora_A.weight")]
+    assert pairs and len(pairs) == len(state) // 2, f"unexpected LoRA state dict {path}"
+    return pairs
+
+
+class LoRALinear(torch.nn.Module):
+    """nn.Linear + separate low-rank update, computed as PEFT does: A, B stay fp32 (PEFT upcasts bf16 / fp16 adapters), the input is
+    cast to fp32 for the low-rank path, base(x) + B(A(x)) * scale is summed in fp32 and cast back to the base output dtype.
+    dropout (training only, PEFT's lora_dropout) acts on the fp32 input of the low-rank path; it is the identity in eval mode."""
+
+    def __init__(self, base: torch.nn.Linear, a: torch.Tensor, b: torch.Tensor, scale: float = 1.0, dropout: float = 0.0):
+        super().__init__()
+        self.base, self.scale = base, scale
+        self.lora_a = torch.nn.Parameter(a.float(), requires_grad=False)
+        self.lora_b = torch.nn.Parameter(b.float(), requires_grad=False)
+        self.dropout = torch.nn.Dropout(dropout) if dropout else torch.nn.Identity()
+
+    @property
+    def weight(self):                   # expose the base weight (dtype / shape lookups by callers)
+        return self.base.weight
+
+    def forward(self, x):
+        y = self.base(x)
+        return (y + F.linear(F.linear(self.dropout(x.to(self.lora_a.dtype)), self.lora_a), self.lora_b) * self.scale).to(y.dtype)
+
+
+def attach_lora(model: torch.nn.Module, path: str, scale: float = 1.0, dropout: float = 0.0) -> int:
+    """Unmerged LoRA: replace each target nn.Linear by a LoRALinear (alpha / r = 1 for the paper's r 32 / alpha 32)."""
+    pairs = _lora_pairs(path)
+    for module, a, b in pairs:
+        parent, _, name = module.rpartition(".")
+        setattr(model.get_submodule(parent), name, LoRALinear(model.get_submodule(module), a, b, scale, dropout))
+    return len(pairs)
+
+
+def lora_state(model: torch.nn.Module) -> dict:
+    """PEFT-format state dict (base_model.model.<module path>.lora_A / lora_B.weight) of the attached LoRALinear modules: the file
+    format every --lora option reads."""
+    state = {}
+    for name, mod in model.named_modules():
+        if isinstance(mod, LoRALinear):
+            state[f"base_model.model.{name}.lora_A.weight"] = mod.lora_a.detach().cpu().clone()
+            state[f"base_model.model.{name}.lora_B.weight"] = mod.lora_b.detach().cpu().clone()
+    return state
+
+
+def store_linear_weights_bf16(model: torch.nn.Module) -> None:
+    """CPU memory saver: nn.Linear weights (and the token-embedding table) stay bf16 and are upcast per call; everything else is fp32.
+
+    bf16 -> fp32 is exact, so outputs equal an fp32 model loaded from the bf16 checkpoint, at ~4.3 GB instead of ~8.5 GB resident.
+    changed: patches nn.Linear / nn.Embedding forward for the whole process (the upcast is inside every timed forward).
+    """
+    nn = torch.nn
+    for mod in model.modules():
+        keep_bf16 = isinstance(mod, nn.Linear) or (isinstance(mod, nn.Embedding) and mod.weight.shape[0] > 100000)
+        for name, p in mod.named_parameters(recurse=False):
+            if p.is_floating_point() and not (keep_bf16 and name == "weight"):
+                p.data = p.data.float()
+        for name, b in mod.named_buffers(recurse=False):
+            if b.is_floating_point():
+                setattr(mod, name, b.float())
+    embedding_forward = nn.Embedding.forward
+
+    def linear_upcast(self, x):
+        w = self.weight
+        return F.linear(x, w if w.dtype == x.dtype else w.to(x.dtype), self.bias)
+
+    nn.Linear.forward = linear_upcast
+    nn.Embedding.forward = lambda self, x: embedding_forward(self, x).float()
+
+
+def embed_patches(m: Model, pixels: torch.Tensor, linear: bool) -> torch.Tensor:
+    """ViT patch embedding of flattened patches (n, 3*2*16*16)."""
+    if linear:
+        # changed: Conv3d patch embed -> F.linear (identical math; MPS Conv3d runs at ~0.02 TFLOP/s)
+        return F.linear(pixels.reshape(-1, m.patch_w.shape[1]).to(m.patch_w.dtype), m.patch_w, m.patch_b)
+    pe = m.vit.patch_embed
+    return type(pe).forward(pe, pixels)        # stock HF Conv3d forward
+
+
+@dataclass
+class Grid:
+    """Constants of one (resolution, prompt) pair; identical for every frame of a stream, so computed once and cached."""
+    key: tuple
+    input_ids: torch.Tensor
+    vit_pos: torch.Tensor           # learned ViT position embedding per patch
+    vit_cos: torch.Tensor           # ViT 2-D RoPE per patch
+    vit_sin: torch.Tensor
+    token_emb: torch.Tensor         # (L_full, H) LLM input embeddings of the prompt (image slots are overwritten by ViT features)
+    llm_pos: torch.Tensor           # (rows, L_full) mRoPE position ids of the full-grid prompt
+    first: int                      # index of the first / last image token in the full-grid prompt
+    last: int
+    n_blocks: int                   # 880 for 1280x704
+    n_patches: int                  # 4 * n_blocks
+    length: int                     # L_full = prompt length with all blocks
+
+    @property
+    def text_tail(self) -> int:
+        """Prompt tokens after the image (instruction end + final token)."""
+        return self.length - self.last - 1
+
+    def llm_len(self, n_kept: int) -> int:
+        """LLM sequence length when only n_kept blocks are fed (Token Pruning drops the rest)."""
+        return self.first + n_kept + self.text_tail
+
+
+@torch.inference_mode()
+def grid(m: Model, inputs: dict) -> Grid:
+    """Position embeddings / RoPE / prompt embeddings, computed once per resolution.
+    changed: hoisted out of the per-frame forward (stock HF recomputes them every call)."""
+    thw, ids = inputs["image_grid_thw"], inputs["input_ids"]
+    key = (tuple(thw[0].tolist()), ids.shape[1])
+    g = m.grids.get(key)
+    if g is not None:
+        assert torch.equal(g.input_ids, ids), "same resolution and prompt length must mean the same prompt"
+        return g
+    V, cv = m.vit, m.vcfg
+    ii, iw = MQ.get_vision_interpolation_indices_and_weights(thw, num_grid_per_side=V.num_grid_per_side, mode=V.interpolation_mode,
+                                                             align_corners=V.interpolation_align_corners,
+                                                             spatial_merge_size=cv.spatial_merge_size, kwargs={})
+    pid = MQ.get_vision_position_ids(thw, V.spatial_merge_size, kwargs={})
+    pos = (V.pos_embed(ii) * iw[:, :, None]).sum(1)
+    cos, sin = V.rotary_pos_emb(pos, pid)
+    mm = inputs.get("mm_token_type_ids", (ids == m.image_token_id).int())
+    x0 = m.core.get_input_embeddings()(ids)
+    m.core.rope_deltas = None
+    p3 = m.core.compute_3d_position_ids(input_ids=ids, inputs_embeds=x0, image_grid_thw=thw, attention_mask=inputs["attention_mask"],
+                                        mm_token_type_ids=mm)
+    ipos = (ids[0] == m.image_token_id).nonzero()[:, 0]
+    first, last, n = int(ipos[0]), int(ipos[-1]), len(ipos)
+    assert last - first + 1 == n and int(thw.prod()) == 4 * n, "expected one contiguous run of image tokens, 4 patches each"
+    g = Grid(key=key, input_ids=ids, vit_pos=pos, vit_cos=cos, vit_sin=sin, token_emb=x0[0], llm_pos=p3[:, 0],
+             first=first, last=last, n_blocks=n, n_patches=4 * n, length=ids.shape[1])
+    m.grids[key] = g
+    return g

streamembed/engine/qwen_embedder.py

@@ -0,0 +1,106 @@
+"""Minimal Qwen3-VL-Embedding wrapper: processor + model + last-token pooling, images and text only.
+
+Trimmed from Qwen's reference `qwen3_vl_embedding.py` (Apache-2.0), which ships with Qwen3-VL-Embedding. Kept identical where it matters
+for the vectors: chat formatting, pixel limits, `process_vision_info(image_patch_size=16)`, processor call arguments, last-token pooling.
+Removed: video inputs, token truncation helper, the automatic move to CUDA in __init__ (the caller picks the device).
+The module tree is unchanged (`Qwen3VLForEmbedding.model` = HF `Qwen3VLModel`), so LoRA state-dict keys of the reference resolve as is.
+"""
+import unicodedata
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.nn.functional as F
+from qwen_vl_utils.vision_process import process_vision_info
+from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLConfig, Qwen3VLModel, Qwen3VLPreTrainedModel
+from transformers.models.qwen3_vl.processing_qwen3_vl import Qwen3VLProcessor
+
+MAX_LENGTH = 8192
+IMAGE_FACTOR = 32
+MIN_PIXELS = 4 * IMAGE_FACTOR * IMAGE_FACTOR
+MAX_PIXELS = 1800 * IMAGE_FACTOR * IMAGE_FACTOR
+DEFAULT_INSTRUCTION = "Represent the user's input."
+
+
+class Qwen3VLForEmbedding(Qwen3VLPreTrainedModel):
+    """Qwen3VLModel without an LM head; the embedding is the final hidden state of the last token."""
+    _checkpoint_conversion_mapping = {}      # (no `from __future__ import annotations` here: transformers reads `config`'s class)
+    accepts_loss_kwargs = False
+    config: Qwen3VLConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = Qwen3VLModel(config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.model.set_input_embeddings(value)
+
+    def forward(self, **inputs):
+        return self.model(**inputs)
+
+
+class Qwen3VLEmbedder:
+    def __init__(self, model_name_or_path: str, dtype: torch.dtype = torch.float32, max_length: int = MAX_LENGTH, local_files_only: bool = False,
+                 min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS, revision: str | None = None):
+        self.max_length, self.min_pixels, self.max_pixels = max_length, min_pixels, max_pixels
+        self.model = Qwen3VLForEmbedding.from_pretrained(model_name_or_path, dtype=dtype, local_files_only=local_files_only,
+                                                         revision=revision).eval()
+        self.processor = Qwen3VLProcessor.from_pretrained(model_name_or_path, padding_side="right", local_files_only=local_files_only,
+                                                          revision=revision)
+
+    def without_model(self) -> "Qwen3VLEmbedder":
+        """Processor-only shallow copy (cheap to send to DataLoader workers started with spawn, e.g. on macOS)."""
+        light = object.__new__(Qwen3VLEmbedder)
+        light.__dict__.update({k: v for k, v in self.__dict__.items() if k != "model"})
+        return light
+
+    def format_model_input(self, text: Optional[str] = None, image: Optional[str] = None, instruction: Optional[str] = None) -> list:
+        """One chat conversation: system = instruction, user = [image] + [text]."""
+        if instruction:
+            instruction = instruction.strip()
+            if instruction and not unicodedata.category(instruction[-1]).startswith("P"):
+                instruction += "."
+        content = []
+        if image is not None:
+            content.append({"type": "image", "image": image if image.startswith(("http://", "https://")) else "file://" + image,
+                            "min_pixels": self.min_pixels, "max_pixels": self.max_pixels})
+        if text is not None:
+            content.append({"type": "text", "text": text})
+        if not content:
+            content.append({"type": "text", "text": "NULL"})
+        return [{"role": "system", "content": [{"type": "text", "text": instruction or DEFAULT_INSTRUCTION}]},
+                {"role": "user", "content": content}]
+
+    def preprocess(self, conversations: List[list]) -> dict:
+        """Conversations -> processor tensors (CPU)."""
+        text = self.processor.apply_chat_template(conversations, add_generation_prompt=True, tokenize=False)
+        images, _, video_kwargs = process_vision_info(conversations, image_patch_size=16, return_video_metadata=True,
+                                                      return_video_kwargs=True)
+        return self.processor(text=text, images=images, videos=None, video_metadata=None, truncation=True, max_length=self.max_length,
+                              padding=True, do_resize=False, return_tensors="pt", **video_kwargs)
+
+    @staticmethod
+    def pool_last(hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        """Hidden state of the last non-padding token of each row (right padding)."""
+        col = attention_mask.shape[1] - attention_mask.flip(dims=[1]).argmax(dim=1) - 1
+        return hidden[torch.arange(hidden.shape[0], device=hidden.device), col]
+
+    @torch.inference_mode()
+    def embed(self, conversations: List[list], normalize: bool = True) -> torch.Tensor:
+        """Stock (unmodified HF) forward of a batch of conversations -> (B, 2048) embeddings."""
+        device = next(self.model.parameters()).device
+        inputs = {k: v.to(device) for k, v in self.preprocess(conversations).items()}
+        hidden = self.model.model(**inputs).last_hidden_state
+        e = self.pool_last(hidden, inputs["attention_mask"])
+        return F.normalize(e, p=2, dim=-1) if normalize else e
+
+    def embed_texts(self, texts: List[str], instruction: Optional[str] = None) -> torch.Tensor:
+        return self.embed([self.format_model_input(text=t, instruction=instruction) for t in texts])
+
+    def process(self, items: List[Dict[str, Any]]) -> torch.Tensor:
+        """Compatibility with the reference API: items = [{"text"|"image"|"instruction": ...}]."""
+        return self.embed([self.format_model_input(text=i.get("text"), image=i.get("image"), instruction=i.get("instruction"))
+                           for i in items])