| |
| """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 |
|
|
| |
| |
| _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()) |
|
|
