iris-3b Turbo on LiteRT: 20 fp16-weight graphs for the phone GPU at FP32, a Python reference pipeline and the Android engine
743fd20 verified | # SPDX-License-Identifier: Apache-2.0 | |
| """iris-3b Turbo text-to-image on LiteRT CompiledModel: the Python reference pipeline. | |
| A prompt becomes a 768x768 image through the 20 graphs in ../model: | |
| tokenize (Qwen3-VL byte-level BPE) -> bf16 embedding lookup | |
| -> 12 text encoder graphs (one per hidden-state tap) -> window rows x mask | |
| -> adapter (once per prompt) | |
| -> 4 steps x (dual0..2 -> [text; image] -> single0..2 -> image rows -> pixel stage -> sampler) | |
| -> clamp -> patch-major pixels to a CHW image | |
| It reads the same files and runs the same host-side steps as the Android engine in ../android | |
| (IrisPipeline.kt, HostOps.kt, PromptTokens.kt, QwenTokenizer.kt). Every graph runs on the CPU | |
| (XNNPACK); the phone runs the same graphs on its GPU at FP32. A seed's starting noise is | |
| java.util.Random(seed).nextGaussian(), as in the Android app. | |
| pip install -r requirements.txt | |
| python iris3b_litert.py --model-dir ../model --prompt "a red apple on a wooden table" --out apple.png | |
| """ | |
| import argparse | |
| import json | |
| import math | |
| import os | |
| import resource | |
| import sys | |
| import time | |
| import unicodedata | |
| from pathlib import Path | |
| import numpy as np | |
| import regex | |
| CONFIG_FILE = "iris3b_config.json" | |
| SEQ = 334 | |
| PREFIX = 34 | |
| WINDOW = 300 | |
| TE_HIDDEN = 2560 | |
| TE_HEADS = 32 | |
| TAPS = 12 | |
| ADAPTER_HEADS = 20 | |
| DIT_HIDDEN = 2560 | |
| IMAGE_SIZE = 768 | |
| PATCH = 16 | |
| GRID = IMAGE_SIZE // PATCH | |
| PATCHES = GRID * GRID | |
| PIXELS = PATCH * PATCH | |
| CHANNELS = 3 | |
| IMAGE_FLOATS = CHANNELS * IMAGE_SIZE * IMAGE_SIZE | |
| MASK_NEG = np.float32(-30000.0) | |
| PAD_TOKEN_ID = 151643 | |
| # The Unicode White_Space set, written out (the Android port does the same so that every regex | |
| # engine agrees on it), and tokenizer.json's Split pattern with `\s` and `(?i:...)` written out. | |
| _WS = "\t\n\x0b\x0c\r \x85\xa0 - " | |
| _PRE_TOKENIZER = regex.compile( | |
| r"'(?:[sSſ]|[tT]|[rR][eE]|[vV][eE]|[mM]|[lL][lL]|[dD])" | |
| r"|[^\r\n\p{L}\p{N}]?\p{L}+" | |
| r"|\p{N}" | |
| r"| ?[^" + _WS + r"\p{L}\p{N}]+[\r\n]*" | |
| r"|[" + _WS + r"]*[\r\n]+" | |
| r"|[" + _WS + r"]+(?![^" + _WS + r"])" | |
| r"|[" + _WS + r"]+" | |
| ) | |
| class QwenTokenizer: | |
| """Byte-level BPE of the Qwen3-VL-4B-Instruct text encoder (a port of QwenTokenizer.kt). | |
| Added tokens are split out first (leftmost, then longest), the text between them is NFC | |
| normalized, pre-tokenized with the Split pattern, byte-level encoded (GPT-2 bytes_to_unicode) | |
| and merged pair by pair in rank order. | |
| """ | |
| def __init__(self, files): | |
| vocab = files["qwen_vocab.txt"].split("\n") | |
| if vocab[-1] == "": | |
| vocab.pop() | |
| self._tokens = vocab | |
| self._token_to_id = {token: i for i, token in enumerate(vocab)} | |
| self._ranks = {} | |
| for rank, line in enumerate(files["qwen_merges.txt"].split("\n")): | |
| space = line.find(" ") | |
| if space < 0: | |
| continue | |
| left = self._token_to_id.get(line[:space]) | |
| right = self._token_to_id.get(line[space + 1:]) | |
| if left is not None and right is not None: | |
| self._ranks[(left, right)] = rank | |
| added = [] | |
| for line in files["qwen_added.txt"].split("\n"): | |
| if line: | |
| tab = line.rindex("\t") | |
| added.append((line[:tab], int(line[tab + 1:]))) | |
| added.sort(key=lambda t: -len(t[0])) | |
| self._added = {} | |
| for token, token_id in added: | |
| self._added.setdefault(token[0], []).append((token, token_id)) | |
| printable = set(range(ord("!"), ord("~") + 1)) | set(range(ord("¡"), ord("¬") + 1)) | set( | |
| range(ord("®"), ord("ÿ") + 1)) | |
| byte_to_char, extra = [], 0 | |
| for byte in range(256): | |
| if byte in printable: | |
| byte_to_char.append(chr(byte)) | |
| else: | |
| byte_to_char.append(chr(256 + extra)) | |
| extra += 1 | |
| self._byte_ids = [self._token_to_id[c] for c in byte_to_char] | |
| self._cache = {} | |
| def encode(self, text): | |
| """Token ids of raw text, no template and no special tokens added.""" | |
| ids, start, i = [], 0, 0 | |
| while i < len(text): | |
| hit = next((a for a in self._added.get(text[i], ()) if text.startswith(a[0], i)), None) | |
| if hit is None: | |
| i += 1 | |
| continue | |
| if start < i: | |
| self._encode_segment(text[start:i], ids) | |
| ids.append(hit[1]) | |
| i += len(hit[0]) | |
| start = i | |
| if start < len(text): | |
| self._encode_segment(text[start:], ids) | |
| return ids | |
| def _encode_segment(self, segment, out): | |
| text = unicodedata.normalize("NFC", segment) | |
| position = 0 | |
| while position < len(text): | |
| match = _PRE_TOKENIZER.search(text, position) | |
| if match is None: | |
| out.extend(self._merge(text[position:])) | |
| break | |
| if match.start() > position: | |
| out.extend(self._merge(text[position:match.start()])) | |
| out.extend(self._merge(match.group())) | |
| position = match.end() | |
| def _merge(self, piece): | |
| cached = self._cache.get(piece) | |
| if cached is not None: | |
| return cached | |
| symbols = [self._byte_ids[b] for b in piece.encode("utf-8")] | |
| while len(symbols) > 1: | |
| best_rank, best_index = None, -1 | |
| for i in range(len(symbols) - 1): | |
| rank = self._ranks.get((symbols[i], symbols[i + 1])) | |
| if rank is not None and (best_rank is None or rank < best_rank): | |
| best_rank, best_index = rank, i | |
| if best_index < 0: | |
| break | |
| merged = self._tokens[symbols[best_index]] + self._tokens[symbols[best_index + 1]] | |
| symbols[best_index:best_index + 2] = [self._token_to_id[merged]] | |
| self._cache[piece] = symbols | |
| return symbols | |
| class PromptTokens: | |
| """The prompt layout and the bf16 embedding lookup (a port of PromptTokens.kt). | |
| System-prompt prefix (34 ids) + caption cut to 295 ids + assistant-turn suffix (5 ids), | |
| right-padded with <|endoftext|> to 334; the real tokens are prefix + caption + suffix. | |
| """ | |
| def __init__(self, files, embed_path): | |
| self.tokenizer = QwenTokenizer(files) | |
| self.prefix = [int(v) for v in files["prompt_prefix_ids.txt"].strip().split(",")] | |
| self.suffix = [int(v) for v in files["prompt_suffix_ids.txt"].strip().split(",")] | |
| self.table = np.memmap(embed_path, dtype="<u2", mode="r").reshape(-1, TE_HIDDEN) | |
| def encode(self, caption): | |
| """(ids [334], real token count, caption tokens before the cut).""" | |
| caption_ids = self.tokenizer.encode(caption) | |
| n = min(len(caption_ids), WINDOW - len(self.suffix)) | |
| ids = self.prefix + caption_ids[:n] + self.suffix | |
| real = len(ids) | |
| return np.array(ids + [PAD_TOKEN_ID] * (SEQ - real), np.int64), real, len(caption_ids) | |
| def embed(self, ids): | |
| """inputs_embeds [334, 2560] fp32 from the bf16 table (exact: bf16 -> fp32 is a 16-bit shift).""" | |
| return (self.table[ids].astype(np.uint32) << 16).view(np.float32) | |
| def te_mask(real): | |
| """Additive causal mask [32, 334, 334]: 0 where key <= query and the key is real.""" | |
| query = np.arange(SEQ)[:, None] | |
| key = np.arange(SEQ)[None, :] | |
| plane = np.where((key <= query) & (key < real), np.float32(0), MASK_NEG).astype(np.float32) | |
| return np.ascontiguousarray(np.broadcast_to(plane, (TE_HEADS, SEQ, SEQ))) | |
| def window_mask(real): | |
| """1 on the caption + suffix rows of the 300-row window.""" | |
| return (PREFIX + np.arange(WINDOW) < real).astype(np.float32) | |
| def adapter_mask_bias(wmask): | |
| """Adapter key mask [20, 300, 300]: a query sees the real keys and itself.""" | |
| keep = (wmask[None, :] != 0) | np.eye(WINDOW, dtype=bool) | |
| plane = np.where(keep, np.float32(0), MASK_NEG).astype(np.float32) | |
| return np.ascontiguousarray(np.broadcast_to(plane, (ADAPTER_HEADS, WINDOW, WINDOW))) | |
| def sampler_step(x, v, sigma, next_sigma): | |
| """Turbo few-step sampler in fp32: clean = x - sigma v, then x at the next sigma (if any).""" | |
| clean = x - sigma * v | |
| if next_sigma is None: | |
| return x, clean | |
| implied = (x - (np.float32(1) - sigma) * clean) / sigma | |
| return (np.float32(1) - next_sigma) * clean + next_sigma * implied, clean | |
| def to_image(x_pix): | |
| """Patch-major pixels [2304, 256, 3] -> CHW image [3, 768, 768].""" | |
| return x_pix.reshape(GRID, GRID, PATCH, PATCH, CHANNELS).transpose(4, 0, 2, 1, 3).reshape( | |
| CHANNELS, IMAGE_SIZE, IMAGE_SIZE) | |
| def seed_noise(seed): | |
| """Standard normal noise [IMAGE_FLOATS] from java.util.Random(seed).nextGaussian(), as the app.""" | |
| mult, add, mask = 0x5DEECE66D, 0xB, (1 << 48) - 1 | |
| unit = 2.0 ** -53 | |
| state = (seed ^ mult) & mask | |
| out = [] | |
| log, sqrt = math.log, math.sqrt | |
| while len(out) < IMAGE_FLOATS: | |
| while True: | |
| state = (state * mult + add) & mask | |
| hi = state >> 22 | |
| state = (state * mult + add) & mask | |
| v1 = 2 * (((hi << 27) + (state >> 21)) * unit) - 1 | |
| state = (state * mult + add) & mask | |
| hi = state >> 22 | |
| state = (state * mult + add) & mask | |
| v2 = 2 * (((hi << 27) + (state >> 21)) * unit) - 1 | |
| s = v1 * v1 + v2 * v2 | |
| if 0 < s < 1: | |
| break | |
| multiplier = sqrt(-2 * log(s) / s) | |
| out.append(v1 * multiplier) | |
| out.append(v2 * multiplier) | |
| return np.array(out[:IMAGE_FLOATS], np.float64).astype(np.float32) | |
| def save_png(image, path): | |
| """Clamped CHW image in [-1, 1] -> 8-bit RGB PNG (the app's rounding).""" | |
| from PIL import Image | |
| level = (image.astype(np.float32) + np.float32(1)) * np.float32(127.5) + np.float32(0.5) | |
| rgb = np.clip(np.nan_to_num(level, nan=0.0), 0, 255).astype(np.uint8).transpose(1, 2, 0) | |
| Image.fromarray(np.ascontiguousarray(rgb)).save(path) | |
| class _Graph: | |
| """One graph file on LiteRT CompiledModel (CPU), driven as its config entry says.""" | |
| def __init__(self, path, spec, threads, environment): | |
| from ai_edge_litert.compiled_model import CompiledModel, CpuOptions, HardwareAccelerator, Options | |
| options = Options(hardware_accelerators=HardwareAccelerator.CPU, | |
| cpu_options=CpuOptions(num_threads=threads)) | |
| self.model = CompiledModel.from_file(str(path), options=options, environment=environment) | |
| key = spec["signature"] | |
| self.index = self.model.get_signature_index(key) | |
| names = self.model.get_signature_list()[key] | |
| self.inputs = self.model.create_input_buffers(self.index) | |
| self.outputs = self.model.create_output_buffers(self.index) | |
| self.feed = [names["inputs"].index(name) for name in spec["inputs"]] | |
| self.read = [(names["outputs"].index(name), count) for name, count in spec["outputs"]] | |
| def run(self, feeds): | |
| for position, values in zip(self.feed, feeds): | |
| self.inputs[position].write(np.ascontiguousarray(values, np.float32)) | |
| self.model.run_by_index(self.index, self.inputs, self.outputs) | |
| return [self.outputs[position].read(count, np.float32) for position, count in self.read] | |
| def close(self): | |
| self.inputs = self.outputs = None | |
| self.model.close() | |
| self.model = None | |
| class Iris3B: | |
| """The pipeline over a model folder (the 24 files of ../model).""" | |
| def __init__(self, model_dir, threads=None, keep_dit=False): | |
| self.dir = Path(model_dir) | |
| self.config = json.loads((self.dir / CONFIG_FILE).read_text(encoding="utf-8")) | |
| graphs = self.config["graphs"] | |
| self.te_specs = graphs["te"] | |
| self.adapter_spec = graphs["adapter"] | |
| self.dual = list(zip(graphs["dual"], graphs["dual_mod_rows"])) | |
| self.single = list(zip(graphs["single"], graphs["single_mod_rows"])) | |
| self.pixel_spec = graphs["pixel"] | |
| self._host_file = self.dir / self.config["host"]["file"] | |
| self._host_index = self.config["host"]["tensors"] | |
| self._host = {} | |
| files = json.loads((self.dir / self.config["tokenizer"]).read_text(encoding="utf-8"))["files"] | |
| self.tokens = PromptTokens(files, self.dir / self.config["embed"]) | |
| self.threads = threads or os.cpu_count() | |
| self.keep_dit = keep_dit | |
| self._environment = None | |
| self._resident = {} | |
| self.timeline = [] | |
| def host(self, name): | |
| """A host tensor (rotary tables, modulation rows, sigmas, reference noise ...) by name.""" | |
| if name not in self._host: | |
| t = self._host_index[name] | |
| self._host[name] = np.fromfile(self._host_file, dtype="<f4", count=t["count"], offset=t["offset"]) | |
| return self._host[name] | |
| def reference_noise(self): | |
| """The upstream reference run's starting noise (torch.randn, seed 0), patch-major.""" | |
| return self.host("noise_pix").copy() | |
| def generate(self, prompt, noise): | |
| """Prompt + noise [IMAGE_FLOATS] -> clamped CHW image [3, 768, 768] in [-1, 1].""" | |
| self.timeline = [] | |
| t0 = time.perf_counter() | |
| ids, real, caption_tokens = self.tokens.encode(prompt) | |
| self.real_len, self.caption_tokens = real, caption_tokens | |
| hidden = self.tokens.embed(ids) | |
| mask, wmask = te_mask(real), window_mask(real) | |
| cos, sin = self.host("te_rope_cos"), self.host("te_rope_sin") | |
| te = np.zeros((WINDOW, TAPS, TE_HIDDEN), np.float32) | |
| for tap, spec in enumerate(self.te_specs): | |
| hidden = self._once(spec, [hidden, mask, cos, sin])[0] | |
| te[:, tap, :] = hidden.reshape(SEQ, TE_HIDDEN)[PREFIX:] * wmask[:, None] | |
| y0 = self._dit(self.adapter_spec, [te, adapter_mask_bias(wmask), self.host("y_pos")])[0] | |
| sigmas = self.host("sigmas") | |
| cos_j, sin_j = self.host("cos_j"), self.host("sin_j") | |
| cos_i, sin_i, pix_pos = self.host("cos_i"), self.host("sin_i"), self.host("pix_pos") | |
| x = np.asarray(noise, np.float32).reshape(-1).copy() | |
| assert x.size == IMAGE_FLOATS, x.size | |
| clean = None | |
| for step in range(sigmas.size): | |
| mods = self.host(f"mods_s{step}").reshape(-1, DIT_HIDDEN) | |
| image, text = x, y0 | |
| for spec, (first, end) in self.dual: | |
| outs = self._dit(spec, [image, text, cos_j, sin_j, mods[first:end]]) | |
| image = _by_size(outs, PATCHES * DIT_HIDDEN) | |
| text = _by_size(outs, WINDOW * DIT_HIDDEN) | |
| joint = np.concatenate([text, image]) | |
| for spec, (first, end) in self.single: | |
| joint = self._dit(spec, [joint, cos_j, sin_j, mods[first:end]])[0] | |
| v = self._dit(self.pixel_spec, [joint[WINDOW * DIT_HIDDEN:], self.host(f"t_emb_s{step}"), x, pix_pos, | |
| cos_i, sin_i])[0] | |
| next_sigma = sigmas[step + 1] if step + 1 < sigmas.size else None | |
| x, clean = sampler_step(x, v, sigmas[step], next_sigma) | |
| print(f"step {step + 1}/{sigmas.size} done at {time.perf_counter() - t0:.1f} s", file=sys.stderr) | |
| for graph in self._resident.values(): | |
| graph.close() | |
| self._resident = {} | |
| self.wall_s = time.perf_counter() - t0 | |
| return to_image(np.clip(clean, -1, 1)) | |
| def _env(self): | |
| if self._environment is None: | |
| import ai_edge_litert.compiled_model as cm | |
| self._environment = cm.Environment.create( | |
| options=cm.EnvironmentOptions(runtime_path=os.path.dirname(os.path.abspath(cm.__file__)))) | |
| return self._environment | |
| def _load(self, spec): | |
| t0 = time.perf_counter() | |
| graph = _Graph(self.dir / spec["file"], spec, self.threads, self._env()) | |
| self.timeline.append({"graph": spec["role"], "event": "load", "s": time.perf_counter() - t0}) | |
| return graph | |
| def _run(self, graph, spec, feeds): | |
| t0 = time.perf_counter() | |
| outs = graph.run(feeds) | |
| self.timeline.append({"graph": spec["role"], "event": "run", "s": time.perf_counter() - t0}) | |
| return outs | |
| def _once(self, spec, feeds): | |
| graph = self._load(spec) | |
| try: | |
| return self._run(graph, spec, feeds) | |
| finally: | |
| graph.close() | |
| def _dit(self, spec, feeds): | |
| if not self.keep_dit: | |
| return self._once(spec, feeds) | |
| graph = self._resident.get(spec["role"]) | |
| if graph is None: | |
| graph = self._resident[spec["role"]] = self._load(spec) | |
| return self._run(graph, spec, feeds) | |
| def _by_size(outputs, count): | |
| hits = [o for o in outputs if o.size == count] | |
| assert len(hits) == 1, [o.size for o in outputs] | |
| return hits[0] | |
| def main(): | |
| here = Path(__file__).resolve().parent | |
| ap = argparse.ArgumentParser(description="iris-3b Turbo text-to-image on LiteRT CompiledModel (CPU)") | |
| ap.add_argument("--model-dir", default=str(here.parent / "model"), help="folder with the 24 model files") | |
| ap.add_argument("--prompt", required=True) | |
| ap.add_argument("--seed", type=int, default=0, help="java.util.Random seed of the starting noise (as the app)") | |
| ap.add_argument("--reference-noise", action="store_true", | |
| help="use the upstream reference run's noise from the host pack instead of --seed") | |
| ap.add_argument("--out", default="iris3b.png", help="PNG to write") | |
| ap.add_argument("--npy", help="also write the clamped CHW image [1, 3, 768, 768] float32 here") | |
| ap.add_argument("--threads", type=int, default=None, help="XNNPACK threads (default: all cores)") | |
| ap.add_argument("--keep-dit", action="store_true", | |
| help="keep the 8 DiT graphs loaded across the 4 steps instead of loading each graph for " | |
| "every use (uses more memory)") | |
| ap.add_argument("--report", help="write timings and memory as JSON here") | |
| args = ap.parse_args() | |
| pipe = Iris3B(args.model_dir, threads=args.threads, keep_dit=args.keep_dit) | |
| noise = pipe.reference_noise() if args.reference_noise else seed_noise(args.seed) | |
| image = pipe.generate(args.prompt, noise) | |
| save_png(image, args.out) | |
| if args.npy: | |
| np.save(args.npy, image[None].astype(np.float32)) | |
| nonfinite = int(np.size(image) - np.count_nonzero(np.isfinite(image))) | |
| rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss | |
| peak_bytes = rss if sys.platform == "darwin" else rss * 1024 | |
| print(f"{args.out}: {pipe.wall_s:.1f} s, real tokens {pipe.real_len}, caption tokens {pipe.caption_tokens}, " | |
| f"non-finite {nonfinite}, peak RSS {peak_bytes / 2**30:.2f} GiB", file=sys.stderr) | |
| if args.report: | |
| loads = sum(e["s"] for e in pipe.timeline if e["event"] == "load") | |
| runs = sum(e["s"] for e in pipe.timeline if e["event"] == "run") | |
| Path(args.report).write_text(json.dumps({ | |
| "prompt": args.prompt, "seed": None if args.reference_noise else args.seed, | |
| "noise": "reference" if args.reference_noise else "java.util.Random", "threads": pipe.threads, | |
| "keep_dit": args.keep_dit, "wall_s": pipe.wall_s, "load_s": loads, "run_s": runs, | |
| "real_len": pipe.real_len, "caption_tokens": pipe.caption_tokens, "nonfinite": nonfinite, | |
| "peak_rss_bytes": peak_bytes, "timeline": pipe.timeline}, indent=1)) | |
| return 0 if nonfinite == 0 else 1 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |