# SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc. # SPDX-License-Identifier: Apache-2.0 import base64 import binascii import collections import io import os import sys import threading import time import urllib.request from concurrent.futures import Future from dataclasses import dataclass, field from pathlib import Path import numpy as np from PIL import Image HF_MODEL = "google/embeddinggemma-2" HF_REVISION = "914f7f89142e33e77833254d9c9b90c3cef7303b" SERVED_MODEL = "google/embeddinggemma-2" MODEL_ALIASES = (SERVED_MODEL, "embeddinggemma-2", "tt-hous/embeddinggemma-2-p150") DIMENSIONS = (768, 512, 256, 128) INPUT_TYPES = ("query", "document") SAMPLING_RATE = 16000 MAX_INPUTS = 2048 MAX_MEDIA_BYTES = 32 * 1024 * 1024 MAX_VIDEO_FRAMES = 1024 URL_TIMEOUT_S = 15 L1_SMALL_SIZE = 32768 TRACE_REGION_SIZE = 2_500_000_000 MEDIA_KEYS = ("image", "audio", "video") class RequestError(ValueError): def __init__(self, message, param=None, status=400, code=None): super().__init__(message) self.message = message self.param = param self.status = status self.code = code def parse_mesh_shape(text): parts = [p for p in text.lower().replace(",", "x").split("x") if p.strip()] if len(parts) != 2 or not all(p.strip().isdigit() for p in parts): raise ValueError(f"EMBEDDINGGEMMA_MESH_SHAPE must look like 1x1, got {text!r}") return int(parts[0]), int(parts[1]) def resolve_snapshot(model, revision): if Path(model).is_dir(): return Path(model) from huggingface_hub import snapshot_download try: return Path(snapshot_download(model, revision=revision, local_files_only=True)) except Exception: return Path(snapshot_download(model, revision=revision)) def strip_data_url(value): if value.startswith("data:"): head, sep, data = value.partition(",") if not sep or ";base64" not in head: raise RequestError("data URLs must be base64 encoded") return data return value def fetch_url(url): req = urllib.request.Request(url, headers={"User-Agent": "embeddinggemma-2-tt-server"}) try: with urllib.request.urlopen(req, timeout=URL_TIMEOUT_S) as resp: data = resp.read(MAX_MEDIA_BYTES + 1) except Exception as exc: raise RequestError(f"could not fetch {url}: {exc}") if len(data) > MAX_MEDIA_BYTES: raise RequestError(f"{url} is larger than {MAX_MEDIA_BYTES} bytes") return data def media_bytes(value, what): if isinstance(value, dict) and isinstance(value.get("url"), str): value = value["url"] if not isinstance(value, str) or not value: raise RequestError(f"{what} must be a non-empty base64 string or URL") if value.startswith(("http://", "https://")): if os.environ.get("EMBEDDINGGEMMA_ALLOW_URLS", "1") == "0": raise RequestError(f"{what}: URL inputs are disabled on this server; send base64") return fetch_url(value) try: data = base64.b64decode(strip_data_url(value), validate=True) except (binascii.Error, ValueError): raise RequestError(f"{what} is not valid base64 and not an http(s) URL") if not data: raise RequestError(f"{what} is empty") if len(data) > MAX_MEDIA_BYTES: raise RequestError(f"{what} is larger than {MAX_MEDIA_BYTES} bytes") return data def decode_image(value, what="image"): data = media_bytes(value, what) try: img = Image.open(io.BytesIO(data)) img.load() except Exception as exc: raise RequestError(f"{what} is not a decodable image: {exc}") return img def decode_audio(value, what="audio"): import soundfile as sf data = media_bytes(value, what) try: wav, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False) except Exception as exc: raise RequestError(f"{what} is not a decodable WAV file: {exc}") if sr != SAMPLING_RATE: raise RequestError(f"{what} sample rate is {sr} Hz; it must be {SAMPLING_RATE} Hz") if wav.ndim == 2: wav = wav.mean(axis=1, dtype=np.float32) if wav.size == 0: raise RequestError(f"{what} has no samples") return {"array": np.ascontiguousarray(wav, dtype=np.float32), "sampling_rate": sr} def decode_video(value, what="video"): if isinstance(value, list): value = {"frames": value} if not isinstance(value, dict) or not isinstance(value.get("frames"), list) or not value["frames"]: raise RequestError(f"{what} must be an object with a non-empty 'frames' list") if len(value["frames"]) > MAX_VIDEO_FRAMES: raise RequestError(f"{what} has {len(value['frames'])} frames; the limit is {MAX_VIDEO_FRAMES}") frames = [np.asarray(decode_image(f, f"{what}.frames[{i}]").convert("RGB")) for i, f in enumerate(value["frames"])] if len({f.shape for f in frames}) != 1: raise RequestError(f"{what} frames must all have the same size") arr = np.stack(frames) fps = value.get("fps") metadata = None if fps is not None: if not isinstance(fps, (int, float)) or isinstance(fps, bool) or fps <= 0: raise RequestError(f"{what}.fps must be a positive number") metadata = {"fps": float(fps), "total_num_frames": len(frames), "duration": len(frames) / float(fps)} return {"array": arr, "video_metadata": metadata} def decode_input(item, index): what = f"input[{index}]" if isinstance(item, str): if not item: raise RequestError(f"{what} is an empty string", param="input") return item if not isinstance(item, dict): raise RequestError(f"{what} must be a string or an object with text/image/audio/video fields", param="input") bad = sorted(set(item) - {"text", *MEDIA_KEYS}) if bad: raise RequestError(f"{what} has unsupported fields {bad}", param="input") media = [k for k in MEDIA_KEYS if item.get(k) is not None] text = item.get("text") if text is not None and not isinstance(text, str): raise RequestError(f"{what}.text must be a string", param="input") if len(media) > 1: raise RequestError(f"{what} mixes {media}; one media kind per input is supported", param="input") if not media: if not text: raise RequestError(f"{what} has no text, image, audio or video", param="input") return text kind = media[0] decoder = {"image": decode_image, "audio": decode_audio, "video": decode_video}[kind] out = {} for key in item: if key == kind: out[kind] = decoder(item[kind], f"{what}.{kind}") elif key == "text" and text: out["text"] = text return out def normalize_inputs(raw): if isinstance(raw, str): raw = [raw] if not isinstance(raw, list): raise RequestError("input must be a string or a list of strings or objects", param="input") if not raw: raise RequestError("input must not be empty", param="input") if len(raw) > MAX_INPUTS: raise RequestError(f"input has {len(raw)} items; the limit is {MAX_INPUTS}", param="input") if all(isinstance(x, int) and not isinstance(x, bool) for x in raw) or any(isinstance(x, list) for x in raw): raise RequestError("token id inputs are not supported; send strings", param="input") return [decode_input(x, i) for i, x in enumerate(raw)] @dataclass class Job: items: list dim: int key: tuple = () future: Future = field(default_factory=Future) enqueued: float = field(default_factory=time.perf_counter) class Engine: def __init__(self, embedder, mesh_device, policy_id=None, max_batch_items=32, batch_wait_ms=0.0): self.embedder = embedder self.mesh_device = mesh_device self.policy_id = policy_id self.max_batch_items = max_batch_items self.batch_wait_s = batch_wait_ms / 1000.0 self.device_lock = threading.Lock() self.text_lock = threading.Lock() self.media_lock = threading.Lock() self.bucket_grouping = os.environ.get("EMBEDDINGGEMMA_BUCKET_GROUPING", "1") != "0" self.pending = collections.deque() self.cv = threading.Condition() self.closing = False sizes = {b for b, _, _ in embedder.traces.slots} if embedder.traces is not None else set() self.slot_sizes = sorted(b for b in sizes if b <= max_batch_items) or [max_batch_items] self.slot_cost = {} self.record = None self.stats = { "requests": 0, "inputs": 0, "device_batches": 0, "encode_groups": 0, "max_jobs_per_batch": 0, "device_busy_s": 0.0, "device_thread_errors": 0, } self.thread = threading.Thread(target=self._loop, name="embeddinggemma-device", daemon=True) self.started = time.time() self.thread.start() @property def max_tokens(self): return self.embedder.max_tokens def prompt_for(self, input_type): if input_type is None: return None if input_type not in self.embedder.prompts: known = sorted(set(INPUT_TYPES) | set(self.embedder.prompts)) raise RequestError(f"unknown input_type {input_type!r}; known: {known}", param="input_type") return self.embedder.prompts[input_type] def prepare(self, inputs, input_type=None): from models.autoports.google_embeddinggemma_2.tt.multimodal import item_kind prompt = self.prompt_for(input_type) emb = self.embedder out = [] for i, x in enumerate(inputs): try: if isinstance(x, str): with self.text_lock: ids = emb.tokenizer([(prompt or "") + x], add_special_tokens=True)["input_ids"][0] it = {"input_ids": ids} else: with self.media_lock: it = emb.process([x], prompt=prompt or "")[0] except RequestError: raise except ValueError as exc: raise RequestError(f"input[{i}]: {exc}", param="input") n = len(it["input_ids"]) if emb.max_tokens is not None and n > emb.max_tokens: raise RequestError( f"input[{i}] has {n} tokens; the maximum is {emb.max_tokens}", param="input", code="context_length_exceeded", ) kind = item_kind(it, emb.pipe.tokens) if kind is not None and emb.traces is not None and emb.media_slot_for(kind, [it]) is None: raise RequestError( f"input[{i}] ({kind}, {n} tokens) does not fit any captured {kind} trace slot", param="input", ) out.append(it) return out def item_key(self, it): from models.autoports.google_embeddinggemma_2.tt.model import uniform_physical_length from models.autoports.google_embeddinggemma_2.tt.multimodal import item_kind, item_patches kind = item_kind(it, self.embedder.pipe.tokens) if kind is not None: return ("media", kind, item_patches(it)) return ("text", uniform_physical_length(self.embedder.seq_len_for([len(it["input_ids"])]))) def submit(self, items, dim): if dim not in DIMENSIONS: raise RequestError(f"dimensions must be one of {list(DIMENSIONS)}", param="dimensions") job = Job(items, dim, tuple(sorted({self.item_key(it) for it in items}))) with self.cv: self.pending.append(job) self.cv.notify() return job.future def queue_depth(self): return len(self.pending) def run_direct(self, fn): with self.device_lock: return fn(self.embedder) def target_items(self, n, phys=None): if n >= self.slot_sizes[-1]: return max(n, self.slot_sizes[-1]) if n > self.max_batch_items else self.slot_sizes[-1] up = min(b for b in self.slot_sizes if b >= n) down = max((b for b in self.slot_sizes if b <= n), default=n) cost_up, cost_down = self.slot_cost.get((up, phys)), self.slot_cost.get((down, phys)) if cost_up is not None and cost_down is not None: return n if cost_up / n <= cost_down / down else down return n if up - n <= n - down else down def text_phys(self, jobs): from models.autoports.google_embeddinggemma_2.tt.model import uniform_physical_length items = [it for j in jobs for it in j.items] if any("pixel_values" in it or "input_features" in it for it in items): return None return uniform_physical_length(self.embedder.seq_len_for([max(len(it["input_ids"]) for it in items)])) def _take(self): with self.cv: while not self.pending and not self.closing: self.cv.wait() if not self.pending: return None deadline = time.perf_counter() + self.batch_wait_s while not self.closing and sum(len(j.items) for j in self.pending) < self.max_batch_items: left = deadline - time.perf_counter() if left <= 0: break self.cv.wait(left) head = self.pending[0].key avail, n = [], 0 for j in self.pending: if self.bucket_grouping and j.key != head: continue if avail and n + len(j.items) > self.max_batch_items: break avail.append(j) n += len(j.items) phys = head[0][1] if len(head) == 1 and head[0][0] == "text" else None k = self.target_items(n, phys if self.bucket_grouping else self.text_phys(avail)) jobs, c = [], 0 for j in avail: if jobs and c + len(j.items) > k: break jobs.append(j) c += len(j.items) taken = {id(j) for j in jobs} self.pending = collections.deque(j for j in self.pending if id(j) not in taken) return jobs def _encode(self, jobs): dims = tuple(sorted({j.dim for j in jobs}, reverse=True)) rows = [(j, k, it) for j in jobs for k, it in enumerate(j.items)] groups = {} for r in rows: key = self.item_key(r[2]) if self.bucket_grouping else () groups.setdefault(key, []).append(r) results = {id(j): np.zeros((len(j.items), j.dim), dtype=np.float32) for j in jobs} for key, grp in groups.items(): items = [it for _, _, it in grp] with self.device_lock: t0 = time.perf_counter() out = self.embedder.encode_processed(items, dims=dims) self.stats["device_busy_s"] += time.perf_counter() - t0 self.stats["encode_groups"] += 1 if self.record is not None: self.record.append( { "key": key, "items": items, "dims": dims, "row_dims": [j.dim for j, _, _ in grp], "out": out, "n_jobs": len(jobs), } ) for pos, (j, k, _) in enumerate(grp): results[id(j)][k] = out[j.dim][pos] for j in jobs: self._resolve(j, results[id(j)]) @staticmethod def _resolve(job, value=None, exc=None): if job.future.done(): return if exc is not None: job.future.set_exception(exc) else: job.future.set_result(value) def _loop(self): while True: jobs = self._take() if jobs is None: return try: self.stats["device_batches"] += 1 self.stats["max_jobs_per_batch"] = max(self.stats["max_jobs_per_batch"], len(jobs)) try: self._encode(jobs) except Exception as exc: if len(jobs) == 1: self._resolve(jobs[0], exc=exc) continue for j in jobs: try: self._encode([j]) except Exception as one: self._resolve(j, exc=one) except Exception as exc: self.stats["device_thread_errors"] += 1 for j in jobs: try: self._resolve(j, exc=exc) except Exception: pass def alive(self): return self.thread.is_alive() def trace_counters(self): traces = self.embedder.traces if traces is None: return {"traced": False} slots = {f"text:{b}x{p}": s for (b, p, _), s in traces.slots.items()} slots.update( {f"{s.shape.kind}:{s.shape.batch}x{s.shape.physical_len}:tower{s.shape.tower_n}": s for s in traces.extra} ) per = { k: {"trace_id": str(s.trace_id), "replays": s.replays, "refreshes": s.refreshes} for k, s in slots.items() } return { "captures": traces.captures, "live_traces": sum(1 for s in slots.values() if s.trace_id is not None), "replays_total": sum(s.replays for s in slots.values()), "slots": per, } def warm(self): traces = self.embedder.traces if traces is None: return 0 import ttnn slots = list(traces.slots.values()) + list(traces.extra) with self.device_lock: for s in slots: s.replay() ttnn.synchronize_device(self.mesh_device) for (batch, phys, exact), s in traces.slots.items(): if exact: continue t0 = time.perf_counter() s.replay() ttnn.synchronize_device(self.mesh_device) self.slot_cost[(batch, phys)] = time.perf_counter() - t0 return len(slots) def slot_summary(self): traces = self.embedder.traces if traces is None: return {"traced": False} text = sorted(traces.slots) media = [s.shape for s in traces.extra] return { "traced": True, "text_slots": len(text), "media_slots": len(media), "text_buckets": sorted({p for _, p, _ in text}), "text_batches": sorted({b for b, _, _ in text}), "media": [f"{m.kind}:b{m.batch}:len{m.physical_len}:tower{m.tower_n}" for m in media], } def close(self): with self.cv: self.closing = True self.cv.notify_all() self.thread.join(timeout=60) def build_engine(): import ttnn from models.autoports.google_embeddinggemma_2.tt import precision_policy as PP from models.autoports.google_embeddinggemma_2.tt.embedder import build_multimodal_embedder sys.setswitchinterval(float(os.environ.get("EMBEDDINGGEMMA_SWITCH_INTERVAL_MS", "0.5")) / 1000.0) rows, cols = parse_mesh_shape(os.environ.get("EMBEDDINGGEMMA_MESH_SHAPE", "1x1")) if (rows, cols) != (1, 1): raise ValueError(f"the embedder runs on one p150 (mesh 1x1); EMBEDDINGGEMMA_MESH_SHAPE={rows}x{cols}") snapshot = resolve_snapshot(os.environ.get("HF_MODEL", HF_MODEL), os.environ.get("HF_REVISION", HF_REVISION)) mesh = ttnn.open_mesh_device( ttnn.MeshShape(rows, cols), l1_small_size=L1_SMALL_SIZE, trace_region_size=TRACE_REGION_SIZE ) try: policy = PP.load() embedder = build_multimodal_embedder(snapshot, mesh, traced=True, precision_policy=policy) report = PP.runtime_report(embedder) if not report["consumed"]: raise RuntimeError(f"precision policy {policy.id} not applied: {report['mismatches'][:4]}") except Exception: ttnn.close_mesh_device(mesh) raise g = mesh.compute_with_storage_grid_size() engine = Engine( embedder, mesh, policy_id=policy.id, max_batch_items=int(os.environ.get("EMBEDDINGGEMMA_MAX_BATCH_ITEMS", "32")), batch_wait_ms=float(os.environ.get("EMBEDDINGGEMMA_BATCH_WAIT_MS", "0")), ) engine.snapshot = str(snapshot) engine.grid = f"{g.x}x{g.y}" engine.mesh_shape = f"{rows}x{cols}" return engine def shutdown_engine(engine): import ttnn engine.close() with engine.device_lock: engine.embedder.release() ttnn.close_mesh_device(engine.mesh_device)