diff --git a/build.zig b/build.zig index 08765bd1df..e9a22d3dfd 100644 --- a/build.zig +++ b/build.zig @@ -12,6 +12,7 @@ const counters_build = @import("zig/build/metal_counters.zig"); const exl3_build = @import("zig/build/exl3.zig"); const exlr8_build = @import("zig/build/exlr8.zig"); const golden_build = @import("zig/build/golden.zig"); +const video_build = @import("zig/build/video.zig"); comptime { const required = std.mem.trim(u8, @embedFile(".zig-version"), "\r\n"); @@ -167,6 +168,7 @@ pub fn build(b: *std.Build) void { } hip_build.steps(b, target, test_step); exlr8_build.hostTests(b, target, test_step); + video_build.steps(b, target); const lanes_cpu = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = b.graph.host, .optimize = .debug, .link_libc = true }); const disk_cpu = b.addTest(.{ .root_module = b.createModule(.{ .root_source_file = b.path("zig/src/glm_disk_tests.zig"), @@ -411,6 +413,8 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .{ .name = "tf-grid-sync-bench", .path = "zig/tests/grid_sync_bench.zig", .about = "A GPU-wide barrier in one persistent dispatch against dependent relaunches" }, .{ .name = "tf-weight-read-check", .path = "zig/tests/weight_read_check.zig", .about = "Check native file reads and failed-read cleanup", .c_source = "zig/tests/pread_fault.c" }, .{ .name = "tf-glm-run", .path = "zig/tests/glm_run.zig", .about = "GLM-5.3-Flash greedy replies at each draft depth against depth 0 and reference tokens" }, + .{ .name = "tf-glm-vision", .path = "zig/tests/glm_vision.zig", .about = "GLM-5.3-Flash's image tower on an image or a reference dump's patches, compared with the dump" }, + .{ .name = "tf-glm-video", .path = "zig/tests/glm_video.zig", .about = "GLM-5.3-Flash's video preparation on a clip (libtfvideo), written for tools/zig/glm_video_fixtures.py" }, .{ .name = "tf-glm-attn-probe", .path = "zig/tests/glm_attn_probe.zig", .about = "GLM-5.3-Flash's dense latent attention kernels on given inputs (scores, probabilities, outputs)" }, .{ .name = "tf-moe-bench", .path = "zig/tests/moe_bench.zig", .about = "Routed-expert prompt kernels from .metal files on a skewed routing: chained speed, bits against the first" }, .{ .name = "tf-mm-bench", .path = "zig/tests/mm_bench.zig", .about = "Prompt matmul kernels from .metal files: chained speed at one shape, bits against the first, error against fp64" }, @@ -422,7 +426,7 @@ fn metalTargets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.b .target = target, .optimize = optimize, .link_libc = true, - .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "cluster", .module = cluster }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "native_engines", .module = engines }, exlr8_build.cliImport(b, target, optimize, engine) }, + .imports = &.{ .{ .name = "metal", .module = metal }, .{ .name = "tensorfold", .module = engine }, .{ .name = "cluster", .module = cluster }, .{ .name = "kernel_sources", .module = sources }, .{ .name = "native_engines", .module = engines }, .{ .name = "engine_api", .module = api }, exlr8_build.cliImport(b, target, optimize, engine) }, }); if (p.c_source) |file| { mod.addCSourceFile(.{ .file = b.path(file), .flags = &.{"-std=c11"} }); @@ -549,6 +553,8 @@ fn metalEngineModules(b: *std.Build, target: std.Build.ResolvedTarget, optimize: metal.linkFramework("Metal", .{}); metal.linkFramework("Foundation", .{}); metal.linkFramework("IOSurface", .{}); // the Neural Engine's prompt share passes IOSurfaces + metal.linkFramework("ImageIO", .{}); // GLM-5.3-Flash's image input: decoding (families/glm/image.zig) + metal.linkFramework("CoreGraphics", .{}); metal.linkSystemLibrary("objc", .{}); const lanes = b.createModule(.{ .root_source_file = b.path("zig/src/core/lanes/lanes.zig"), .target = target, .optimize = optimize, .link_libc = true }); const fabric = b.createModule(.{ .root_source_file = b.path("zig/src/fabric/fabric.zig"), .target = target, .optimize = optimize, .link_libc = true }); diff --git a/docs/api.md b/docs/api.md index a0563b0048..058aee52fe 100644 --- a/docs/api.md +++ b/docs/api.md @@ -66,9 +66,12 @@ take a string `prompt` (`add_special_tokens`, default false) or a list of token chat fields that shape the prompt (`tools`, `reasoning_effort`, `chat_template_kwargs`) and `add_generation_prompt` (default true). It returns `count`, `max_model_len` (the context window; null when none is set on MLX) and `tokens`, and `token_strs` with `return_token_strs: true`. The IDs are the ones the chat route runs for the same -request, images expanded. `/detokenize` takes `tokens` and returns `prompt`. +request, images and videos expanded. `/detokenize` takes `tokens` and returns `prompt`. In the Python 0.6.6 engine, `--vision` lets supported Qwen3.5/3.8 dense checkpoints accept `image_url` content parts -alongside text in user messages and in tool results (`role: "tool"`), such as an agent's screenshots. +alongside text in user messages and in tool results (`role: "tool"`), such as an agent's screenshots. The native +server reads images for GLM-5.3-Flash on Metal, with data URLs, and public HTTPS URLs under `--vision-urls` +([image input](vision.md#the-native-server)); it also takes `video_url` parts in user messages +([video](vision.md#video-on-the-native-server)). See [image input](vision.md) for data URLs, public image URLs, limits and cache behavior. The native binary refuses `--vision` and answers image parts with HTTP 400. Unsupported image input, audio, video and non-text output requests receive HTTP 400. diff --git a/docs/vision.md b/docs/vision.md index 9536f7e7fc..c719fd999a 100644 --- a/docs/vision.md +++ b/docs/vision.md @@ -142,6 +142,63 @@ This prevents identical image-placeholder token IDs from reusing another image's Multi-turn image conversations work when the request includes the original image content parts, but image-prefix reuse and persisted image KV are not implemented. For Qwen, each image request carries its own multimodal rotary positions and continuation offset, including during concurrent lane rounds. GLM uses its native KDA/NoPE attention state. +## The native server + +`tensorfold-native serve` takes `--vision`, `--vision-max-images` (default 4) and `--vision-image-tokens` (default +4,096) for engines that read images; GLM-5.3-Flash on one Mac reads them so far, and other native engines refuse the +flag at startup. + +- Images arrive as `image_url` data URLs in user messages and tool results. With `--vision-urls` (it needs + `--vision`), a part may name a public HTTPS URL instead, fetched as 0.6.6 fetched it: port 443 without credentials + or a fragment; every address the host resolves to public (no private, loopback, link-local, carrier-grade NAT, + documentation, multicast or metadata address, and for IPv6 global unicast only, without 6to4 or Teredo); TLS + verified for the host against the system's certificates, on the address that was checked; up to 3 redirects, + each checked again; a declared `image/jpeg`, `image/png` or `image/webp`, uncompressed; 10 MB an image and 20 MB + a request across data URLs and downloads; 10 s a download and 30 s a request's downloads. The refusals are + 0.6.6's words; a TLS failure, like any other failed connection, is "image download failed or timed out". Host + names that aren't ASCII and IPv6 address literals are refused. Without the flag, an HTTPS URL is refused with the + flag's name. +- PNG and JPEG are decoded by the platform's decoder (ImageIO on macOS). WebP is refused by name until a decoder + exists; other formats are refused. +- The limits and the processing are 0.6.6's: EXIF orientation, alpha composited over white, the GLM5-Next processor's + sizing and PIL's bicubic resampler. Six drawn PNG fixtures give patches bit-equal to 0.6.6 with mlx-vlm + (`tools/zig/glm_vision_fixtures.py`). JPEG patches differ slightly, because ImageIO and libjpeg-turbo decode + differently. +- Images are counted from the request's image parts. A user or tool message whose text spells the model's image + markers (`<|image|>`, `<|begin_of_image|>`, `<|end_of_image|>`) is tokenized as ordinary text, so only the chat + template's own markers, one wrapper for each image part, become image tokens. +- The prompt cache keys each image's placeholder rows by the image's content, so a kept state resumes the same image + and never another one under the same placeholders. Image conversations reuse their history like text ones. + +### Video on the native server + +With `--vision`, GLM-5.3-Flash also reads `video_url` parts in user messages, as transformers 5.19's +`Glm5NextProcessor` prepares them: +- **Sampling:** 2 frames a second, deduplicated and padded to an even count, at most 256 frames. +- **Frame pairs:** each pair of frames is one image to the tower. +- **Prompt:** the template's `<|begin_of_video|><|video|><|end_of_video|>` holds one + `<|begin_of_image|>…<|end_of_image|>T.T seconds` block a pair. +- **Budget:** a request's videos share `--vision-video-tokens` visual tokens (default 16,384; the processor's + `max_image_tokens`). The frames are sized under it and resized with torch's antialiased uint8 bicubic. + +Video needs `libtfvideo`, an optional library over FFmpeg 9.0.2 (the release PyAV 19.0.1 ships). The server loads it +when the first video arrives and refuses videos by name without it: + +```sh +tools/zig/build_ffmpeg.sh ~/ffmpeg-9.0.2 # static FFmpeg: H.264, HEVC, VP8, VP9, MPEG-4; MP4/MOV, MKV/WebM +zig build video -Dffmpeg=$HOME/ffmpeg-9.0.2 # zig-out/native/lib/libtfvideo.dylib, beside the server +``` + +- **Limits (0.6.6's):** 2 videos a request, 16 MB each and 20 MB together, an hour of footage, 8192 pixels a side. + Data URLs work, and with `--vision-urls` public HTTPS URLs too (MP4, WebM, QuickTime or Matroska; 20 s a download, + 120 s a request). +- **Too short:** a clip too short for one sampled frame (under half a second) is refused, as the processor samples + nothing. +- **Frame counts:** a container without one (Matroska, WebM) is counted by its video packets. transformers' PyAV loader + can't read these at all. +- **Checks:** `tools/zig/glm_video_fixtures.py` checks the patches bit for bit against the processor on frames PyAV + decodes, over 13 clips. `tools/zig/video_served.py` checks a served model's prompt ids against the processor's. + ## Verification Compare image requests with `draft: true` and `draft: false` at identical sampling settings and seed, then compare concurrent requests with their solo results. diff --git a/tools/zig/build_ffmpeg.sh b/tools/zig/build_ffmpeg.sh new file mode 100755 index 0000000000..fd58a371f5 --- /dev/null +++ b/tools/zig/build_ffmpeg.sh @@ -0,0 +1,26 @@ +#!/usr/bin/env bash +# FFmpeg 9.0.2 (PyAV 19.0.1's) for libtfvideo, static PIC, PyAV's decoders only: build_ffmpeg.sh PREFIX [SOURCE_DIR] +set -euo pipefail +version=9.0.2 +sha256=8c3850283eb25fa026482078a04051e0be17347b09ef81a0849bec15a96e002e +prefix="$(mkdir -p "$1" && cd "$1" && pwd)" +work="${2:-$prefix/src}" +mkdir -p "$work" +cd "$work" +if [ ! -d "ffmpeg-$version" ]; then + [ -f "ffmpeg-$version.tar.xz" ] || curl -sSfLO "https://ffmpeg.org/releases/ffmpeg-$version.tar.xz" + echo "$sha256 ffmpeg-$version.tar.xz" | shasum -a 256 -c - + tar xf "ffmpeg-$version.tar.xz" +fi +cd "ffmpeg-$version" +./configure --prefix="$prefix" --enable-static --disable-shared --enable-pic \ + --disable-programs --disable-doc --disable-network --disable-autodetect --disable-everything \ + --disable-avdevice --disable-avfilter --disable-swresample \ + --enable-avcodec --enable-avformat --enable-swscale \ + --enable-decoder=h264,hevc,vp8,vp9,mpeg4 \ + --enable-parser=h264,hevc,vp8,vp9,mpeg4video \ + --enable-demuxer=mov,matroska \ + --enable-bsf=vp9_superframe_split +make -j"$(getconf _NPROCESSORS_ONLN)" +make install +echo "FFmpeg $version in $prefix" diff --git a/tools/zig/glm_video_fixtures.py b/tools/zig/glm_video_fixtures.py new file mode 100644 index 0000000000..3118237351 --- /dev/null +++ b/tools/zig/glm_video_fixtures.py @@ -0,0 +1,154 @@ +"""GLM-5.3-Flash video input against transformers' processor, on a fixed clip set (#565's video gate).""" + +from __future__ import annotations + +import argparse +import json +import re +import subprocess +import sys +from fractions import Fraction +from pathlib import Path + +import av +import numpy as np + + +def scene(i: int, w: int, h: int) -> np.ndarray: + y, x = np.mgrid[0:h, 0:w] + f = np.empty((h, w, 3), np.uint8) + f[..., 0] = (x * 255 // max(w - 1, 1)).astype(np.uint8) + f[..., 1] = ((y + 3 * i) % 256).astype(np.uint8) + f[..., 2] = ((x + y + 7 * i) % 200 + 30).astype(np.uint8) + s = max(8, min(w, h) // 4) + cx, cy = (5 + 4 * i) % max(w - s, 1), (h - s) // 2 + f[cy:cy + s, cx:cx + s] = (250, 240, 20) + return f + + +CLIPS = [ # name, container, codec, pixel format, size, rate, frames, options + ("h264.mp4", "mp4", "libx264", "yuv420p", (320, 180), 25, 80, {}), + ("h264-odd.mp4", "mp4", "libx264", "yuv420p", (318, 182), 30, 45, {}), + ("h264-full.mp4", "mp4", "libx264", "yuvj420p", (256, 144), 24, 30, {}), + ("h264-709.mov", "mov", "libx264", "yuv420p", (640, 360), Fraction(30000, 1001), 60, {"bt709": "1"}), + ("h264-444.mp4", "mp4", "libx264", "yuv444p", (200, 120), 25, 25, {}), + ("h264-hd.mp4", "mp4", "libx264", "yuv420p", (1280, 720), 30, 90, {}), + ("hevc.mp4", "mp4", "libx265", "yuv420p", (320, 240), 25, 50, {"x265-params": "log-level=error"}), + ("hevc10.mov", "mov", "libx265", "yuv420p10le", (320, 240), 25, 26, {"x265-params": "log-level=error"}), + ("vp9.webm", "webm", "libvpx-vp9", "yuv420p", (320, 180), 25, 50, {}), + ("vp8.mkv", "matroska", "libvpx", "yuv420p", (240, 160), 20, 40, {}), + ("mpeg4.mp4", "mp4", "mpeg4", "yuv420p", (352, 288), 25, 25, {}), + ("short.mp4", "mp4", "libx264", "yuv420p", (160, 120), 25, 10, {}), + ("long.mp4", "mp4", "libx264", "yuv420p", (64, 48), 10, 1500, {}), +] + + +def make(folder: Path) -> None: + folder.mkdir(parents=True, exist_ok=True) + for name, fmt, codec, pix, (w, h), rate, n, opts in CLIPS: + opts = dict(opts) + bt709 = opts.pop("bt709", None) + c = av.open(str(folder / name), "w", format=fmt) + s = c.add_stream(codec, rate=rate, options=opts) + s.width, s.height, s.pix_fmt = w, h, pix + if bt709: + s.codec_context.colorspace = s.codec_context.color_primaries = s.codec_context.color_trc = 1 + for i in range(n): + for p in s.encode(av.VideoFrame.from_ndarray(scene(i, w, h), format="rgb24").reformat(format=pix)): + c.mux(p) + for p in s.encode(): + c.mux(p) + c.close() + + +def load(path: Path, vp, max_frames: int): + """read_video_pyav, with a counted frame total where the container keeps none.""" + from transformers.video_utils import VideoMetadata + + with av.open(str(path)) as c: + st = c.streams.video[0] + total = int(st.frames) + if total == 0: + total = sum(1 for p in c.demux(st) if p.size) + c.seek(0) + fps = st.average_rate + meta = VideoMetadata(total_num_frames=total, fps=float(fps), duration=float(total / fps) if fps else 0.0, + video_backend="pyav", height=st.height, width=st.width) + vp.max_frames = max_frames + indices = vp.sample_frames(meta) + if len(indices) == 0: + return None, meta + frames = [] + c.seek(0) + for i, frame in enumerate(c.decode(video=0)): + if i > indices[-1]: + break + if i in indices: + frames.append(frame) + meta.frames_indices = indices + return np.stack([f.to_ndarray(format="rgb24") for f in frames]), meta + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--model", required=True) + ap.add_argument("--clips", required=True) + ap.add_argument("--native", default="zig-out/bin/tf-glm-video") + ap.add_argument("--max-tokens", type=int, default=16384) + ap.add_argument("--max-frames", type=int, default=256) + ap.add_argument("--make", action="store_true") + args = ap.parse_args() + from transformers import AutoProcessor + + clips = Path(args.clips) + if args.make: + make(clips) + proc = AutoProcessor.from_pretrained(args.model) + vp = proc.video_processor + failures = 0 + for path in sorted(p for p in clips.iterdir() if p.suffix in {".mp4", ".mov", ".webm", ".mkv"}): + frames, meta = load(path, vp, args.max_frames) + out = clips / "out" / path.stem + out.mkdir(parents=True, exist_ok=True) + run = subprocess.run([args.native, args.model, str(path), str(out), str(args.max_tokens), str(args.max_frames)], + capture_output=True, text=True) + if frames is None: + ok = run.returncode != 0 and "VideoTooShort" in run.stderr + failures += not ok + print(f"{'ok ' if ok else 'FAIL'} {path.name}: the processor samples no frames; native {'refuses' if run.returncode else 'answers'}") + continue + if run.returncode != 0: + failures += 1 + print(f"FAIL {path.name}: tf-glm-video: {run.stderr.strip()[-300:]}") + continue + native = json.loads((out / "native.json").read_text()) + msgs = [{"role": "user", "content": [{"type": "video"}, {"type": "text", "text": "What moves?"}]}] + text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) + enc = proc(text=[text], videos=[frames], video_metadata=[meta], do_sample_frames=False, + max_image_tokens=args.max_tokens, return_tensors="np") + ref = enc["pixel_values_videos"].astype(np.float32) + grid = [int(v) for v in enc["video_grid_thw"][0]] + mine = np.fromfile(out / "native_patches.f32", dtype=np.float32).reshape(-1, ref.shape[1]) if ref.ndim == 2 else None + decoded = sorted(set(int(i) for i in meta.frames_indices)) + prompt = proc.tokenizer.decode(enc["input_ids"][0]) + texts = re.findall(r"<\|end_of_image\|>([^<]*)", prompt) + same = mine is not None and mine.shape == ref.shape and np.array_equal(mine.view(np.uint32), ref.view(np.uint32)) + diff = 0 if mine is None or mine.shape != ref.shape else int((mine.view(np.uint32) != ref.view(np.uint32)).sum()) + checks = { + "indices": native["indices"] == decoded, + "grid": native["grid"] == grid, + "patches": same, + "texts": native["texts"] == texts, + } + ok = all(checks.values()) + failures += not ok + print(f"{'ok ' if ok else 'FAIL'} {path.name}: frames {meta.total_num_frames}{' (counted)' if native['counted'] else ''} at {float(meta.fps):.3f}, " + f"{len(decoded)} decoded, grid {grid}, {ref.size} values {'bit-equal' if same else f'({diff} differ)'}; " + f"times {texts[0]}..{texts[-1]}; {native['ms']} ms" + + ("" if ok else f" {[k for k, v in checks.items() if not v]} native {native['indices'][:8]} {native['grid']} {native['texts'][:3]}")) + print("PASS" if failures == 0 else f"FAIL: {failures}") + return 0 if failures == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/zig/glm_vision_fixtures.py b/tools/zig/glm_vision_fixtures.py new file mode 100644 index 0000000000..72e222fe93 --- /dev/null +++ b/tools/zig/glm_vision_fixtures.py @@ -0,0 +1,151 @@ +"""GLM-5.3-Flash image input against the Python frontend, on a fixed image set (#565's gate).""" + +from __future__ import annotations + +import argparse +import io +import json +import re +import subprocess +import sys +from pathlib import Path + +import mlx.core as mx +import numpy as np +from PIL import Image, ImageDraw + + +def draw_shapes(size: tuple[int, int]) -> Image.Image: + w, h = size + im = Image.new("RGB", size, (245, 240, 230)) + d = ImageDraw.Draw(im) + d.ellipse((w * 0.08, h * 0.25, w * 0.42, h * 0.7), fill=(210, 40, 40)) + d.rectangle((w * 0.55, h * 0.25, w * 0.9, h * 0.7), fill=(40, 70, 200)) + for x in range(int(w * 0.1), int(w * 0.9), 6): + d.line((x, h * 0.8, x, h * 0.92), fill=(20, 20, 20), width=2) + d.text((w * 0.4, h * 0.06), "GLM 7", fill=(0, 0, 0)) + return im + + +def fixtures(folder: Path) -> list[Path]: + folder.mkdir(parents=True, exist_ok=True) + out = [] + + def save(name: str, im: Image.Image, **kw) -> None: + p = folder / name + im.save(p, **kw) + out.append(p) + + save("shapes.png", draw_shapes((517, 389))) # padding only + big = draw_shapes((2000, 1500)) + px = big.load() + for y in range(0, 1500, 3): # fine detail for the antialias + for x in range(0, 2000, 7): + px[x, y] = ((x * 3) % 256, (y * 5) % 256, (x + y) % 256) + save("large.png", big) # antialiased downscale + save("tiny.png", draw_shapes((20, 15))) # upscale to the token floor + logo = Image.new("RGBA", (96, 96), (0, 0, 0, 0)) + d = ImageDraw.Draw(logo) + d.rounded_rectangle((4, 4, 92, 92), 16, fill=(30, 30, 40, 255)) + d.polygon([(26, 28), (70, 28), (30, 68), (70, 68), (70, 74), (24, 74), (64, 34), (26, 34)], fill=(255, 255, 255, 230)) + save("logo.png", logo) # RGBA upscale, alpha over white + pal = draw_shapes((517, 389)).quantize(32) + pal.info["transparency"] = 0 + save("palette.png", pal, transparency=0) # palette, a transparent entry + la = Image.new("LA", (240, 180), (128, 0)) + d = ImageDraw.Draw(la) + for r in range(80, 0, -8): + d.ellipse((120 - r, 90 - r, 120 + r, 90 + r), fill=(255 - 3 * r, 255 - 2 * r)) + save("gray-alpha.png", la) # gray with alpha + rot = draw_shapes((300, 200)) + exif = Image.Exif() + exif[0x0112] = 6 # rotate 90 degrees clockwise + save("rotated.jpg", rot, quality=90, exif=exif.tobytes()) # EXIF orientation + save("shapes.jpg", draw_shapes((517, 389)), quality=85) # baseline JPEG + return out + + +def reference(model: Path, image_path: Path, out: Path, tower, processor, max_tokens: int) -> dict: + from tensorfold.vision.images import DEFAULT_LIMITS, _decode + + out.mkdir(parents=True, exist_ok=True) + image = _decode(image_path.read_bytes(), "auto", DEFAULT_LIMITS, DEFAULT_LIMITS.max_total_pixels).to_pil() + done = processor([image], return_tensors="np", min_image_tokens=min(16, max_tokens), max_image_tokens=max_tokens) + patches = np.asarray(done["pixel_values"], dtype=np.float32) + grid = np.asarray(done["image_grid_thw"], dtype=np.int64) + dtype = tower.patch_embed.proj.weight.dtype + x = mx.array(patches).astype(dtype) + g = mx.array(grid, dtype=mx.int32) + h = tower.patch_embed(x) + embed = h + pos = tower._rotary_embeddings(g) + cu = mx.array([0, int(h.shape[0])]) + block0 = None + for i, block in enumerate(tower.blocks): + h = block(h, cu, pos) + if i == 0: + block0 = h + blocks = h + features = tower(x, g) + mx.eval(embed, block0, blocks, features) + for name, a in (("patches", patches), ("embed", embed), ("block0", block0), ("blocks", blocks), ("features", features)): + np.asarray(a.astype(mx.float32) if isinstance(a, mx.array) else a, dtype=np.float32).tofile(out / f"{name}.f32") + meta = {"image": str(image_path), "size": list(image.size), "mode": Image.open(image_path).mode, "grid_thw": grid[0].tolist(), + "max_tokens": max_tokens, "tokens": int(features.shape[0]), "patch_width": int(patches.shape[1]), "dtype": str(dtype)} + (out / "meta.json").write_text(json.dumps(meta, indent=1) + "\n") + return meta + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--model", required=True, help="the GLM-5.3-Flash checkpoint folder") + ap.add_argument("--out", required=True, help="where the images and their references go") + ap.add_argument("--max-tokens", type=int, default=4096, help="the visual-token budget an image is sized to") + ap.add_argument("--tower", help="tf-glm-vision: also compare the native patches and tower with each reference") + args = ap.parse_args() + + from mlx_vlm.models.glm5_next.config import VisionConfig + from mlx_vlm.models.glm5_next.processing import Glm5NextImageProcessor + from mlx_vlm.models.glm5_next.vision import VisionModel + + model, out = Path(args.model), Path(args.out) + config = json.loads((model / "config.json").read_text()) + proc_cfg = json.loads((model / "processor_config.json").read_text())["image_processor"] + proc_cfg.pop("image_processor_type", None) + processor = Glm5NextImageProcessor(**proc_cfg) + tower = VisionModel(VisionConfig.from_dict(config["vision_config"])) + index = json.loads((model / "model.safetensors.index.json").read_text())["weight_map"] + weights = {} + for shard in sorted({f for k, f in index.items() if k.startswith("model.visual.")}): + for k, v in mx.load(str(model / shard)).items(): + if k.startswith("model.visual."): + weights[k[len("model.visual."):]] = v + tower.load_weights(list(tower.sanitize(weights).items()), strict=True) + mx.eval(tower.parameters()) + + failures = 0 + for image in fixtures(out / "images"): + ref = out / image.name.replace(".", "-") + meta = reference(model, image, ref, tower, processor, args.max_tokens) + line = f"{image.name:15} {meta['size'][0]}x{meta['size'][1]} {meta['mode']:5} grid {meta['grid_thw']} {meta['tokens']} tokens" + if args.tower: + run = subprocess.run([args.tower, str(model), str(ref), "features", str(image)], capture_output=True, text=True) + text = run.stdout + run.stderr + p = re.search(r"patches: (\d+) of (\d+) values differ", text) + c = re.search(r"cosine mean ([0-9.]+) worst ([0-9.]+)", text) + if run.returncode != 0 or not p or not c: + failures += 1 + line += f" FAIL: {text.strip().splitlines()[-1] if text.strip() else run.returncode}" + else: + differ = int(p.group(1)) + exact = differ == 0 + # PNG patches must be bit-exact; JPEG patches differ by decoder (ImageIO against libjpeg-turbo) + failures += (not exact) and image.suffix == ".png" + line += f" patches {'equal' if exact else f'{differ} of {p.group(2)} differ'}, rows' cosine {c.group(1)} (worst {c.group(2)})" + print(line) + print("PASS" if failures == 0 else f"FAIL: {failures}") + return 0 if failures == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/zig/video_served.py b/tools/zig/video_served.py new file mode 100644 index 0000000000..87d012b7fc --- /dev/null +++ b/tools/zig/video_served.py @@ -0,0 +1,151 @@ +"""Video input served end to end on a running ``tensorfold-native --vision`` (#565's video checks).""" + +from __future__ import annotations + +import argparse +import base64 +import concurrent.futures as cf +import json +import sys +import urllib.error +import urllib.request +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from glm_video_fixtures import load # noqa: E402 + + +def post(url: str, path: str, body: dict) -> tuple[int, dict]: + req = urllib.request.Request(url + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=1800) as r: + return r.status, json.loads(r.read()) + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read() or b"{}") + + +def data_url(raw: bytes, mime: str) -> str: + return f"data:{mime};base64,{base64.b64encode(raw).decode()}" + + +MIME = {".mp4": "video/mp4", ".mov": "video/quicktime", ".webm": "video/webm", ".mkv": "video/x-matroska"} + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--url", default="http://127.0.0.1:8000") + ap.add_argument("--model", required=True) + ap.add_argument("--clips", required=True) + ap.add_argument("--image", help="a PNG for the video-and-image request") + ap.add_argument("--video-tokens", type=int, default=16384) + ap.add_argument("--tokens", type=int, default=32) + args = ap.parse_args() + from transformers import AutoProcessor + + proc = AutoProcessor.from_pretrained(args.model) + vp = proc.video_processor + with urllib.request.urlopen(args.url + "/v1/models") as r: + model = json.loads(r.read())["data"][0]["id"] + clips = Path(args.clips) + failures = 0 + + def check(label: str, ok: bool, detail: str = "") -> None: + nonlocal failures + failures += not ok + print(f"{'ok ' if ok else 'FAIL'} {label}{': ' + detail if detail else ''}") + + def video(p: Path) -> dict: + return {"type": "video_url", "video_url": {"url": data_url(p.read_bytes(), MIME[p.suffix])}} + + def text(t: str) -> dict: + return {"type": "text", "text": t} + + def user(*parts) -> dict: + return {"role": "user", "content": list(parts)} + + off = {"enable_thinking": False} + + def ask(messages: list, draft: bool = True, tokens: int | None = None) -> tuple[int, dict]: + b = {"model": model, "messages": messages, "max_tokens": tokens or args.tokens, "temperature": 0, "chat_template_kwargs": off} + if not draft: + b["draft"] = False + return post(args.url, "/v1/chat/completions", b) + + def reply(r: dict) -> str: + return r["choices"][0]["message"].get("content") or "" + + # the prompt ids per clip against the processor's (thinking on: no template switch; off, the server adds ) + q = "What moves?" + same = 0 + clip_files = sorted(p for p in clips.iterdir() if p.suffix in MIME and p.name != "short.mp4") + for p in clip_files: + frames, meta = load(p, vp, 256) + msgs = [{"role": "user", "content": [{"type": "video"}, {"type": "text", "text": q}]}] + prompt = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True) # the default: thinking on + ref = proc(text=[prompt], videos=[frames], video_metadata=[meta], do_sample_frames=False, max_image_tokens=args.video_tokens, return_tensors="np")["input_ids"][0].tolist() + code, r = post(args.url, "/tokenize", {"messages": [user(video(p), text(q))]}) + ok = code == 200 and r.get("tokens") == ref + same += ok + if not ok: + got = r.get("tokens", []) + at = next((i for i, (x, y) in enumerate(zip(got, ref)) if x != y), min(len(got), len(ref))) + print(f" {p.name}: HTTP {code} {len(got)} ids against {len(ref)}, first difference at {at}: {got[at:at + 6]} vs {ref[at:at + 6]} {r.get('error', '')}") + check("a video request's prompt ids equal the processor's", same == len(clip_files), f"{same} of {len(clip_files)} clips") + + a, b = clips / "h264.mp4", clips / "hevc.mp4" + # the cache, before any concurrent run (a slot another stream takes leaves GLM's resident state stale) + system = {"role": "system", "content": " ".join(f"Rule {i}: describe what you see plainly and briefly." for i in range(120))} + head = ask([system, user(text("What is shown?"))], tokens=1)[1]["usage"]["prompt_tokens"] + long = " ".join(f"Detail {i}: the clip shows shapes, colours and motion." for i in range(40)) + + def turns(c: Path) -> tuple[list, list]: + two = [system, user(video(c), text("What is shown?")), {"role": "assistant", "content": long}, user(text("Again, in three words."))] + return two, two + [{"role": "assistant", "content": "Three words here."}, user(text("Once more."))] + + two_a, three_a = turns(a) + _, three_b = turns(clips / "h264-odd.mp4") + ask(two_a, tokens=8) + code, ra = ask(three_a, tokens=16) + cached_a = ra.get("usage", {}).get("prompt_tokens_details", {}).get("cached_tokens", 0) if code == 200 else -1 + check("a video conversation's later turn resumes past its video", code == 200 and cached_a > head, f"cached {cached_a} of {ra.get('usage', {}).get('prompt_tokens')}; the text alone is {head}") + code, rb = ask(three_b, tokens=16) + cached_b = rb.get("usage", {}).get("prompt_tokens_details", {}).get("cached_tokens", 0) if code == 200 else -1 + check("another video under the same text resumes only the text before it", code == 200 and cached_b < head, f"cached {cached_b} of {rb.get('usage', {}).get('prompt_tokens')}") + + convo = [user(video(a), text("Describe the video in one sentence."))] + code, r = ask(convo) + code2, r2 = ask(convo, draft=False) + check("a video request answers", code == 200 and code2 == 200, f"HTTP {code}/{code2} {r.get('error', '')}") + if code == 200 and code2 == 200: + check("drafted equals \"draft\": false with a video", reply(r) == reply(r2), repr(reply(r)[:60])) + runs = [[user(video(c), text(t))] for c in (a, b) for t in ("What colours appear?", "Does anything move?")] + alone = [reply(ask(m)[1]) for m in runs] + with cf.ThreadPoolExecutor(len(runs)) as pool: + together = [reply(f.result()[1]) for f in [pool.submit(ask, m) for m in runs]] + check("concurrent video requests equal their runs alone", alone == together, f"{sum(x == y for x, y in zip(alone, together))} of {len(runs)}") + if args.image: + img = {"type": "image_url", "image_url": {"url": data_url(Path(args.image).read_bytes(), "image/png")}} + code, r = ask([user(img, video(a), text("Does the picture appear in the video?"))]) + check("a video and an image in one request answer", code == 200, f"HTTP {code} {r.get('error', '')}") + + refusals = [ + ("three videos", [user(video(a), video(a), video(a), text("hi"))], "at most 2 videos"), + ("a video in an assistant message", [user(text("hi")), {"role": "assistant", "content": [video(a)]}, user(text("ok"))], "only in user messages"), + ("a video in a tool result", [user(text("hi")), {"role": "assistant", "content": "", "tool_calls": [{"id": "c", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "c", "content": [video(a)]}], "only in user messages"), + ("a clip too short to sample", [user(video(clips / "short.mp4"), text("hi"))], "too short"), + ("bytes that aren't a video", [user({"type": "video_url", "video_url": {"url": data_url(b"not a video at all" * 20, "video/mp4")}}, text("hi"))], "invalid or unsupported"), + ("a URL without --vision-urls", [user({"type": "video_url", "video_url": {"url": "https://example.com/a.mp4"}}, text("hi"))], "--vision-urls"), + ("another scheme", [user({"type": "video_url", "video_url": {"url": "ftp://example.com/a.mp4"}}, text("hi"))], "data URLs or public HTTPS URLs"), + ("a video_url without a url", [user({"type": "video_url", "video_url": {}}, text("hi"))], "non-empty url"), + ] + for label, messages, words in refusals: + code, r = ask(messages) + message = r.get("error", {}).get("message", json.dumps(r)) + check(f"refused {label}", code == 400 and words in message, f"HTTP {code} {message[:110]!r}") + print("PASS" if failures == 0 else f"FAIL: {failures}") + return 0 if failures == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/zig/vision_served.py b/tools/zig/vision_served.py new file mode 100644 index 0000000000..8d9a645e7f --- /dev/null +++ b/tools/zig/vision_served.py @@ -0,0 +1,164 @@ +"""Image input served end to end on a running ``tensorfold-native --vision`` (#565's served checks).""" + +from __future__ import annotations + +import argparse +import base64 +import concurrent.futures as cf +import json +import sys +import urllib.error +import urllib.request +from pathlib import Path + + +def post(url: str, body: dict) -> tuple[int, dict]: + req = urllib.request.Request(url + "/v1/chat/completions", data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=1800) as r: + return r.status, json.loads(r.read()) + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read() or b"{}") + + +def data_url(raw: bytes, mime: str) -> str: + return f"data:{mime};base64,{base64.b64encode(raw).decode()}" + + +def image_part(url: str, detail: str | None = None) -> dict: + inner = {"url": url} + if detail: + inner["detail"] = detail + return {"type": "image_url", "image_url": inner} + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--url", default="http://127.0.0.1:8000") + ap.add_argument("--images", required=True) + ap.add_argument("--tokens", type=int, default=48) + ap.add_argument("--fresh-url", help="the same model served with --prompt-cache-gib 0: the cache check's reference reply") + args = ap.parse_args() + with urllib.request.urlopen(args.url + "/v1/models") as r: + model = json.loads(r.read())["data"][0]["id"] + folder = Path(args.images) + pngs = sorted(folder.glob("*.png")) + jpgs = sorted(folder.glob("*.jpg")) + first, second = pngs[0], pngs[1] + url_a = data_url(first.read_bytes(), "image/png") + url_b = data_url(second.read_bytes(), "image/png") + failures = 0 + + def ask(messages: list, draft: bool = True, tokens: int | None = None) -> tuple[int, dict]: + b = {"model": model, "messages": messages, "max_tokens": tokens or args.tokens, "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}} + if not draft: + b["draft"] = False + return post(args.url, b) + + def user(*parts) -> dict: + return {"role": "user", "content": list(parts)} + + def text(t: str) -> dict: + return {"type": "text", "text": t} + + def reply(r: dict) -> str: + return r["choices"][0]["message"].get("content") or "" + + def check(label: str, ok: bool, detail: str = "") -> None: + nonlocal failures + failures += not ok + print(f"{'ok ' if ok else 'FAIL'} {label}{': ' + detail if detail else ''}") + + # one image, drafted and plain; its usage counts the image's tokens + convo = [user(image_part(url_a), text("Describe this image in one sentence."))] + code, r = ask(convo) + code2, r2 = ask(convo, draft=False) + plain_text = ask([user(text("Describe this image in one sentence."))])[1] + check("an image request answers", code == 200 and code2 == 200, f"HTTP {code}/{code2}") + if code == 200 and code2 == 200: + check("drafted equals \"draft\": false with an image", reply(r) == reply(r2), repr(reply(r)[:60])) + check("the image's tokens are in the prompt", r["usage"]["prompt_tokens"] > plain_text["usage"]["prompt_tokens"] + 8, + f"{r['usage']['prompt_tokens']} prompt tokens with the image, {plain_text['usage']['prompt_tokens']} without") + + # #511: the marker quoted in a user's text and in a tool result, beside one real image + quoted = [user(image_part(url_a), text("A file I opened contains the string <|image|> and also <|begin_of_image|>. What is in the picture?")), + {"role": "assistant", "content": "", "tool_calls": [{"id": "c1", "type": "function", "function": {"name": "read", "arguments": "{\"path\": \"a.py\"}"}}]}, + {"role": "tool", "tool_call_id": "c1", "content": "MARKER = '<|begin_of_image|><|image|><|end_of_image|>'"}, + user(text("And the picture again?"))] + tools = [{"type": "function", "function": {"name": "read", "parameters": {"type": "object", "properties": {"path": {"type": "string"}}}}}] + code, r = post(args.url, {"model": model, "messages": quoted, "tools": tools, "max_tokens": args.tokens, "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}}) + check("quoted image markers beside one image are text (#511)", code == 200, f"HTTP {code} {r.get('error', {}).get('message', '')[:100]}") + + # a tool result carrying an image (a screenshot), rendered inside the tool response + shot = [user(text("Take a screenshot.")), + {"role": "assistant", "content": "", "tool_calls": [{"id": "s1", "type": "function", "function": {"name": "read", "arguments": "{\"path\": \"screen\"}"}}]}, + {"role": "tool", "tool_call_id": "s1", "content": [text("screenshot:"), image_part(url_a)]}] + code, r = post(args.url, {"model": model, "messages": shot, "tools": tools, "max_tokens": args.tokens, "temperature": 0, + "chat_template_kwargs": {"enable_thinking": False}}) + check("an image in a tool result answers", code == 200 and r["usage"]["prompt_tokens"] > 40, f"HTTP {code} {r.get('usage', r.get('error', {}).get('message', ''))}") + + # concurrent equals alone, images in each request + runs = [[user(image_part(u), text(q))] for u in (url_a, url_b) for q in ("What colours do you see?", "Count the objects.")] + alone = [reply(ask(m)[1]) for m in runs] + with cf.ThreadPoolExecutor(len(runs)) as pool: + together = [reply(f.result()[1]) for f in [pool.submit(ask, m) for m in runs]] + same = sum(x == y for x, y in zip(alone, together)) + check("concurrent image requests equal their runs alone", same == len(runs), f"{same} of {len(runs)}") + + # the prompt cache keeps state at planned cuts (system, a 256+ token reply): turn three resumes past turn 2's image + system = {"role": "system", "content": " ".join(f"Rule {i}: describe what you see plainly and briefly." for i in range(120))} + q, q2, q3 = "What is shown?", "Say that again in three words.", "And once more." + head = ask([system, user(text(q))], tokens=1)[1]["usage"]["prompt_tokens"] # the text alone: past where an image starts + first = " ".join(f"Detail {i}: the picture shows shapes, colours and lines." for i in range(40)) # a reply of 256+ tokens + + def turns(url: str) -> tuple[list, list]: + two = [system, user(image_part(url), text(q)), {"role": "assistant", "content": first}, user(text(q2))] + return two, two + [{"role": "assistant", "content": "Three words here."}, user(text(q3))] + + two_a, three_a = turns(url_a) + two_b, three_b = turns(url_b) + if args.fresh_url: # B's reply with nothing kept + rf = post(args.fresh_url, {"model": model, "messages": three_b, "max_tokens": 24, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}})[1] + ask(two_a, tokens=8) # keeps A's state at the long reply's start, past its image + code, ra = ask(three_a, tokens=24) + cached_a = ra.get("usage", {}).get("prompt_tokens_details", {}).get("cached_tokens", 0) if code == 200 else -1 + check("an image conversation's later turn resumes past its image", code == 200 and cached_a > head, f"cached {cached_a} of {ra.get('usage', {}).get('prompt_tokens')}; the text alone is {head}") + # image B under the same text and placeholder ids resumes at most the text before its image + code, rb = ask(three_b, tokens=24) + cached_b = rb.get("usage", {}).get("prompt_tokens_details", {}).get("cached_tokens", 0) if code == 200 else -1 + check("another image under the same placeholders resumes only the text before it", code == 200 and cached_b < head, + f"cached {cached_b} of {rb.get('usage', {}).get('prompt_tokens')}") + if args.fresh_url: + check("and answers as a server with nothing kept", reply(rf) == reply(rb), repr(reply(rb)[:60])) + + # refusals, before any token + webp = b"RIFF\x24\x00\x00\x00WEBPVP8 \x18\x00\x00\x00" + bytes(24) + gif = b"GIF89a\x01\x00\x01\x00\x80\x00\x00\xff\xff\xff\x00\x00\x00!\xf9\x04\x00\x00\x00\x00\x00,\x00\x00\x00\x00\x01\x00\x01\x00\x00\x02\x02D\x01\x00;" + refusals = [ + ("WebP, by name", [user(image_part(data_url(webp, "image/webp")), text("hi"))], "WebP"), + ("GIF", [user(image_part(data_url(gif, "image/gif")), text("hi"))], "not PNG or JPEG"), + ("an image in an assistant message", [user(text("hi")), {"role": "assistant", "content": [image_part(url_a)]}, user(text("ok"))], "user and tool messages only"), + ("too many images", [user(*[image_part(url_a)] * 5, text("hi"))], "at most"), + ("a remote URL without --vision-urls", [user(image_part("https://example.com/cat.png"), text("hi"))], "--vision-urls"), + ("a URL that is neither data nor https", [user(image_part("http://example.com/cat.png"), text("hi"))], "data URLs or public HTTPS URLs"), + ("bad base64", [user(image_part("data:image/png;base64,@@@@"), text("hi"))], "encoding"), + ("a bad detail", [user(image_part(url_a, "medium"), text("hi"))], "detail"), + ("bytes that decode as nothing", [user(image_part(data_url(b"\x89PNG\r\n\x1a\n" + bytes(32), "image/png")), text("hi"))], "could not be decoded"), + ] + for label, messages, words in refusals: + code, r = ask(messages) + message = r.get("error", {}).get("message", json.dumps(r)) + check(f"refused {label}", code == 400 and words in message, f"HTTP {code} {message[:100]!r}") + if jpgs: + code, r = ask([user(image_part(data_url(jpgs[0].read_bytes(), "image/jpeg")), text("Describe it."))]) + check("a JPEG answers", code == 200, f"HTTP {code}") + code, r = ask([user(image_part(url_a, "low"), text("Describe it."))]) + check("detail low answers", code == 200, f"HTTP {code}") + print("PASS" if failures == 0 else f"FAIL: {failures}") + return 0 if failures == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/zig/vision_urls_served.py b/tools/zig/vision_urls_served.py new file mode 100644 index 0000000000..3ab3c5773a --- /dev/null +++ b/tools/zig/vision_urls_served.py @@ -0,0 +1,99 @@ +"""Image URLs served end to end on a running ``tensorfold-native --vision --vision-urls`` (#565's URL checks).""" + +from __future__ import annotations + +import argparse +import base64 +import json +import ssl +import sys +import urllib.error +import urllib.request + +IMAGE = "https://raw.githubusercontent.com/ashhart/TensorFold/main/assets/tensorfold-hero.png" +REDIRECT = "https://github.com/ashhart/TensorFold/raw/main/assets/tensorfold-hero.png" # 302 to IMAGE's host +OTHER = "https://raw.githubusercontent.com/ashhart/TensorFold/main/assets/tensorfold-cuda-banner.png" + + +def post(url: str, body: dict) -> tuple[int, dict]: + req = urllib.request.Request(url + "/v1/chat/completions", data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + try: + with urllib.request.urlopen(req, timeout=1800) as r: + return r.status, json.loads(r.read()) + except urllib.error.HTTPError as e: + return e.code, json.loads(e.read() or b"{}") + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + ap.add_argument("--url", default="http://127.0.0.1:8000") + ap.add_argument("--tokens", type=int, default=48) + args = ap.parse_args() + with urllib.request.urlopen(args.url + "/v1/models") as r: + model = json.loads(r.read())["data"][0]["id"] + try: # python.org's macOS builds ship no CA store until "Install Certificates" runs + import certifi + tls = ssl.create_default_context(cafile=certifi.where()) + except ImportError: + tls = ssl.create_default_context() + with urllib.request.urlopen(IMAGE, timeout=30, context=tls) as r: + inline = "data:image/png;base64," + base64.b64encode(r.read()).decode() + failures = 0 + + def ask(*parts) -> tuple[int, dict]: + return post(args.url, {"model": model, "messages": [{"role": "user", "content": list(parts)}], "max_tokens": args.tokens, + "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}}) + + def image(url: str) -> dict: + return {"type": "image_url", "image_url": {"url": url}} + + def text(t: str) -> dict: + return {"type": "text", "text": t} + + def check(label: str, ok: bool, detail: str = "") -> None: + nonlocal failures + failures += not ok + print(f"{'ok ' if ok else 'FAIL'} {label}{': ' + detail if detail else ''}") + + def same(label: str, a: tuple[int, dict], b: tuple[int, dict]) -> None: + ok = a[0] == 200 and b[0] == 200 and a[1]["usage"]["prompt_tokens"] == b[1]["usage"]["prompt_tokens"] and \ + a[1]["choices"][0]["message"]["content"] == b[1]["choices"][0]["message"]["content"] + detail = f"HTTP {a[0]}/{b[0]}" if a[0] != 200 or b[0] != 200 else \ + f"{a[1]['usage']['prompt_tokens']} prompt tokens, {a[1]['choices'][0]['message']['content'][:60]!r}" + if a[0] != 200 or b[0] != 200: + detail += " " + json.dumps(a[1].get("error") or b[1].get("error"))[:120] + check(label, ok, detail) + + q = text("Describe this image in one sentence.") + by_data = ask(image(inline), q) + same("an HTTPS image URL answers as its bytes in a data URL", ask(image(IMAGE), q), by_data) + same("a URL that redirects to the image answers the same", ask(image(REDIRECT), q), by_data) + two = ask(image(OTHER), image(inline), text("How do these two images differ?")) + check("a URL and a data URL in one request answer", two[0] == 200, f"HTTP {two[0]} {json.dumps(two[1].get('error', ''))[:100]}") + + refusals = [ + ("an HTTP error", "https://raw.githubusercontent.com/ashhart/TensorFold/main/assets/no-such-image.png", "image download returned HTTP 404"), + ("a page that isn't an image", "https://example.com/", "image URL content type must be JPEG, PNG or WebP"), + ("a host that resolves to a private address", "https://localtest.me/x.png", "image URLs must resolve only to public internet addresses"), + ("a private address literal", "https://169.254.169.254/latest/meta-data", "image URLs must resolve only to public internet addresses"), + ("localhost", "https://localhost/x.png", "image URLs must use public internet hosts"), + ("cloud metadata by name", "https://metadata.google.internal/x.png", "image URLs must use public internet hosts"), + ("an expired certificate", "https://expired.badssl.com/", "image download failed or timed out"), + ("a certificate for another host", "https://wrong.host.badssl.com/", "image download failed or timed out"), + ("a self-signed certificate", "https://self-signed.badssl.com/", "image download failed or timed out"), + ("plain HTTP", "http://example.com/x.png", "images require data URLs or public HTTPS URLs"), + ("another port", "https://example.com:8443/x.png", "image URL must be HTTPS on port 443, without credentials or a fragment"), + ("credentials", "https://user:pw@example.com/x.png", "image URL must be HTTPS on port 443, without credentials or a fragment"), + ("a fragment", "https://example.com/x.png#a", "image URL must be HTTPS on port 443, without credentials or a fragment"), + ("a host that never resolves", "https://no-such-host.invalid/x.png", "image host could not be resolved"), + ] + for label, url, words in refusals: + code, r = ask(image(url), text("hi")) + message = r.get("error", {}).get("message", json.dumps(r)) + check(f"refused {label}", code == 400 and message == words, f"HTTP {code} {message[:100]!r}") + print("PASS" if failures == 0 else f"FAIL: {failures}") + return 0 if failures == 0 else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/zig/build/video.zig b/zig/build/video.zig new file mode 100644 index 0000000000..5594959832 --- /dev/null +++ b/zig/build/video.zig @@ -0,0 +1,22 @@ +//! libtfvideo: video input's optional runtime library (tf_video.c over the static FFmpeg build_ffmpeg.sh installs). + +const std = @import("std"); + +/// The FFmpeg release the decoder is written and checked against (its frames against PyAV's). +pub const pinned = "9.0.2"; + +pub fn steps(b: *std.Build, target: std.Build.ResolvedTarget) void { + const step = b.step("video", "libtfvideo (video input) over FFmpeg " ++ pinned ++ ": -Dffmpeg=PREFIX from tools/zig/build_ffmpeg.sh"); + const prefix = b.option([]const u8, "ffmpeg", "A static FFmpeg " ++ pinned ++ " install (tools/zig/build_ffmpeg.sh PREFIX), for `zig build video`") orelse { + step.dependOn(&b.addFail("zig build video needs -Dffmpeg=PREFIX: tools/zig/build_ffmpeg.sh PREFIX builds FFmpeg " ++ pinned ++ " there").step); + return; + }; + const root: std.Build.LazyPath = .{ .cwd_relative = prefix }; + const mod = b.createModule(.{ .target = target, .optimize = .ReleaseFast, .link_libc = true }); + mod.addIncludePath(root.path(b, "include")); + mod.addCSourceFile(.{ .file = b.path("zig/src/video/tf_video.c"), .flags = &.{ "-std=c11", "-Wall" } }); + for ([_][]const u8{ "libavformat.a", "libavcodec.a", "libswscale.a", "libavutil.a" }) |a| mod.addObjectFile(root.path(b, b.pathJoin(&.{ "lib", a }))); + mod.linkSystemLibrary("m", .{}); + const lib = b.addLibrary(.{ .name = "tfvideo", .linkage = .dynamic, .root_module = mod }); + step.dependOn(&b.addInstallArtifact(lib, .{ .dest_dir = .{ .override = .{ .custom = "native/lib" } } }).step); +} diff --git a/zig/kernels/metal/glm_vision.metal b/zig/kernels/metal/glm_vision.metal new file mode 100644 index 0000000000..cf79f43b03 --- /dev/null +++ b/zig/kernels/metal/glm_vision.metal @@ -0,0 +1,225 @@ +// GLM-5.3-Flash's image tower, MLX's per-op bf16 arithmetic as mlx-vlm runs it; needs prefill/gemm_nax.metal first. + +// ---- the projections and attention products: tfp::gemm_nax (bm64 bn128 bk256 wm2 wn4) by operand layout ---- +#define GLMV_GEMM(NAME, T, TA, TB, AM, AN) \ + [[max_total_threads_per_threadgroup(256)]] [[kernel]] void NAME( \ + const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], const device int32_t* P [[buffer(2)]], \ + device T* D [[buffer(3)]], uint sg [[simdgroup_index_in_threadgroup]], \ + uint lane [[thread_index_in_simdgroup]], uint3 tg [[threadgroup_position_in_grid]]) { \ + tfp::gemm_nax(A, B, D, P, tg, sg, lane); \ + } +// x [M, K] w^T: every linear layer (bf16 in and out, fp32 sums) +GLMV_GEMM(glmv_gemm_bf16_nt_00, bfloat16_t, false, true, false, false) +GLMV_GEMM(glmv_gemm_bf16_nt_01, bfloat16_t, false, true, false, true) +GLMV_GEMM(glmv_gemm_bf16_nt_10, bfloat16_t, false, true, true, false) +GLMV_GEMM(glmv_gemm_bf16_nt_11, bfloat16_t, false, true, true, true) +// q k^T per head in fp32 (the scores MLX's attention keeps in fp32) +GLMV_GEMM(glmv_gemm_f32_nt_00, float, false, true, false, false) +GLMV_GEMM(glmv_gemm_f32_nt_01, float, false, true, false, true) +GLMV_GEMM(glmv_gemm_f32_nt_10, float, false, true, true, false) +GLMV_GEMM(glmv_gemm_f32_nt_11, float, false, true, true, true) +// p v per head in fp32 +GLMV_GEMM(glmv_gemm_f32_nn_00, float, false, false, false, false) +GLMV_GEMM(glmv_gemm_f32_nn_01, float, false, false, false, true) +GLMV_GEMM(glmv_gemm_f32_nn_10, float, false, false, true, false) +GLMV_GEMM(glmv_gemm_f32_nn_11, float, false, false, true, true) + +struct GlmvRows { + uint rows; + uint width; + float eps; +}; + +inline float glmv_bf(float x) { return float(bfloat(x)); } + +// One threadgroup's sum of every thread's `v` (at most 1024 threads), returned to all of them. +inline float glmv_group_sum(float v, threadgroup float* part, uint t, uint n, uint lane, uint sg) { + v = simd_sum(v); + if (lane == 0) part[sg] = v; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float s = lane < (n + 31) / 32 ? part[lane] : 0.0f; + s = simd_sum(s); + if (lane == 0) part[32] = s; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + const float out = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + return out; +} + +// y[r, c] = y + b[c] (a linear layer's bias, added after its product as nn.Linear's bf16 add rounds it) +[[kernel]] void glmv_bias(device bfloat* y [[buffer(0)]], const device bfloat* b [[buffer(1)]], + constant uint& width [[buffer(2)]], uint2 at [[thread_position_in_grid]]) { + const ulong i = ulong(at.y) * width + at.x; + y[i] = bfloat(float(y[i]) + float(b[at.x])); +} + +// x += y (a block's residual) +[[kernel]] void glmv_add(device bfloat* x [[buffer(0)]], const device bfloat* y [[buffer(1)]], + uint i [[thread_position_in_grid]]) { + x[i] = bfloat(float(x[i]) + float(y[i])); +} + +// mx.fast.rms_norm, one threadgroup a row: w * bf16(x * rsqrt(mean(x^2) + eps)) +[[kernel]] void glmv_rms(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], + device bfloat* y [[buffer(2)]], constant GlmvRows& p [[buffer(3)]], + uint row [[threadgroup_position_in_grid]], uint t [[thread_position_in_threadgroup]], + uint n [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]], + uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device bfloat* xr = x + ulong(row) * p.width; + device bfloat* yr = y + ulong(row) * p.width; + float s = 0.0f; + for (uint i = t; i < p.width; i += n) { + const float v = float(xr[i]); + s += v * v; + } + s = glmv_group_sum(s, part, t, n, lane, sg); + const float inv = metal::precise::rsqrt(s / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) yr[i] = bfloat(float(w[i]) * glmv_bf(float(xr[i]) * inv)); +} + +// mx.fast.layer_norm with weight and bias, one threadgroup a row: w * bf16((x - mean) * rsqrt(var + eps)) + b +[[kernel]] void glmv_layer_norm(const device bfloat* x [[buffer(0)]], const device bfloat* w [[buffer(1)]], + const device bfloat* b [[buffer(2)]], device bfloat* y [[buffer(3)]], + constant GlmvRows& p [[buffer(4)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint n [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + const device bfloat* xr = x + ulong(row) * p.width; + device bfloat* yr = y + ulong(row) * p.width; + float s = 0.0f; + for (uint i = t; i < p.width; i += n) s += float(xr[i]); + const float mean = glmv_group_sum(s, part, t, n, lane, sg) / float(p.width); + float v = 0.0f; + for (uint i = t; i < p.width; i += n) { + const float d = float(xr[i]) - mean; + v += d * d; + } + const float inv = metal::precise::rsqrt(glmv_group_sum(v, part, t, n, lane, sg) / float(p.width) + p.eps); + for (uint i = t; i < p.width; i += n) { + yr[i] = bfloat(glmv_bf(float(w[i]) * glmv_bf((float(xr[i]) - mean) * inv)) + float(b[i])); + } +} + +struct GlmvQkv { + uint L; // rows (patches) of this image + float eps; // the q and k norms' +}; + +// A simdgroup a (row, head, q|k|v) of qkv [L, 3, 16, 64]: q and k RMS-normed and 2-D rotated; out head-major fp32. +[[kernel]] void glmv_qkv(const device bfloat* qkv [[buffer(0)]], const device bfloat* qn [[buffer(1)]], + const device bfloat* kn [[buffer(2)]], const device float* cs [[buffer(3)]], + const device float* sn [[buffer(4)]], device float* Q [[buffer(5)]], + device float* K [[buffer(6)]], device float* V [[buffer(7)]], + constant GlmvQkv& p [[buffer(8)]], uint g [[threadgroup_position_in_grid]], + uint lane [[thread_index_in_simdgroup]]) { + const uint which = g % 3, h = (g / 3) % 16, r = g / 48; + const device bfloat* src = qkv + ulong(r) * 3072 + which * 1024 + h * 64; + const ulong out = (ulong(h) * p.L + r) * 64; + const float a = float(src[lane]), b = float(src[lane + 32]); + if (which == 2) { + V[out + lane] = a; + V[out + lane + 32] = b; + return; + } + const device bfloat* w = which == 0 ? qn : kn; + const float inv = metal::precise::rsqrt(simd_sum(a * a + b * b) / 64.0f + p.eps); + const float na = glmv_bf(float(w[lane]) * glmv_bf(a * inv)); + const float nb = glmv_bf(float(w[lane + 32]) * glmv_bf(b * inv)); + // cos and sin [L, 64]: the height's 16 frequencies, then the width's, twice + const device float* c = cs + ulong(r) * 64; + const device float* s = sn + ulong(r) * 64; + device float* dst = which == 0 ? Q : K; + dst[out + lane] = glmv_bf(na * c[lane] - nb * s[lane]); + dst[out + lane + 32] = glmv_bf(nb * c[lane + 32] + na * s[lane + 32]); +} + +// Each row of S [rows, n] (fp32 scores) as softmax(scale * s), in place; one threadgroup a row. +[[kernel]] void glmv_softmax(device float* S [[buffer(0)]], constant uint& n [[buffer(1)]], + constant float& scale [[buffer(2)]], uint row [[threadgroup_position_in_grid]], + uint t [[thread_position_in_threadgroup]], uint tn [[threads_per_threadgroup]], + uint lane [[thread_index_in_simdgroup]], uint sg [[simdgroup_index_in_threadgroup]]) { + threadgroup float part[33]; + device float* s = S + ulong(row) * n; + float m = -INFINITY; + for (uint i = t; i < n; i += tn) m = max(m, s[i] * scale); + m = simd_max(m); + if (lane == 0) part[sg] = m; + threadgroup_barrier(mem_flags::mem_threadgroup); + if (sg == 0) { + float v = lane < (tn + 31) / 32 ? part[lane] : -INFINITY; + v = simd_max(v); + if (lane == 0) part[32] = v; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + m = part[32]; + threadgroup_barrier(mem_flags::mem_threadgroup); + float z = 0.0f; + for (uint i = t; i < n; i += tn) { + const float e = metal::precise::exp(s[i] * scale - m); + s[i] = e; + z += e; + } + const float inv = 1.0f / glmv_group_sum(z, part, t, tn, lane, sg); + for (uint i = t; i < n; i += tn) s[i] *= inv; +} + +// O [16, rows, 64] (fp32, one chunk's attention) into rows [r0, r0 + rows) of out [L, 1024] (bf16, heads side by side) +[[kernel]] void glmv_heads(const device float* O [[buffer(0)]], device bfloat* out [[buffer(1)]], + constant uint2& p [[buffer(2)]], uint2 at [[thread_position_in_grid]]) { + const uint rows = p.x, r0 = p.y, h = at.x / 64, d = at.x % 64, r = at.y; + out[ulong(r0 + r) * 1024 + at.x] = bfloat(O[(ulong(h) * rows + r) * 64 + d]); +} + +// The vision MLP's and merger's gate: silu(min(g, limit)) * clip(u, -limit, limit), each op rounded as MLX does. +[[kernel]] void glmv_swiglu(const device bfloat* g [[buffer(0)]], const device bfloat* u [[buffer(1)]], + device bfloat* act [[buffer(2)]], constant float& limit [[buffer(3)]], + uint i [[thread_position_in_grid]]) { + const float gate = glmv_bf(min(float(g[i]), limit)); + const float up = glmv_bf(clamp(float(u[i]), -limit, limit)); + const float sig = glmv_bf(1.0f / (1.0f + metal::precise::exp(-gate))); + act[i] = bfloat(glmv_bf(gate * sig) * up); +} + +// erf to fp32 precision (W. J. Cody's rational approximations, as libm's erff). +inline float glmv_erf(float x) { + const float ax = fabs(x); + if (ax < 0.84375f) { + const float z = x * x; + const float r = 1.28379167e-01f + z * (-3.25042108e-01f + z * (-2.84817495e-02f + z * (-5.77027029e-03f + z * -2.37630166e-05f))); + const float s = 1.0f + z * (3.97917233e-01f + z * (6.50222499e-02f + z * (5.08130628e-03f + z * (1.32494738e-04f + z * -3.96022827e-06f)))); + return x + x * (r / s); + } + if (ax < 1.25f) { + const float s = ax - 1.0f; + const float P = -2.36211857e-03f + s * (4.14856106e-01f + s * (-3.72207481e-01f + s * (3.18346620e-01f + s * (-1.10894694e-01f + s * (3.54783031e-02f + s * -2.16637559e-03f))))); + const float Q = 1.0f + s * (1.06420882e-01f + s * (5.40397942e-01f + s * (7.18286547e-02f + s * (1.26171216e-01f + s * (1.36370836e-02f + s * 1.19844998e-02f))))); + const float e = 8.45062911e-01f + P / Q; + return x >= 0 ? e : -e; + } + if (ax >= 6.0f) return x >= 0 ? 1.0f : -1.0f; + const float s = 1.0f / (ax * ax); + float R, S; + if (ax < 1.0f / 0.35f) { + R = -9.86494403e-03f + s * (-6.93858570e-01f + s * (-1.05586262e+01f + s * (-6.23753357e+01f + s * (-1.62396660e+02f + s * (-1.84605087e+02f + s * (-8.12874374e+01f + s * -9.81432934e+00f)))))); + S = 1.0f + s * (1.96512714e+01f + s * (1.37657754e+02f + s * (4.34565887e+02f + s * (6.45387268e+02f + s * (4.29008148e+02f + s * (1.08635002e+02f + s * (6.57024977e+00f + s * -6.04244141e-02f))))))); + } else { + R = -9.86494310e-03f + s * (-7.99283242e-01f + s * (-1.77579556e+01f + s * (-1.60636383e+02f + s * (-6.37566467e+02f + s * (-1.02509509e+03f + s * -4.83519196e+02f))))); + S = 1.0f + s * (3.03380604e+01f + s * (3.25792511e+02f + s * (1.53672961e+03f + s * (3.19985820e+03f + s * (2.55305029e+03f + s * (4.74528534e+02f + s * -2.24409515e+01f)))))); + } + const float zh = as_type(as_type(ax) & 0xfffff000u); + const float r = metal::precise::exp(-zh * zh - 0.5625f) * metal::precise::exp((zh - ax) * (zh + ax) + R / S); + const float e = 1.0f - r / ax; + return x >= 0 ? e : -e; +} + +// nn.gelu (exact): x * (1 + erf(x / sqrt(2))) / 2, each op rounded to bf16 as the compiled MLX graph does +[[kernel]] void glmv_gelu(device bfloat* x [[buffer(0)]], uint i [[thread_position_in_grid]]) { + const float v = float(x[i]); + const float t = glmv_bf(v / 1.41421356237309504880f); + const float e = glmv_bf(glmv_erf(t)); + const float w = glmv_bf(v * glmv_bf(1.0f + e)); + x[i] = bfloat(w / 2.0f); +} diff --git a/zig/kernels/metal/sources.zig b/zig/kernels/metal/sources.zig index 74d24354e9..015e12383a 100644 --- a/zig/kernels/metal/sources.zig +++ b/zig/kernels/metal/sources.zig @@ -91,6 +91,8 @@ pub const glm_glue = @embedFile("glm_glue.metal"); pub const glm_rows = @embedFile("glm_rows.metal"); /// GLM-5.3-Flash's latent attention for rows that read every key: 64 heads as one matrix on the tensor units. pub const glm_attn = @embedFile("glm_attn.metal"); +/// GLM-5.3-Flash's image tower: the prefill GEMM template's projections and attention products, then its own steps. +pub const glm_vision = @embedFile("prefill/gemm_nax.metal") ++ @embedFile("glm_vision.metal"); /// A prompt chunk's KDA layer in three passes, appended to the generated kda_rows source (its helpers). pub const glm_kda_prompt = @embedFile("glm_kda_prompt.metal"); /// GLM-5.3-Flash speed-up mode's exchange kernels (families/glm/ep.zig). diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index 903dc4d40b..9fcfd7656c 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -4,11 +4,16 @@ const lanes = @import("lanes"); const Allocator = std.mem.Allocator; /// The server's name for one request, unique while the process lives. +pub const tfvideo = @import("tfvideo.zig"); + pub const Id = u64; /// One request's sampling (Python's exact_sampling.Sampling); a request without one decodes greedily. pub const Sampling = lanes.Sampling; +/// One image of a request (``Request.images``): the tower's input and the placeholder rows its output replaces. +pub const Image = lanes.stream.Image; + /// A prompt's reproducible seed (``seed_for``), for requests that name none. pub const seedFor = lanes.sampling.seedFor; @@ -88,8 +93,77 @@ pub const Request = struct { logprobs: ?u8 = null, /// Rows for each prompt token after the first, with ranks and this many best tokens, when ``Info.prompt_logprobs``. prompt_logprobs: ?u8 = null, + /// Images whose rows replace placeholder tokens of ``prompt``, in order; only when ``Info.vision`` is set. + images: []const Image = &.{}, + /// The prompt as the prompt cache matches it (each image's placeholders by its content); empty: ``prompt``. + cache_key: []const u32 = &.{}, + + /// The prompt as the prompt cache matches it: an image request's placeholders stand for their images. + pub fn cacheKey(r: *const Request) []const u32 { + return if (r.cache_key.len > 0) r.cache_key else r.prompt; + } +}; + +/// An image prepared for the engine's tower (``Vision.prepare``): patches, patch grid, tokens and a content hash. +pub const PreparedImage = struct { pixels: []f32, gh: u32, gw: u32, tokens: u32, hash: u64 }; + +/// Image input an engine serves (``Info.vision``): image bytes to tower input, and what a request may hold. +pub const Vision = struct { + ctx: *anyopaque, + /// The image's bytes (any format the platform decodes) at most ``max_tokens`` tokens; allocated in ``a``. + prepare: *const fn (ctx: *anyopaque, a: Allocator, bytes: []const u8, max_tokens: u32) anyerror!PreparedImage, + /// The token the chat template writes once per image, expanded to the image's tokens. + image_token: u32, + /// The tokens the template wraps an image in (the image token among them): text that spells one stays text. + markers: []const u32 = &.{}, + /// Image tokens one request may hold in all, and images. + image_tokens: u32, + max_images: u32, + /// Video input, when the engine's family reads it. + video: ?VideoOffer = null, +}; + +/// A video's facts as the server's decoder reads them, and its frames on request. +pub const VideoSource = struct { + frames: u64, // the stream's frame count + rate_num: u64, // its average frame rate, num / den + rate_den: u64, + width: u32, + height: u32, + ctx: *anyopaque, + /// The frames at ``indices`` (ascending, distinct), each to ``sink`` in order. + decode: *const fn (ctx: *anyopaque, indices: []const u32, sink: FrameSink) anyerror!void, +}; + +/// Where a decoded frame goes: RGB rows of ``width`` * 3 bytes, ``stride`` bytes apart. +pub const FrameSink = struct { + ctx: *anyopaque, + take: *const fn (ctx: *anyopaque, k: usize, rgb: []const u8, width: u32, height: u32, stride: usize) anyerror!void, +}; + +/// A video prepared for the tower: each group of frames an image, and each group's time in seconds. +pub const PreparedVideo = struct { groups: []PreparedImage, times: []f64 }; + +/// Video input an engine serves (``Vision.video``): how a decoded video becomes images and how the prompt holds them. +pub const VideoOffer = struct { + /// ``src`` at most ``max_tokens`` tokens and ``max_frames`` sampled frames; allocated in ``a``. + prepare: *const fn (ctx: *anyopaque, a: Allocator, src: VideoSource, max_tokens: u32, max_frames: u32) anyerror!PreparedVideo, + /// The token the chat template writes once per video, expanded to its groups. + video_token: u32, + /// The tokens around each group's image tokens; the group's time follows as text. + group_open: u32, + group_close: u32, + /// The tokens the template wraps a video in (the video token among them): text that spells one stays text. + markers: []const u32 = &.{}, + /// Video tokens one request may hold in all, videos, and frames sampled a video. + video_tokens: u32, + max_videos: u32, + max_frames: u32, }; +/// --vision: image and video input's limits (``Open.vision``). +pub const VisionOpen = struct { image_tokens: u32 = 4096, max_images: u32 = 4, video_tokens: u32 = 16384, max_videos: u32 = 2, max_frames: u32 = 256 }; + pub const Reason = enum { stop, length, cancelled, failed }; /// What a finished reply's rounds did, for its runtime fields and /metrics. @@ -160,6 +234,8 @@ pub const Info = struct { plain_only: bool = false, /// The immutable retained-prefix plan applied at startup, not a process memory limit or current cache occupancy. prompt_cache_plan: ?PromptCachePlan = null, + /// Image input (--vision); null: text only. + vision: ?Vision = null, /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, /// Whether a request's caches fit beside the running streams' now (null: a free lane does); else it waits in order. @@ -225,6 +301,8 @@ pub const Open = struct { master_port: u16 = 29551, /// --policy: the engine's policy as `key=value,...`; empty: its defaults. policy: []const u8 = "", + /// --vision: load the checkpoint's image tower (null: text only; an engine without one refuses). + vision: ?VisionOpen = null, }; /// An opened engine; ``close`` stops its thread and frees its backend. diff --git a/zig/src/core/lane_host.zig b/zig/src/core/lane_host.zig index 4394fdac05..6fac9e4668 100644 --- a/zig/src/core/lane_host.zig +++ b/zig/src/core/lane_host.zig @@ -77,7 +77,7 @@ pub const LaneHost = struct { fn kept(ptr: *anyopaque, s: *lanes.Stream, at: u32) void { const job: *Job = @ptrCast(@alignCast(ptr)); job.reported(); // before a keep can evict the entry the pass restored - if (job.host.cache) |store| _ = store.keepBorn(job.request.prompt, born(job.request), at, s, job.request.chunks, job.host.spans(job.request)); + if (job.host.cache) |store| _ = store.keepBorn(job.request.cacheKey(), born(job.request), at, s, job.request.chunks, job.host.spans(job.request)); } /// Report a restored prefix or its failed copy; an untouched prefix remains kept. @@ -87,7 +87,7 @@ pub const LaneHost = struct { const store = job.host.cache orelse return; if (!job.started) return; job.stream.reuse.saved = null; // the backend restores before its first chunk; a later keep may free it - if (job.stream.reuse_failed) store.resumed(e, job.request.prompt, false) else if (job.stream.cached == e.at) store.resumed(e, job.request.prompt, true); + if (job.stream.reuse_failed) store.resumed(e, job.request.cacheKey(), false) else if (job.stream.cached == e.at) store.resumed(e, job.request.cacheKey(), true); } }; @@ -309,7 +309,7 @@ pub const LaneHost = struct { if (r.background and r.max_tokens == 0 and h.warm_quiet > 0 and !h.quiet(job)) return h.cancel(job); // other work came first: it never ran var reuse: lanes.stream.Reuse = .{}; // the entry stays alive until the backend restores it: nothing keeps between here and this stream's own pass - if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.prompt, born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { + if (h.cache) |store| if (r.prompt_logprobs == null and !r.untaught) if (store.lookupRewind(h.gpa, r.cacheKey(), born(r), r.history_len, r.rewind_len, r.shared_prefixes, r.chunks, h.spans(r))) |l| { job.entry = l.entry; job.kept0 = store.counts.kept; job.marks = l.marks; @@ -344,6 +344,7 @@ pub const LaneHost = struct { .decode_spans = h.spans(r), .decide = r.decide, .untaught = r.untaught, + .images = r.images, }) catch { job.proposer.deinit(); return h.drop(job, "out of memory"); diff --git a/zig/src/core/lane_host_reuse_test.zig b/zig/src/core/lane_host_reuse_test.zig index 01b1fc6bec..69fca5ae1f 100644 --- a/zig/src/core/lane_host_reuse_test.zig +++ b/zig/src/core/lane_host_reuse_test.zig @@ -106,6 +106,94 @@ test "a lane host resumes a conversation from its kept prompt state and reports try std.testing.expectEqual(@as(usize, 2), store.entries.items.len); } +test "an image's placeholders are cached by its content: another image under the same ids resumes none of it" { + const gpa = std.testing.allocator; + var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); + defer cfg.deinit(gpa); + var target: lanes.fake.Fake = .{ .gpa = gpa }; + defer target.deinit(); + var clock: lanes.fake.FixedClock = .{}; + var core = lanes.Engine.init(gpa, &cfg, target.backend(), clock.clock()); + defer core.deinit(); + const Snaps = struct { + fn bytes(_: *anyopaque, at: u32) u64 { + return 64 + 4 * @as(u64, at); + } + fn save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!pc.Saved { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + return f.save(@ptrCast(@alignCast(owner.?)), at); + } + fn restore(_: *anyopaque, _: ?*anyopaque, _: pc.Saved) anyerror!void { + return error.BackendRestores; // lane-core backends restore inside their own prompt pass + } + fn drop(ptr: *anyopaque, saved: pc.Saved) void { + const f: *lanes.fake.Fake = @ptrCast(@alignCast(ptr)); + f.drop(saved); + } + }; + var store = pc.Store.init(gpa, .{ .ptr = &target, .vtable = &.{ .bytes = Snaps.bytes, .save = Snaps.save, .restore = Snaps.restore, .drop = Snaps.drop } }, .{ .min_gap = 4, .min_prompt = 0 }, 1 << 20); + defer store.deinit(); + var host = LaneHost.init(gpa, std.testing.io, &core, .{ .lanes = 2 }); + host.cache = &store; + try host.start(); + defer host.stop(); + const Box = struct { + mutex: std.Io.Mutex = .init, + tokens: std.ArrayList(u32) = .empty, + cached: ?u32 = null, + done: ?Reason = null, + fn event(ctx: *anyopaque, _: Id, e: *const Event) void { + const b: *@This() = @ptrCast(@alignCast(ctx)); + b.mutex.lockUncancelable(std.testing.io); + defer b.mutex.unlock(std.testing.io); + switch (e.*) { + .prefilled => |c| b.cached = c, + .logprobs, .prompt_logprobs => {}, + .tokens => |t| b.tokens.appendSlice(gpa, t) catch {}, + .finished => |f| b.done = f.reason, + } + } + fn wait(b: *@This()) Reason { + while (true) { + b.mutex.lockUncancelable(std.testing.io); + const d = b.done; + b.mutex.unlock(std.testing.io); + if (d) |r| return r; + std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {}; + } + } + }; + const e = host.engine(); + // prompt ids 3..5 are one image's placeholders (token 9); cache_key holds ids derived from each image's content + const t1 = [_]u32{ 3, 1, 9, 9, 9, 9, 2, 6, 5, 3 }; + var key_a = t1; + var key_b = t1; + for (2..6) |k| { + key_a[k] = 0x8000_0000 + @as(u32, @intCast(k)); + key_b[k] = 0x9000_0000 + @as(u32, @intCast(k)); + } + var b1: Box = .{}; + defer b1.tokens.deinit(gpa); + const r1: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_a }; + try e.submit(1, &r1, .{ .ctx = &b1, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b1.wait()); + try std.testing.expectEqual(@as(usize, 1), store.entries.items.len); // image A's state at 8, past its image + // the same text and placeholder ids with image B: the kept state is A's, so B resumes none of it + var b2: Box = .{}; + defer b2.tokens.deinit(gpa); + const r2: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_b }; + try e.submit(2, &r2, .{ .ctx = &b2, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b2.wait()); + try std.testing.expectEqual(@as(?u32, 0), b2.cached); + // image A again resumes its own state + var b3: Box = .{}; + defer b3.tokens.deinit(gpa); + const r3: Request = .{ .prompt = &t1, .max_tokens = 6, .history_len = 8, .cache_key = &key_a }; + try e.submit(3, &r3, .{ .ctx = &b3, .event = Box.event }); + try std.testing.expectEqual(Reason.length, b3.wait()); + try std.testing.expectEqual(@as(?u32, 8), b3.cached); +} + test "a turn's mark evicting the state it resumed from: the hit is reported first and the entry freed once" { const gpa = std.testing.allocator; var cfg = try lanes.Config.init(gpa, .{ .exact_width = 8, .gpu_tokens = true, .hidden_rows = true }, 8, 7); diff --git a/zig/src/core/lanes/stream.zig b/zig/src/core/lanes/stream.zig index f028d10cc4..906c568ead 100644 --- a/zig/src/core/lanes/stream.zig +++ b/zig/src/core/lanes/stream.zig @@ -97,6 +97,9 @@ pub const Held = struct { parents: ?[]i32 = null, }; +/// One image of a prompt: tower input (patches [gh * gw, width]) for the `tokens` placeholder rows from `at`. +pub const Image = struct { pixels: []const f32, gh: u32, gw: u32, at: u32, tokens: u32 }; + pub const Spec = struct { id: []const u8, prompt: []const u32, @@ -123,6 +126,7 @@ pub const Spec = struct { prompt_logprobs: ?u8 = null, // and for each prompt token after the first, with ranks (the prompt pass starts at 0) grammar: ?*grammar.Constraint = null, // structured output: the reply stays in it (the creator frees it) decode_spans: []const [2]u32 = &.{}, // prompt rows an earlier reply decoded (Request.decode_spans), for backends that keep their bits + images: []const Image = &.{}, // in prompt order; the backend's tower fills their rows }; pub const Stream = struct { @@ -155,6 +159,7 @@ pub const Stream = struct { decode_spans: []const [2]u32 = &.{}, decide: ?u32 = null, untaught: bool = false, + images: []const Image = &.{}, cached: u32 = 0, // prompt tokens the backend restored from `reuse` (its prompt pass started there) reuse_failed: bool = false, // the backend's restore of `reuse` failed: it prefilled from 0 context: std.ArrayList(u32) = .empty, @@ -216,6 +221,7 @@ pub const Stream = struct { .prompt_logprobs = spec.prompt_logprobs, .grammar = spec.grammar, .decode_spans = spec.decode_spans, + .images = spec.images, }; // under a grammar a forced close ends at the think end: the grammar takes the reply from the next token if (spec.grammar) |g| if (g.after) |end| if (std.mem.indexOfScalar(u32, s.think_close, end)) |at| { diff --git a/zig/src/core/tfvideo.zig b/zig/src/core/tfvideo.zig new file mode 100644 index 0000000000..5de99bcabd --- /dev/null +++ b/zig/src/core/tfvideo.zig @@ -0,0 +1,107 @@ +//! libtfvideo (tf_video.c over FFmpeg 9.0.2), loaded on the first video: stream facts as PyAV reads them, RGB frames. +const std = @import("std"); +const builtin = @import("builtin"); + +pub const Info = extern struct { frames: i64, rate_num: c_int, rate_den: c_int, width: c_int, height: c_int, counted: c_int }; +const FrameFn = *const fn (ctx: ?*anyopaque, k: c_int, index: c_int, rgb: [*]const u8, width: c_int, height: c_int, stride: c_int) callconv(.c) c_int; + +const Fns = struct { + open: *const fn ([*]const u8, usize, *?*anyopaque, *Info) callconv(.c) c_int, + frames: *const fn (*anyopaque, [*]const c_int, c_int, FrameFn, ?*anyopaque) callconv(.c) c_int, + close: *const fn (*anyopaque) callconv(.c) void, + version: *const fn () callconv(.c) [*:0]const u8, +}; + +/// Beside the server (zig-out/native/lib, where `zig build video` installs it), then the loader's own search path. +const names: []const [:0]const u8 = if (builtin.os.tag == .macos) + &.{ "@executable_path/../lib/libtfvideo.dylib", "libtfvideo.dylib" } +else + &.{"libtfvideo.so"}; + +var lib: ?Fns = null; +var lib_failed = false; +var lib_mutex: std.atomic.Mutex = .unlocked; + +fn load() ?Fns { + while (!lib_mutex.tryLock()) std.atomic.spinLoopHint(); + defer lib_mutex.unlock(); + if (lib) |l| return l; + if (lib_failed) return null; + lib_failed = true; + var dl = for (names) |name| { + break std.DynLib.open(name) catch continue; + } else return null; + var f: Fns = undefined; + const info = @typeInfo(Fns).@"struct"; + inline for (info.field_names, info.field_types) |name, T| { + @field(f, name) = dl.lookup(T, "tf_video_" ++ name) orelse return null; + } + lib_failed = false; + lib = f; + return f; +} + +/// The FFmpeg release the loaded library was built against, or null when it can't be loaded. +pub fn version() ?[]const u8 { + const f = load() orelse return null; + return std.mem.span(f.version()); +} + +pub const Error = error{ NoLibrary, Invalid, NoStream, NoDecoder, OutOfMemory, Decode, Short, Stopped }; + +fn err(rc: c_int) Error { + return switch (rc) { + -2 => error.NoStream, + -3 => error.NoDecoder, + -4 => error.OutOfMemory, + -5 => error.Decode, + -6 => error.Short, + -7 => error.Stopped, + else => error.Invalid, + }; +} + +/// An open container over bytes the caller keeps alive until `close`. +pub const Video = struct { + handle: *anyopaque, + info: Info, + fns: Fns, + + pub fn open(bytes: []const u8) Error!Video { + const f = load() orelse return error.NoLibrary; + var h: ?*anyopaque = null; + var info: Info = undefined; + const rc = f.open(bytes.ptr, bytes.len, &h, &info); + if (rc != 0) return err(rc); + return .{ .handle = h.?, .info = info, .fns = f }; + } + + pub fn close(v: Video) void { + v.fns.close(v.handle); + } + + /// The frames at `indices` (ascending, distinct), each to `sink.take(k, rgb, width, height, stride)` in order. + pub fn frames(v: Video, indices: []const u32, sink: anytype) anyerror!void { + const S = @TypeOf(sink); + const Ctx = struct { sink: S, failure: ?anyerror = null }; + var ctx: Ctx = .{ .sink = sink }; + const cb = struct { + fn take(p: ?*anyopaque, k: c_int, _: c_int, rgb: [*]const u8, w: c_int, h: c_int, stride: c_int) callconv(.c) c_int { + const c: *Ctx = @ptrCast(@alignCast(p.?)); + const n: usize = @as(usize, @intCast(stride)) * @as(usize, @intCast(h)); + c.sink.take(@intCast(k), rgb[0..n], @intCast(w), @intCast(h), @intCast(stride)) catch |e| { + c.failure = e; + return 1; + }; + return 0; + } + }.take; + const idx = try std.heap.page_allocator.alloc(c_int, indices.len); + defer std.heap.page_allocator.free(idx); + for (idx, indices) |*d, s| d.* = @intCast(s); + const rc = v.fns.frames(v.handle, idx.ptr, @intCast(idx.len), cb, &ctx); + if (ctx.failure) |e| return e; + if (rc < 0) return err(rc); + if (rc != indices.len) return error.Short; + } +}; diff --git a/zig/src/core/tokenizer/tokenizer.zig b/zig/src/core/tokenizer/tokenizer.zig index d856838549..183f043ee3 100644 --- a/zig/src/core/tokenizer/tokenizer.zig +++ b/zig/src/core/tokenizer/tokenizer.zig @@ -143,10 +143,32 @@ pub const Tokenizer = struct { m.put(text, first, ids.items[from..]); } + /// `encode` with no added token matched: text that spells one (a quoted marker) is ordinary text. + pub fn encodePlain(t: *const Tokenizer, a: Allocator, text: []const u8) Error![]u32 { + var arena = std.heap.ArenaAllocator.init(a); + defer arena.deinit(); + var ctx = steps.Context{ .a = arena.allocator(), .matcher = .{ .a = arena.allocator() } }; + var valid = text; + if (!std.unicode.utf8ValidateSlice(text)) { + var fixed: std.ArrayList(u8) = .empty; + try unicode.appendLossy(&fixed, ctx.a, text); + valid = fixed.items; + } + var ids: std.ArrayList(u32) = .empty; + if (valid.len > 0) try t.encodeTextWith(&ctx, valid, true, &ids, false); + return a.dupe(u32, ids.items); + } + fn encodeText(t: *const Tokenizer, ctx: *steps.Context, text: []const u8, first: bool, ids: *std.ArrayList(u32)) Error!void { + return t.encodeTextWith(ctx, text, first, ids, true); + } + + fn encodeTextWith(t: *const Tokenizer, ctx: *steps.Context, text: []const u8, first: bool, ids: *std.ArrayList(u32), added: bool) Error!void { var zero: usize = if (first) unicode.decodeAt(text, 0).len else 0; const normalized = if (t.normalizer) |*n| try n.apply(ctx, text, &zero) else text; - for (try t.splitAdded(ctx, &t.normalized, normalized)) |seg| { + const whole = [_]Segment{.{ .start = 0, .end = normalized.len, .id = null }}; + const segs: []const Segment = if (added) try t.splitAdded(ctx, &t.normalized, normalized) else &whole; + for (segs) |seg| { if (seg.id) |id| { try ids.append(ctx.a, id); continue; diff --git a/zig/src/families/glm/backend.zig b/zig/src/families/glm/backend.zig index 85b6f8c78e..416ea57b7a 100644 --- a/zig/src/families/glm/backend.zig +++ b/zig/src/families/glm/backend.zig @@ -3,6 +3,7 @@ const std = @import("std"); const lanes = @import("lanes"); const st = @import("state.zig"); const slots_mod = @import("slots.zig"); +const slot_images = @import("slot_images.zig"); const pool = @import("pool.zig"); const mirror = @import("mirror.zig"); const snapshot = @import("snapshot.zig"); @@ -101,6 +102,11 @@ pub const Backend = struct { try mirror.send(e, .begin, b.words.items); try b.sl.begin(i, prompt, s.drafts); s.cached = 0; + if (s.images.len > 0) { // the images through the tower first: the chunks inject their rows + if (e.ep != null) return error.VisionOnPair; + try slot_images.encode(b.sl, i, s.images); + } + defer if (s.images.len > 0) slot_images.end(b.sl); var at: u32 = 0; if (s.reuse.saved) |saved| { // a kept state of this prompt's prefix: the pass starts there const snap: *snapshot.Snap = @ptrCast(@alignCast(saved)); diff --git a/zig/src/families/glm/checks.zig b/zig/src/families/glm/checks.zig index 98404ca22a..50c0285c40 100644 --- a/zig/src/families/glm/checks.zig +++ b/zig/src/families/glm/checks.zig @@ -168,7 +168,7 @@ pub fn checkMatmul(e: *Engine, prompt: []const u32) !void { const out = try e.arena.buffer(@as(usize, rows) * N * 2); var x = e.ctx(); const b = e.begin(); - fwd.embed(&x, b.enc, e.prompt_ids, rows); + fwd.embed(&x, b.enc, e.prompt_ids, rows, 0); fwd.boundary(&x, b.enc, rows, false, L.hc.?[0], L.in_norm); fwd.qmv(&x, b.enc, e.k.qmv_kda_in, e.sc.normed, a.in_proj, out, rows); affine_mm.rowSums(b.enc, e.k.mm_bf16, 64, e.sc.normed, pr.sums, rows, a.in_proj.k); diff --git a/zig/src/families/glm/engine.zig b/zig/src/families/glm/engine.zig index e1ffe5ac41..9b33850a46 100644 --- a/zig/src/families/glm/engine.zig +++ b/zig/src/families/glm/engine.zig @@ -4,6 +4,7 @@ const mtl = @import("metal"); const cfg = @import("config.zig"); const wts = @import("weights.zig"); const page_cache = @import("page_cache.zig"); +const vision_mod = @import("vision.zig"); const st = @import("state.zig"); const fwd = @import("forward.zig"); const mtp = @import("mtp.zig"); @@ -84,6 +85,9 @@ pub const Engine = struct { rank_log: bool = false, // GLM_RANKS=1: each MTP depth's logits kept, and the target's rank in them where drafts miss ep_arena: std.heap.ArenaAllocator, // the link settings, alive as long as the link residency: ?mtl.ResidencySet = null, + vision: ?*vision_mod.Vision = null, // --vision: the image tower + vision_rows: Ref = undefined, // bf16 [image tokens, hidden]: a prompt's image rows (vision only) + vision_limits: vision_mod.Limits = .{}, keepalive_sets: [1]mtl.ResidencySet = undefined, // the residency set, for the idle keepalive's commit keepalive_target: mtl.keepalive.Target = undefined, // the engine's queue, set at load for the server's ticker load_seconds: f64 = 0, @@ -95,11 +99,11 @@ pub const Engine = struct { /// The checkpoint in `dir`, caches for `cap` tokens; GLM_LAYERS=N: the first N layers only; GLM_EP=settings: half the experts. pub fn load(gpa: std.mem.Allocator, dir: []const u8, cap: u32) !*Engine { - return loadWith(gpa, dir, cap, if (std.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false); + return loadWith(gpa, dir, cap, if (std.c.getenv("GLM_EP")) |v| std.mem.span(v) else null, false, null); } - /// `load` with expert parallel over the link in `ep_path` (this Mac's settings), or on one Mac when null. - pub fn loadWith(gpa: std.mem.Allocator, dir: []const u8, cap: u32, ep_path: ?[]const u8, learn: bool) !*Engine { + /// `load` with expert parallel over `ep_path`'s link (one Mac when null) and `vision`'s image tower and rows. + pub fn loadWith(gpa: std.mem.Allocator, dir: []const u8, cap: u32, ep_path: ?[]const u8, learn: bool, vision: ?vision_mod.Limits) !*Engine { const e = try gpa.create(Engine); // undefined memory: every field is set below errdefer gpa.destroy(e); const pool = mtl.objc.Pool.push(); @@ -108,6 +112,7 @@ pub const Engine = struct { e.gpa = gpa; e.ev = 0; e.residency = null; + e.vision = null; e.gpu = .{ 0, 0 }; e.fused_route = if (std.c.getenv("GLM_ROUTE")) |v| v[0] != '0' else true; e.draft_vocab = if (std.c.getenv("GLM_DRAFT_VOCAB")) |v| std.fmt.parseInt(u32, std.mem.span(v), 10) catch 0 else 0; @@ -215,10 +220,13 @@ pub const Engine = struct { me.rest[2] = @intFromBool(learn); // --learn on both Macs or neither: each holds its half of a learned state e.ep = try ep_mod.Ep.init(gpa, e.device, s, me); } + if (vision) |lim| try vision_mod.attach(e, gpa, dir, f.bytes[0..f.size], lim); // --vision: the tower and its rows + errdefer if (e.vision) |v| v.deinit(); // opt-in: wiring 181 GB leaves macOS nothing to reclaim if another model shares the Mac (Flash Next runs without) - if (std.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len)) |set| { + if (std.c.getenv("GLM_RESIDENCY") == null) {} else if (e.device.residencySet(e.w.buffers.items.len + e.arena.buffers.items.len + if (e.vision) |v| v.buffers().len else 0)) |set| { for (e.w.buffers.items) |b| set.add(b); for (e.arena.buffers.items) |b| set.add(b); + if (e.vision) |v| for (v.buffers()) |b| set.add(b); set.commit(); set.requestResidency(); e.queue.addResidencySet(set); @@ -268,6 +276,7 @@ pub const Engine = struct { set.deinit(); } if (e.ep) |ep| ep.deinit(gpa); + if (e.vision) |v| v.deinit(); e.ep_arena.deinit(); if (e.pr) |*p| p.deinit(); if (e.round_fence) |f| f.deinit(); diff --git a/zig/src/families/glm/forward.zig b/zig/src/families/glm/forward.zig index 330d9324e6..5d1801caee 100644 --- a/zig/src/families/glm/forward.zig +++ b/zig/src/families/glm/forward.zig @@ -31,8 +31,29 @@ pub const Ctx = struct { draft_vocab: u32 = 154880, // the MTP head's tokens: the vocabulary's first this many m_row: u32 = 0, // the MTP head's logits row in m_logits (a rank log keeps each depth's in its own row) segs: []const Seg = &.{}, // a shared window's streams in row order (empty: every row is `s`'s) + inject: ?*const Inject = null, // a prompt pass's image rows over its placeholder embeddings }; +/// An image prompt's tower output: prompt rows [at, at + n) of each span embed as rows [row, row + n) of `rows`. +pub const Inject = struct { rows: Ref, spans: []const Span }; +pub const Span = struct { at: u32, n: u32, row: u32 }; + +/// The image rows among `rows` embedding rows (positions pos..) written over `out`; nothing without an injection. +pub fn inject(x: *const Ctx, e: mtl.ComputeEncoder, out: Ref, rows: u32, pos: u32) void { + const inj = x.inject orelse return; + const D: usize = x.c.hidden; + for (inj.spans) |sp| { + const lo = @max(sp.at, pos); + const hi = @min(sp.at + sp.n, pos + rows); + if (lo >= hi) continue; + const n: usize = hi - lo; + e.setPipeline(x.k.copy_u32); + bind(e, 0, .{ inj.rows.at((sp.row + lo - sp.at) * D * 2), out.at((lo - pos) * D * 2) }); + e.setValue(@as(u32, @intCast(n * D / 2)), 2); + e.dispatchThreads(size(n * D / 2, 1, 1), size(256, 1, 1)); + } +} + /// One stream's rows in a shared window: its caches, its first row, its rows and the position of the first. pub const Seg = struct { s: *st.State, row0: u32, rows: u32, pos: u32 }; @@ -136,11 +157,12 @@ pub fn scale(x: *const Ctx, e: mtl.ComputeEncoder, in: Ref, out: Ref, rows: u32, e.dispatchThreads(size(width, rows, 1), size(@min(width, 256), 1, 1)); } -/// The window's tokens (`ids`, u32) as embedding rows in all four streams. -pub fn embed(x: *Ctx, e: mtl.ComputeEncoder, ids: Ref, rows: u32) void { +/// The window's tokens (`ids`, u32, at positions pos..) as embedding rows in all four streams, image rows injected. +pub fn embed(x: *Ctx, e: mtl.ComputeEncoder, ids: Ref, rows: u32, pos: u32) void { const sc = x.sc; const D = x.c.hidden; embedRows(x, e, ids, sc.h, rows); + inject(x, e, sc.h, rows, pos); e.setPipeline(x.k.streams); bind(e, 0, .{ sc.h, sc.x[0] }); e.setValue([2]u32{ D, rows }, 2); @@ -438,7 +460,7 @@ pub fn backbone(x: *Ctx, e: mtl.ComputeEncoder, ids: Ref, rows: u32, pos: u32) v const c = x.c; const sc = x.sc; const k = x.skip; - if (k & Class.ends == 0) embed(x, e, ids, rows); + if (k & Class.ends == 0) embed(x, e, ids, rows, pos); const plane = @as(usize, rows) * c.hidden * 2; snap(x, e, sc.h, plane); var pending = false; diff --git a/zig/src/families/glm/glm.zig b/zig/src/families/glm/glm.zig index 11a7502c8f..eafc634757 100644 --- a/zig/src/families/glm/glm.zig +++ b/zig/src/families/glm/glm.zig @@ -15,6 +15,9 @@ pub const mirror = @import("mirror.zig"); pub const timing = @import("timing.zig"); pub const affine = @import("affine.zig"); pub const wide = @import("wide.zig"); +pub const vision = @import("vision.zig"); +pub const image = @import("image.zig"); +pub const video = @import("video.zig"); test { _ = config; @@ -38,4 +41,7 @@ test { _ = @import("names.zig"); _ = @import("weights_test.zig"); _ = @import("wide.zig"); + _ = vision; + _ = image; + _ = video; } diff --git a/zig/src/families/glm/image.zig b/zig/src/families/glm/image.zig new file mode 100644 index 0000000000..0ff08383ca --- /dev/null +++ b/zig/src/families/glm/image.zig @@ -0,0 +1,445 @@ +//! An image's bytes to GLM5-Next vision patches as TensorFold 0.6.6's --vision makes them (PIL, then mlx-vlm). +const std = @import("std"); +const vision = @import("vision.zig"); + +pub const Rgb = struct { + w: u32, + h: u32, + px: []u8, // row-major RGB, 3 bytes a pixel + + pub fn deinit(r: Rgb, gpa: std.mem.Allocator) void { + gpa.free(r.px); + } +}; + +/// One image ready for the tower: its patches, its patch grid and its tokens (gh * gw / 4). +pub const Prepared = struct { + pixels: []f32, // [gh * gw, 1176] + gh: u32, + gw: u32, + tokens: u32, + hash: u64, // of the image's bytes: two prompts with the same placeholders but other images are other prompts +}; + +pub const Error = error{ BadImage, ImageTooLarge, Animated, ImageTooSmall, OutOfMemory }; + +/// images.py's ImageLimits: an image's encoded bytes, sides and pixels. +pub const max_bytes = 10 << 20; +pub const max_side = 8192; +pub const max_pixels = 16 << 20; + +// ---- ImageIO and CoreGraphics (C APIs; the server binary links both frameworks) ---- +const CF = *anyopaque; +const CGRect = extern struct { x: f64, y: f64, w: f64, h: f64 }; +extern "c" fn CFDataCreate(alloc: ?CF, bytes: [*]const u8, len: isize) ?CF; +extern "c" fn CFDataGetBytePtr(data: CF) [*]const u8; +extern "c" fn CFDataGetLength(data: CF) isize; +extern "c" fn CFRelease(cf: CF) void; +extern "c" fn CGImageSourceCreateWithData(data: CF, options: ?CF) ?CF; +extern "c" fn CGImageSourceGetCount(src: CF) usize; +extern "c" fn CGImageSourceCreateImageAtIndex(src: CF, index: usize, options: ?CF) ?CF; +extern "c" fn CGImageGetWidth(img: CF) usize; +extern "c" fn CGImageGetHeight(img: CF) usize; +extern "c" fn CGImageGetBitsPerComponent(img: CF) usize; +extern "c" fn CGImageGetBitsPerPixel(img: CF) usize; +extern "c" fn CGImageGetBytesPerRow(img: CF) usize; +extern "c" fn CGImageGetBitmapInfo(img: CF) u32; +extern "c" fn CGImageGetColorSpace(img: CF) ?CF; +extern "c" fn CGColorSpaceGetModel(cs: CF) i32; +extern "c" fn CGImageGetDataProvider(img: CF) ?CF; +extern "c" fn CGDataProviderCopyData(p: CF) ?CF; +extern "c" fn CGColorSpaceCreateDeviceRGB() ?CF; +extern "c" fn CGColorSpaceRelease(cs: CF) void; +extern "c" fn CGBitmapContextCreate(data: ?*anyopaque, w: usize, h: usize, bpc: usize, bpr: usize, cs: CF, info: u32) ?CF; +extern "c" fn CGContextDrawImage(ctx: CF, rect: CGRect, img: CF) void; +extern "c" fn CGContextRelease(ctx: CF) void; +extern "c" fn CGContextSetRGBFillColor(ctx: CF, r: f64, g: f64, b: f64, a: f64) void; +extern "c" fn CGContextFillRect(ctx: CF, rect: CGRect) void; +extern "c" fn CGImageSourceCopyPropertiesAtIndex(src: CF, index: usize, options: ?CF) ?CF; +extern "c" fn CFDictionaryGetValue(dict: CF, key: ?*const anyopaque) ?CF; +extern "c" fn CFNumberGetValue(num: CF, kind: isize, out: *anyopaque) bool; +extern var kCGImagePropertyOrientation: ?*const anyopaque; + +const alpha_mask: u32 = 0x1f; +const order_mask: u32 = 0x7000; +const order_32_little: u32 = 0x2000; +const order_16_little: u32 = 0x1000; +const float_components: u32 = 0x100; + +/// Any single-frame image ImageIO reads as 8-bit RGB, upright by its EXIF orientation and over white where alpha. +pub fn decode(gpa: std.mem.Allocator, bytes: []const u8) Error!Rgb { + if (bytes.len == 0) return error.BadImage; + if (bytes.len > max_bytes) return error.ImageTooLarge; + const data = CFDataCreate(null, bytes.ptr, @intCast(bytes.len)) orelse return error.OutOfMemory; + defer CFRelease(data); + const src = CGImageSourceCreateWithData(data, null) orelse return error.BadImage; + defer CFRelease(src); + const frames = CGImageSourceGetCount(src); + if (frames == 0) return error.BadImage; + if (frames > 1) return error.Animated; // images.py: animated and multipage images are refused + const img = CGImageSourceCreateImageAtIndex(src, 0, null) orelse return error.BadImage; + defer CFRelease(img); + const w = CGImageGetWidth(img); + const h = CGImageGetHeight(img); + if (w == 0 or h == 0) return error.BadImage; + if (w > max_side or h > max_side or w * h > max_pixels) return error.ImageTooLarge; + const out = try gpa.alloc(u8, w * h * 3); + errdefer gpa.free(out); + if (!(try raw(img, w, h, out))) try drawn(img, w, h, out); + const rgb: Rgb = .{ .w = @intCast(w), .h = @intCast(h), .px = out }; + const o = orientation(src); + if (o <= 1 or o > 8) return rgb; + defer rgb.deinit(gpa); + return orient(gpa, rgb, o); +} + +/// The first image's EXIF orientation (1-8; 1 or 0: as stored). +fn orientation(src: CF) u32 { + const props = CGImageSourceCopyPropertiesAtIndex(src, 0, null) orelse return 1; + defer CFRelease(props); + const num = CFDictionaryGetValue(props, kCGImagePropertyOrientation) orelse return 1; + var v: i32 = 1; + if (!CFNumberGetValue(num, 9, &v)) return 1; // kCFNumberIntType + return if (v < 1) 1 else @intCast(v); +} + +/// PIL's exif_transpose: EXIF orientation 2-8 applied, each output pixel read from its source pixel. +pub fn orient(gpa: std.mem.Allocator, src: Rgb, o: u32) Error!Rgb { + const w: usize = src.w; + const h: usize = src.h; + const swap = o >= 5; + const ow = if (swap) h else w; + const oh = if (swap) w else h; + const out = try gpa.alloc(u8, ow * oh * 3); + for (0..oh) |y| for (0..ow) |x| { + const s: [2]usize = switch (o) { + 2 => .{ w - 1 - x, y }, + 3 => .{ w - 1 - x, h - 1 - y }, + 4 => .{ x, h - 1 - y }, + 5 => .{ y, x }, + 6 => .{ y, h - 1 - x }, + 7 => .{ w - 1 - y, h - 1 - x }, + else => .{ w - 1 - y, x }, + }; + @memcpy(out[(y * ow + x) * 3 ..][0..3], src.px[(s[1] * w + s[0]) * 3 ..][0..3]); + }; + return .{ .w = @intCast(ow), .h = @intCast(oh), .px = out }; +} + +/// PIL's paste of a pixel with alpha `a` onto white: DIV255(255 * (255 - a) + c * a). +fn overWhite(c: u8, a: u8) u8 { + const t: u32 = 255 * (255 - @as(u32, a)) + @as(u32, c) * a + 128; + return @intCast(((t >> 8) + t) >> 8); +} + +/// The decoder's own values for 8- and 16-bit gray and RGB layouts (what PIL reads); false for any other layout. +fn raw(img: CF, w: usize, h: usize, out: []u8) Error!bool { + const info = CGImageGetBitmapInfo(img); + if (info & float_components != 0) return false; + const cs = CGImageGetColorSpace(img) orelse return false; + const model = CGColorSpaceGetModel(cs); // 0 monochrome, 1 RGB + if (model != 0 and model != 1) return false; + const bpc = CGImageGetBitsPerComponent(img); + const bpp = CGImageGetBitsPerPixel(img); + if (bpc != 8 and bpc != 16) return false; + const alpha = info & alpha_mask; // 0 none, 1/2 premultiplied last/first, 3/4 last/first, 5/6 skip last/first + if (alpha == 1 or alpha == 2 or alpha > 6) return false; + const colours: usize = if (model == 1) 3 else 1; + const comps: usize = colours + @intFromBool(alpha != 0); + if (bpp != comps * bpc) return false; + const first = alpha == 2 or alpha == 4 or alpha == 6; // the alpha (or skipped) component leads + const order = info & order_mask; + const little = (bpc == 8 and order == order_32_little) or (bpc == 16 and order == order_16_little); + if (bpc == 8 and order != 0 and order != order_32_little) return false; + if (bpc == 8 and little and comps != 4) return false; + const provider = CGImageGetDataProvider(img) orelse return false; + const data = CGDataProviderCopyData(provider) orelse return false; + defer CFRelease(data); + const base = CFDataGetBytePtr(data); + const len: usize = @intCast(CFDataGetLength(data)); + const bpr = CGImageGetBytesPerRow(img); + if (bpr * (h - 1) + w * bpp / 8 > len) return false; + const step = bpc / 8; + for (0..h) |y| { + const row = base + y * bpr; + for (0..w) |x| { + const px = row + x * comps * step; + var c: [4]u8 = undefined; + for (0..comps) |i| { + // a 32-bit little-endian pixel stores its components reversed; a 16-bit one each component's bytes + const slot = if (bpc == 8 and little) comps - 1 - i else i; + const at = px + slot * step; + c[i] = if (step == 1) at[0] else if (little) at[1] else at[0]; // 16 bits: the high byte, as PIL's 8-bit view + } + const lead: usize = @intFromBool(first and alpha != 0); + const o = out[(y * w + x) * 3 ..][0..3]; + if (colours == 3) { + o.* = .{ c[lead], c[lead + 1], c[lead + 2] }; + } else { + o.* = .{ c[lead], c[lead], c[lead] }; + } + if (alpha == 3 or alpha == 4) { // a real alpha band: over white, as images.py pastes it + const a = if (first) c[0] else c[comps - 1]; + for (o) |*v| v.* = overWhite(v.*, a); + } + } + } + return true; +} + +/// Any other layout (palette, CMYK, float, premultiplied): drawn into an RGB context over white. +fn drawn(img: CF, w: usize, h: usize, out: []u8) Error!void { + const cs = CGColorSpaceCreateDeviceRGB() orelse return error.BadImage; + defer CGColorSpaceRelease(cs); + const tmp = std.heap.page_allocator.alloc(u8, w * h * 4) catch return error.OutOfMemory; + defer std.heap.page_allocator.free(tmp); + @memset(tmp, 0); + const ctx = CGBitmapContextCreate(tmp.ptr, w, h, 8, w * 4, cs, 5) orelse return error.BadImage; // RGBX + defer CGContextRelease(ctx); + CGContextSetRGBFillColor(ctx, 1, 1, 1, 1); + CGContextFillRect(ctx, .{ .x = 0, .y = 0, .w = @floatFromInt(w), .h = @floatFromInt(h) }); + CGContextDrawImage(ctx, .{ .x = 0, .y = 0, .w = @floatFromInt(w), .h = @floatFromInt(h) }, img); + for (0..w * h) |i| @memcpy(out[i * 3 ..][0..3], tmp[i * 4 ..][0..3]); +} + +// ---- the processor's geometry ---- + +pub const Geometry = struct { target_h: u32, target_w: u32, content_h: u32, content_w: u32 }; + +fn alignUp(v: u64, f: u64) u64 { + return (v + f - 1) / f * f; +} + +/// smart_resize and _resize_geometry: the padded canvas under `max_tokens` (at least `min_tokens`), the content size. +pub fn geometry(h: u32, w: u32, min_tokens: u32, max_tokens: u32) Error!Geometry { + const factor: u64 = vision.patch * vision.merge; + const frames: u64 = vision.temporal; + const per_token = frames * factor * factor; + const min_px = @as(u64, min_tokens) * per_token; + const max_px = @as(u64, max_tokens) * per_token; + if (min_tokens == 0 or max_tokens < min_tokens) return error.ImageTooSmall; + const H: u64 = h; + const W: u64 = w; + var ah = alignUp(H, factor); + var aw = alignUp(W, factor); + var pixels = frames * ah * aw; + if (pixels < min_px) { + const scale = @sqrt(@as(f64, @floatFromInt(min_px)) / @as(f64, @floatFromInt(frames * H * W))); + ah = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(H)) * scale)))), factor); + aw = alignUp(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(W)) * scale)))), factor); + pixels = frames * ah * aw; + } + if (pixels > max_px) { // the tallest content that fits, by bisection on its height + if (max_px < frames * factor * factor) return error.ImageTooSmall; + var best: [2]u64 = .{ factor, factor }; + var low: i64 = 1; + var high: i64 = @intCast(H); + while (low <= high) { + const ch: u64 = @intCast(@divFloor(low + high, 2)); + const cw: u64 = @max(1, W * ch / H); // math.floor(width * content_h / height) + const cand: [2]u64 = .{ alignUp(ch, factor), alignUp(cw, factor) }; + if (frames * cand[0] * cand[1] <= max_px) { + best = cand; + low = @as(i64, @intCast(ch)) + 1; + } else high = @as(i64, @intCast(ch)) - 1; + } + ah = best[0]; + aw = best[1]; + } + var scale = @min(@as(f64, @floatFromInt(ah)) / @as(f64, @floatFromInt(H)), @as(f64, @floatFromInt(aw)) / @as(f64, @floatFromInt(W))); + if (frames * H * W >= min_px) scale = @min(1.0, scale); + const chf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(H)) * scale)); + const cwf: u64 = @intFromFloat(@floor(@as(f64, @floatFromInt(W)) * scale)); + return .{ .target_h = @intCast(ah), .target_w = @intCast(aw), .content_h = @intCast(@max(1, @min(ah, chf))), .content_w = @intCast(@max(1, @min(aw, cwf))) }; +} + +// ---- PIL's ImagingResample for 8-bit images (Resample.c), bicubic (a = -0.5) ---- + +const precision_bits = 32 - 8 - 2; + +fn bicubic(x0: f64) f64 { + const a = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1; + if (x < 2.0) return (((x - 5) * x + 8) * x - 4) * a; + return 0; +} + +const Coeffs = struct { bounds: []u32, kk: []i32, ksize: usize }; + +/// precompute_coeffs and normalize_coeffs_8bpc: each output pixel's first input pixel, count and fixed-point weights. +fn coeffs(gpa: std.mem.Allocator, in_size: u32, out_size: u32) Error!Coeffs { + const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); + const filterscale = @max(scale, 1.0); + const support = 2.0 * filterscale; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const bounds = try gpa.alloc(u32, @as(usize, out_size) * 2); + errdefer gpa.free(bounds); + const kk = try gpa.alloc(i32, @as(usize, out_size) * ksize); + errdefer gpa.free(kk); + var pre = try gpa.alloc(f64, ksize); + defer gpa.free(pre); + const ss = 1.0 / filterscale; + for (0..out_size) |xx| { + const center = (@as(f64, @floatFromInt(xx)) + 0.5) * scale; + var xmin_i: i64 = @intFromFloat(center - support + 0.5); // (int) truncates toward zero, as C does + if (xmin_i < 0) xmin_i = 0; + var xmax_i: i64 = @intFromFloat(center + support + 0.5); + if (xmax_i > in_size) xmax_i = in_size; + xmax_i -= xmin_i; + const xmin: usize = @intCast(xmin_i); + const xmax: usize = @intCast(xmax_i); + var ww: f64 = 0; + for (0..xmax) |x| { + const w = bicubic((@as(f64, @floatFromInt(x + xmin)) - center + 0.5) * ss); + pre[x] = w; + ww += w; + } + for (0..ksize) |x| { + const k = if (x < xmax) (if (ww != 0) pre[x] / ww else pre[x]) else 0; + const v = k * @as(f64, @floatFromInt(@as(i64, 1) << precision_bits)); + kk[xx * ksize + x] = @intFromFloat(if (v < 0) v - 0.5 else v + 0.5); // C's (int) of the rounded value + } + bounds[xx * 2] = @intCast(xmin); + bounds[xx * 2 + 1] = @intCast(xmax); + } + return .{ .bounds = bounds, .kk = kk, .ksize = ksize }; +} + +fn clip8(v: i64) u8 { + const s = v >> precision_bits; + return if (s < 0) 0 else if (s > 255) 255 else @intCast(s); +} + +/// Image.resize((w, h), BICUBIC) on an 8-bit RGB image: horizontal pass, then vertical, each rounded to 8 bits. +pub fn resize(gpa: std.mem.Allocator, src: Rgb, w: u32, h: u32) Error!Rgb { + var cur = src; + var owned = false; + errdefer if (owned) cur.deinit(gpa); + if (w != src.w) { // horizontal (PIL limits it to the rows the vertical pass reads: all of them here) + const c = try coeffs(gpa, src.w, w); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, w) * cur.h * 3); + for (0..cur.h) |y| for (0..w) |xx| { + const xmin = c.bounds[xx * 2]; + const xmax = c.bounds[xx * 2 + 1]; + const k = c.kk[xx * c.ksize ..]; + for (0..3) |ch| { + var ss: i64 = 1 << (precision_bits - 1); + for (0..xmax) |x| ss += @as(i64, cur.px[(y * cur.w + x + xmin) * 3 + ch]) * k[x]; + out[(y * w + xx) * 3 + ch] = clip8(ss); + } + }; + if (owned) cur.deinit(gpa); + cur = .{ .w = w, .h = cur.h, .px = out }; + owned = true; + } + if (h != src.h) { + const c = try coeffs(gpa, src.h, h); + defer gpa.free(c.bounds); + defer gpa.free(c.kk); + const out = try gpa.alloc(u8, @as(usize, cur.w) * h * 3); + for (0..h) |yy| { + const ymin = c.bounds[yy * 2]; + const ymax = c.bounds[yy * 2 + 1]; + const k = c.kk[yy * c.ksize ..]; + for (0..cur.w) |x| for (0..3) |ch| { + var ss: i64 = 1 << (precision_bits - 1); + for (0..ymax) |y| ss += @as(i64, cur.px[((y + ymin) * cur.w + x) * 3 + ch]) * k[y]; + out[(yy * cur.w + x) * 3 + ch] = clip8(ss); + }; + } + if (owned) cur.deinit(gpa); + cur = .{ .w = cur.w, .h = h, .px = out }; + owned = true; + } + if (!owned) cur = .{ .w = src.w, .h = src.h, .px = try gpa.dupe(u8, src.px) }; + return cur; +} + +/// The resized content padded with black, scaled, normalized and cut into 2x2 windows of 14x14 patches, row-major. +pub fn patchify(gpa: std.mem.Allocator, content: Rgb, g: Geometry, mean: [3]f32, sd: [3]f32) Error![]f32 { + const P = vision.patch; + const M = vision.merge; + const gh = g.target_h / P; + const gw = g.target_w / P; + const out = try gpa.alloc(f32, @as(usize, gh) * gw * vision.patch_width); + const rescale: f32 = @floatCast(@as(f64, 1.0) / 255.0); + var row: usize = 0; + for (0..gh / M) |wy| for (0..gw / M) |wx| for (0..M) |my| for (0..M) |mx| { + const py0 = (wy * M + my) * P; + const px0 = (wx * M + mx) * P; + const dst = out[row * vision.patch_width ..][0..vision.patch_width]; + for (0..3) |ch| for (0..vision.temporal) |t| for (0..P) |dy| for (0..P) |dx| { + const y = py0 + dy; + const x = px0 + dx; + const v: u8 = if (y < content.h and x < content.w) content.px[(y * content.w + x) * 3 + ch] else 0; + const f = @as(f32, @floatFromInt(v)) * rescale; + dst[((ch * vision.temporal + t) * P + dy) * P + dx] = (f - mean[ch]) / sd[ch]; + }; + row += 1; + }; + return out; +} + +/// Bytes to patches: decoded, sized for `max_tokens` (at least min(16, max_tokens)), resized, padded and cut. +pub fn prepare(gpa: std.mem.Allocator, bytes: []const u8, c: vision.Config, max_tokens: u32) Error!Prepared { + const img = try decode(gpa, bytes); + defer img.deinit(gpa); + return prepareRgb(gpa, img, std.hash.Wyhash.hash(0x1a6e, bytes), c, max_tokens); +} + +pub fn prepareRgb(gpa: std.mem.Allocator, img: Rgb, hash: u64, c: vision.Config, max_tokens: u32) Error!Prepared { + const g = try geometry(img.h, img.w, @min(c.min_tokens, max_tokens), max_tokens); + const content = try resize(gpa, img, g.content_w, g.content_h); + defer content.deinit(gpa); + const pixels = try patchify(gpa, content, g, c.mean, c.std); + const gh = g.target_h / vision.patch; + const gw = g.target_w / vision.patch; + return .{ .pixels = pixels, .gh = gh, .gw = gw, .tokens = gh * gw / vision.patches_per_token, .hash = hash }; +} + +test "geometry follows smart_resize" { + // the reference dumps' grids: 517x389 pads to 532x392 unscaled; 96x96 grows to 112; 20x15 to the 16-token floor + try std.testing.expectEqual(Geometry{ .target_h = 392, .target_w = 532, .content_h = 389, .content_w = 517 }, try geometry(389, 517, 16, 4096)); + try std.testing.expectEqual(Geometry{ .target_h = 112, .target_w = 112, .content_h = 112, .content_w = 112 }, try geometry(96, 96, 16, 4096)); + const tiny = try geometry(15, 20, 16, 4096); + try std.testing.expectEqual(@as(u32, 8 * 14), tiny.target_h); + try std.testing.expectEqual(@as(u32, 10 * 14), tiny.target_w); + // a large image under a small budget: the canvas fits it + const big = try geometry(1500, 2000, 16, 256); + try std.testing.expect(@as(u64, big.target_h) * big.target_w * 2 <= 256 * 2 * 28 * 28); + try std.testing.expect(big.content_h <= big.target_h and big.content_w <= big.target_w); +} + +test "orientations move pixels as PIL's transposes do" { + const gpa = std.testing.allocator; + // a 3x2 image of pixel indices: 0 1 2 / 3 4 5 + var px: [18]u8 = undefined; + for (0..6) |i| @memset(px[i * 3 ..][0..3], @intCast(i)); + const src: Rgb = .{ .w = 3, .h = 2, .px = &px }; + const want = [_][]const u8{ &.{ 2, 1, 0, 5, 4, 3 }, &.{ 5, 4, 3, 2, 1, 0 }, &.{ 3, 4, 5, 0, 1, 2 }, &.{ 0, 3, 1, 4, 2, 5 }, &.{ 3, 0, 4, 1, 5, 2 }, &.{ 5, 2, 4, 1, 3, 0 }, &.{ 2, 5, 1, 4, 0, 3 } }; + for (want, 2..) |w, o| { + const out = try orient(gpa, src, @intCast(o)); + defer out.deinit(gpa); + for (w, 0..) |v, i| try std.testing.expectEqual(v, out.px[i * 3]); + } +} + +test "alpha over white follows PIL's DIV255" { + try std.testing.expectEqual(@as(u8, 255), overWhite(0, 0)); + try std.testing.expectEqual(@as(u8, 10), overWhite(10, 255)); + try std.testing.expectEqual(@as(u8, 132), overWhite(10, 128)); // DIV255(255 * 127 + 10 * 128) +} + +test "resize keeps a flat image flat and its size" { + const gpa = std.testing.allocator; + const px = try gpa.alloc(u8, 7 * 5 * 3); + defer gpa.free(px); + @memset(px, 200); + const out = try resize(gpa, .{ .w = 7, .h = 5, .px = px }, 13, 3); + defer out.deinit(gpa); + try std.testing.expectEqual(@as(u32, 13), out.w); + for (out.px) |v| try std.testing.expectEqual(@as(u8, 200), v); +} diff --git a/zig/src/families/glm/mtp.zig b/zig/src/families/glm/mtp.zig index 2f1234c1fb..ee88c7e7f8 100644 --- a/zig/src/families/glm/mtp.zig +++ b/zig/src/families/glm/mtp.zig @@ -35,6 +35,7 @@ pub fn absorb(x: *fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, next: Ref, rows: u32, const L = &x.w.layers[c.layers]; const m = x.w.mtp.?; fwd.embedRows(x, e, next, sc.m_emb, rows); + fwd.inject(x, e, sc.m_emb, rows, pos + 1); // the head reads each row's next token: an image's rows one place on fwd.rms(x, e, sc.m_emb, m.enorm, sc.m_eh, rows, D, D, 2 * D, c.eps); fwd.rms(x, e, h, m.hnorm, sc.m_eh.at(@as(usize, D) * 2), rows, D, D, 2 * D, c.eps); if (m.eh_bf16) |eh| { // stored unquantized (mlx-lm's): fp32 sums of its bf16 rows diff --git a/zig/src/families/glm/prompt.zig b/zig/src/families/glm/prompt.zig index 4f73d916b0..cb7134e4f3 100644 --- a/zig/src/families/glm/prompt.zig +++ b/zig/src/families/glm/prompt.zig @@ -470,7 +470,7 @@ pub fn backbone(p: *const Prompt, x: *fwd.Ctx, e: mtl.ComputeEncoder, ids: Ref, std.debug.assert(x.sc == ss and M <= p.rows); const s = x.skip; // classes a profile leaves out const Class = fwd.Class; - if (s & Class.ends == 0) fwd.embed(x, e, ids, M); + if (s & Class.ends == 0) fwd.embed(x, e, ids, M, pos); const plane = @as(usize, M) * c.hidden * 2; // a trace's capture points: the decode backbone's fwd.snap(x, e, ss.h, plane); var pending = false; @@ -532,6 +532,7 @@ pub fn mtp(p: *const Prompt, x: *const fwd.Ctx, e: mtl.ComputeEncoder, h: Ref, n const L = &x.w.layers[c.layers]; const m = x.w.mtp.?; fwd.embedRows(x, e, next, p.m_emb, M); + fwd.inject(x, e, p.m_emb, M, pos + 1); // the head reads each row's next token: an image's rows one place on fwd.rms(x, e, p.m_emb, m.enorm, p.m_eh, M, D, D, 2 * D, c.eps); fwd.rms(x, e, h, m.hnorm, p.m_eh.at(@as(usize, D) * 2), M, D, D, 2 * D, c.eps); if (m.eh_bf16) |eh| wide.bf16(&p.mm.wide.?, e, eh, D, 2 * D, .{ .in = p.m_eh, .ldx = 2 * D, .out = p.m_x, .ldo = D, .rows = M }) else qmm(p, e, p.m_eh, m.eh_proj, p.m_x, M); diff --git a/zig/src/families/glm/slot_images.zig b/zig/src/families/glm/slot_images.zig new file mode 100644 index 0000000000..8b41b53d05 --- /dev/null +++ b/zig/src/families/glm/slot_images.zig @@ -0,0 +1,35 @@ +//! A stream's images through the tower into the engine's image rows, injected over placeholders by its prompt chunks. +const std = @import("std"); +const mtl = @import("metal"); +const vision_mod = @import("vision.zig"); +const Slots = @import("slots.zig").Slots; +const Image = @import("lanes").stream.Image; + +/// Slot `i`'s images into the engine's image rows, one image a command buffer: the tower's host scratch is shared. +pub fn encode(sl: *Slots, i: u32, imgs: []const Image) !void { + const e = sl.e; + const v = e.vision orelse return error.NoVision; + if (imgs.len > e.vision_limits.max_images) return error.TooManyImages; + if (i >= sl.slots.len or !sl.slots[i].used) return error.SlotOutOfStep; + v.spans.clearRetainingCapacity(); + try sl.flush(); + e.sync(); + var row: u32 = 0; + for (imgs) |img| { + if (row + img.tokens > e.vision_limits.image_tokens or img.gh * img.gw != img.tokens * vision_mod.patches_per_token) return error.ImageTooLarge; + if (img.at + img.tokens > sl.slots[i].prompt_len) return error.ImageOutOfPrompt; + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const b = e.begin(); + try v.encode(b.enc, img.pixels, img.gh, img.gw, e.vision_rows.at(@as(usize, row) * e.c.hidden * 2), null); + try e.finish(b.cb, b.enc); + try v.spans.append(v.gpa, .{ .at = img.at, .n = img.tokens, .row = row }); + row += img.tokens; + } + v.inject = .{ .slot = i, .v = .{ .rows = e.vision_rows, .spans = v.spans.items } }; +} + +/// The stream's prompt pass is over: the chunks after it inject nothing. +pub fn end(sl: *Slots) void { + if (sl.e.vision) |v| v.inject = null; +} diff --git a/zig/src/families/glm/slots.zig b/zig/src/families/glm/slots.zig index 90aa566962..98a69baf7d 100644 --- a/zig/src/families/glm/slots.zig +++ b/zig/src/families/glm/slots.zig @@ -219,6 +219,7 @@ pub const Slots = struct { const pool = mtl.objc.Pool.push(); defer pool.pop(); var x = sl.ctx(slot); + x.inject = if (e.vision) |v| v.injected(i) else null; // an image prompt's rows over its placeholders (slot_images.zig) const b = e.begin(); const ids = e.prompt_ids.at(@as(usize, at) * 4); const next = e.prompt_ids.at(@as(usize, at + 1) * 4); diff --git a/zig/src/families/glm/video.zig b/zig/src/families/glm/video.zig new file mode 100644 index 0000000000..84f03754c5 --- /dev/null +++ b/zig/src/families/glm/video.zig @@ -0,0 +1,398 @@ +//! A video's frames to GLM5-Next vision patches as transformers 5.19's Glm5NextVideoProcessor makes them from PyAV's. +const std = @import("std"); +const vision = @import("vision.zig"); + +pub const Error = error{ VideoTooShort, VideoTooLarge, BadVideo, OutOfMemory }; + +/// The processor's frame rate and sampling cap. +pub const sample_fps = 2.0; + +/// A decoded video's facts (the PyAV loader's VideoMetadata) and its frames on request. +pub const Source = struct { + frames: u64, // the stream's frame count + rate_num: u64, // its average frame rate, num / den + rate_den: u64, + width: u32, + height: u32, + ctx: *anyopaque, + /// The frames at `indices` (ascending, distinct), each handed to `sink` in order as RGB rows `stride` bytes apart. + decode: *const fn (ctx: *anyopaque, indices: []const u32, sink: *Sink) anyerror!void, + + fn fps(s: Source) f64 { + return @as(f64, @floatFromInt(s.rate_num)) / @as(f64, @floatFromInt(s.rate_den)); + } + + /// total / fps as the loader computes it from the Fraction: the exact ratio, rounded once. + fn duration(s: Source) f64 { + return @as(f64, @floatFromInt(s.frames * s.rate_den)) / @as(f64, @floatFromInt(s.rate_num)); + } +}; + +/// Where decoded frames go: each resized into its slot of the canvas. +pub const Sink = struct { + plan: *const Plan, + frames: []u8, // [n, content_h, content_w, 3] + scratch: std.mem.Allocator, + got: usize = 0, + + /// Frame `k` of the plan, `w` x `h` RGB rows `stride` bytes apart. + pub fn take(s: *Sink, k: usize, rgb: []const u8, w: u32, h: u32, stride: usize) Error!void { + const p = s.plan; + if (k != s.got or k >= p.indices.len or w != p.width or h != p.height) return error.BadVideo; + const dst = s.frames[k * p.frameBytes() ..][0..p.frameBytes()]; + try resizeInto(s.scratch, rgb, w, h, stride, p.content_w, p.content_h, dst); + s.got += 1; + } +}; + +/// What a video becomes before any frame is decoded: the frames to decode and the canvas they land on. +pub const Plan = struct { + indices: []u32, // distinct, ascending: the frames decoded + times: []f64, // one a pair: the pair's first sampled frame's index over the rate (metadata.timestamps[::2]) + width: u32, + height: u32, + target_h: u32, + target_w: u32, + content_h: u32, + content_w: u32, + + fn frameBytes(p: Plan) usize { + return @as(usize, p.content_h) * p.content_w * 3; + } + + pub fn pairs(p: Plan) usize { + return (p.indices.len + 1) / 2; + } + + pub fn tokensPerPair(p: Plan) u32 { + return (p.target_h / vision.patch) * (p.target_w / vision.patch) / vision.patches_per_token; + } +}; + +/// numpy.linspace(start, stop, num, dtype=int): the float steps, the last exactly stop, truncated toward zero. +fn linspace(a: std.mem.Allocator, start: f64, stop: f64, num: usize) ![]u32 { + const out = try a.alloc(u32, num); + const step = if (num > 1) (stop - start) / @as(f64, @floatFromInt(num - 1)) else 0; + for (out, 0..) |*v, i| { + var y = if (num > 1 and step == 0) start else @as(f64, @floatFromInt(i)) * step + start; + if (num > 1 and i == num - 1) y = stop; + v.* = @intFromFloat(@trunc(y)); + } + return out; +} + +/// Python's round(): halves to even. +fn pyRound(x: f64) f64 { + if (@abs(x - @trunc(x)) == 0.5) return 2.0 * @round(x / 2.0); + return @round(x); +} + +/// Glm5NextVideoProcessor.sample_frames on the loader's metadata: distinct indices (its even-count pad decoded once). +pub fn sample(a: std.mem.Allocator, s: Source, max_frames: u32) Error![]u32 { + if (s.frames == 0 or s.rate_num == 0 or s.rate_den == 0) return error.BadVideo; + const total = s.frames; + const fps = s.fps(); + var duration = s.duration(); + if (duration == 0) duration = pyRound(@as(f64, @floatFromInt(total - 1)) / fps) + 1; + const max_seconds: f64 = @trunc(duration); + var extract_t: u64 = @intFromFloat(@trunc(duration * sample_fps)); + extract_t = @min(extract_t, max_frames); + if (extract_t == 0) return error.VideoTooShort; // the processor samples nothing and the loader fails + var picked: std.ArrayList(u32) = .empty; + defer picked.deinit(a); + var indices: []u32 = undefined; + if (total < extract_t) { + indices = try linspace(a, 0, @floatFromInt(total - 1), extract_t); + } else { + const per_frame = 1.0 / fps; + const inv: f64 = 1.0 / sample_fps; + var current: f64 = 0; + for (0..total) |i| { + if (@as(f64, @floatFromInt(i)) * per_frame >= current) { + current += inv; + try picked.append(a, @intCast(i)); + if (current >= max_seconds) break; + } + } + if (picked.items.len < extract_t) { + const start: f64 = if (picked.items.len == 0) 0 else @floatFromInt(picked.items[0]); + const end: f64 = if (picked.items.len == 0) @floatFromInt(total - 1) else @floatFromInt(picked.items[picked.items.len - 1]); + indices = try linspace(a, start, end, extract_t); + } else if (picked.items.len > extract_t) { + indices = try linspace(a, 0, @floatFromInt(total - 1), extract_t); + } else indices = try a.dupe(u32, picked.items); + } + // the order kept, repeats dropped (linspace's are adjacent, the loop's ascending) + var n: usize = 0; + for (indices) |v| { + if (n > 0 and indices[n - 1] == v) continue; + indices[n] = v; + n += 1; + } + return indices[0..n]; +} + +/// The processor's smart_resize: a canvas in multiples of 28, aspect kept, whose frame pairs fit min..max tokens. +pub fn canvas(frames: u64, height: u32, width: u32, min_tokens: u64, max_tokens: u64) Error![2]u32 { + const factor: u64 = vision.patch * vision.merge; + const per_token = vision.temporal * factor * factor; + const min_pixels = min_tokens * per_token; + const max_pixels = max_tokens * per_token; + const align_ = struct { + fn f(v: u64, by: u64) u64 { + return std.math.divCeil(u64, v, by) catch unreachable; + } + }.f; + const frames_f: f64 = @floatFromInt(frames); + const aligned_frames: u64 = @max(vision.temporal, @as(u64, @intFromFloat(pyRound(frames_f / vision.temporal))) * vision.temporal); + var ah = align_(height, factor) * factor; + var aw = align_(width, factor) * factor; + if (aligned_frames * ah * aw < min_pixels) { + const scale = @sqrt(@as(f64, @floatFromInt(min_pixels)) / (frames_f * @as(f64, @floatFromInt(height)) * @as(f64, @floatFromInt(width)))); + ah = align_(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(height)) * scale)))), factor) * factor; + aw = align_(@max(1, @as(u64, @intFromFloat(@ceil(@as(f64, @floatFromInt(width)) * scale)))), factor) * factor; + } + if (aligned_frames * ah * aw > max_pixels) { + if (max_pixels < aligned_frames * factor * factor) return error.VideoTooLarge; + var low: i64 = 1; + var high: i64 = height; + var best_h: u64 = factor; + var best_w: u64 = factor; + while (low <= high) { + const ch: u64 = @intCast(@divFloor(low + high, 2)); + const cw: u64 = @max(1, @as(u64, @intFromFloat(@floor(@as(f64, @floatFromInt(width)) * @as(f64, @floatFromInt(ch)) / @as(f64, @floatFromInt(height)))))); + const th = align_(ch, factor) * factor; + const tw = align_(cw, factor) * factor; + if (aligned_frames * th * tw <= max_pixels) { + best_h = th; + best_w = tw; + low = @as(i64, @intCast(ch)) + 1; + } else high = @as(i64, @intCast(ch)) - 1; + } + ah = best_h; + aw = best_w; + } + return .{ @intCast(ah), @intCast(aw) }; +} + +/// What `s` becomes at most `max_tokens` tokens and `max_frames` sampled frames: the frames to decode and the canvas. +pub fn plan(a: std.mem.Allocator, s: Source, c: vision.Config, max_tokens: u32, max_frames: u32) Error!Plan { + const indices = try sample(a, s, max_frames); + const n: u64 = indices.len; // the frames decoded: the processor's num_frames + const t = try canvas(n, s.height, s.width, @min(c.min_tokens, max_tokens), max_tokens); + const per_token: u64 = vision.temporal * (vision.patch * vision.merge) * (vision.patch * vision.merge); + const hf: f64 = @floatFromInt(s.height); + const wf: f64 = @floatFromInt(s.width); + var scale = @min(@as(f64, @floatFromInt(t[0])) / hf, @as(f64, @floatFromInt(t[1])) / wf); + if (n * s.height * s.width >= per_token * c.min_tokens) scale = @min(1.0, scale); + const content_h: u32 = @intCast(@max(1, @min(t[0], @as(u64, @intFromFloat(@floor(hf * scale)))))); + const content_w: u32 = @intCast(@max(1, @min(t[1], @as(u64, @intFromFloat(@floor(wf * scale)))))); + // metadata.timestamps[::2] over the sampled indices, the even-making duplicate among them + const pairs = (indices.len + 1) / 2; + const times = try a.alloc(f64, pairs); + for (times, 0..) |*v, g| v.* = @as(f64, @floatFromInt(indices[2 * g])) / s.fps(); + return .{ .indices = indices, .times = times, .width = s.width, .height = s.height, .target_h = t[0], .target_w = t[1], .content_h = content_h, .content_w = content_w }; +} + +/// One pair of frames ready for the tower, and the video's. +pub const Pair = struct { pixels: []f32, gh: u32, gw: u32, tokens: u32, hash: u64 }; +pub const Prepared = struct { groups: []Pair, times: []f64 }; + +/// `s` decoded, resized as it decodes, then normalized and cut into pairs (allocated in `a`). +pub fn prepare(a: std.mem.Allocator, s: Source, c: vision.Config, max_tokens: u32, max_frames: u32) anyerror!Prepared { + const p = try plan(a, s, c, max_tokens, max_frames); + const frames = try a.alloc(u8, p.indices.len * p.frameBytes()); + var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator); + defer arena.deinit(); + var sink: Sink = .{ .plan = &p, .frames = frames, .scratch = arena.allocator() }; + try s.decode(s.ctx, p.indices, &sink); + if (sink.got != p.indices.len) return error.BadVideo; + return .{ .groups = try patchify(a, p, frames, c.mean, c.std), .times = p.times }; +} + +/// Frames padded black to the canvas, normalized as the fused rescale, odd last repeated, pairs cut into 2x2 windows. +pub fn patchify(a: std.mem.Allocator, p: Plan, frames: []const u8, mean: [3]f32, sd: [3]f32) Error![]Pair { + const P = vision.patch; + const M = vision.merge; + const gh = p.target_h / P; + const gw = p.target_w / P; + var m255: [3]f32 = undefined; + var s255: [3]f32 = undefined; + const by: f32 = @floatCast(@as(f64, 1.0) / (@as(f64, 1.0) / 255.0)); + for (0..3) |ch| { + m255[ch] = mean[ch] * by; + s255[ch] = sd[ch] * by; + } + const n = p.indices.len; + const out = try a.alloc(Pair, p.pairs()); + for (out, 0..) |*pair, g| { + const px = try a.alloc(f32, @as(usize, gh) * gw * vision.patch_width); + var row: usize = 0; + for (0..gh / M) |wy| for (0..gw / M) |wx| for (0..M) |my| for (0..M) |mx| { + const py0 = (wy * M + my) * P; + const px0 = (wx * M + mx) * P; + const dst = px[row * vision.patch_width ..][0..vision.patch_width]; + for (0..3) |ch| for (0..vision.temporal) |t| { + const f = @min(2 * g + t, n - 1); + const frame = frames[f * p.frameBytes() ..][0..p.frameBytes()]; + for (0..P) |dy| for (0..P) |dx| { + const y = py0 + dy; + const x = px0 + dx; + const v: u8 = if (y < p.content_h and x < p.content_w) frame[(y * p.content_w + x) * 3 + ch] else 0; + dst[((ch * vision.temporal + t) * P + dy) * P + dx] = (@as(f32, @floatFromInt(v)) - m255[ch]) / s255[ch]; + }; + }; + row += 1; + }; + var h = std.hash.Wyhash.init(0x71de0); + h.update(std.mem.sliceAsBytes(px)); + pair.* = .{ .pixels = px, .gh = gh, .gw = gw, .tokens = gh * gw / vision.patches_per_token, .hash = h.final() }; + } + return out; +} + +// ---- torch's antialiased bicubic on uint8 (aten UpSampleKernel's separable int16 path, as torchvision calls it) ---- + +fn cubic(x0: f64) f64 { + const a = -0.5; + const x = @abs(x0); + if (x < 1.0) return ((a + 2.0) * x - (a + 3.0)) * x * x + 1; + if (x < 2.0) return (((x - 5) * x + 8) * x - 4) * a; + return 0.0; +} + +const Coeffs = struct { ksize: usize, bounds: []usize, kk: []i32, precision: u5 }; + +/// The axis's weights (PIL's precompute_coeffs) as int16 fixed point, the most precision under 22 bits that fits. +fn coeffs(a: std.mem.Allocator, in_size: usize, out_size: usize) Error!Coeffs { + const scale = @as(f64, @floatFromInt(in_size)) / @as(f64, @floatFromInt(out_size)); + const filterscale = @max(scale, 1.0); + const support = 2.0 * filterscale; + const ksize: usize = @as(usize, @intFromFloat(@ceil(support))) * 2 + 1; + const bounds = try a.alloc(usize, out_size * 2); + const kk = try a.alloc(i32, out_size * ksize); + const pre_all = try a.alloc(f64, out_size * ksize); + defer a.free(pre_all); + @memset(pre_all, 0); + for (0..out_size) |xx| { + const pre = pre_all[xx * ksize ..][0..ksize]; + const center = (@as(f64, @floatFromInt(xx)) + 0.5) * scale; + const ss = 1.0 / filterscale; + const xmin_i: i64 = @max(0, @as(i64, @intFromFloat(center - support + 0.5))); + const xmax_i: i64 = @min(@as(i64, @intCast(in_size)), @as(i64, @intFromFloat(center + support + 0.5))); + const xmin: usize = @intCast(xmin_i); + const xmax: usize = @intCast(xmax_i - xmin_i); + var ww: f64 = 0; + for (0..xmax) |x| { + pre[x] = cubic((@as(f64, @floatFromInt(x + xmin)) - center + 0.5) * ss); + ww += pre[x]; + } + if (ww != 0) for (0..xmax) |x| { + pre[x] /= ww; + }; + bounds[xx * 2] = xmin; + bounds[xx * 2 + 1] = xmax; + } + var w_max: f64 = 0; + for (pre_all) |v| w_max = @max(w_max, @abs(v)); + var precision: u5 = 0; + while (precision < 22) : (precision += 1) { + const next: i64 = @intFromFloat(0.5 + w_max * @as(f64, @floatFromInt(@as(i64, 1) << (precision + 1)))); + if (next >= (1 << 15)) break; + } + const one: f64 = @floatFromInt(@as(i64, 1) << precision); + for (pre_all, 0..) |v, i| kk[i] = @intFromFloat(if (v < 0) -0.5 + v * one else 0.5 + v * one); + return .{ .ksize = ksize, .bounds = bounds, .kk = kk, .precision = precision }; +} + +fn clip8(v: i64, precision: u5) u8 { + return @intCast(std.math.clamp(v >> precision, 0, 255)); +} + +/// `src` (w x h RGB, rows `stride` apart) resized to ow x oh into `dst`, horizontal then vertical, a kept axis copied. +pub fn resizeInto(a: std.mem.Allocator, src: []const u8, w: u32, h: u32, stride: usize, ow: u32, oh: u32, dst: []u8) Error!void { + // the horizontal pass (or the rows as they are) into a packed w' x h buffer + const mid_w: usize = ow; + const mid = if (ow == w and oh == h) dst else try a.alloc(u8, mid_w * h * 3); + if (ow == w) { + for (0..h) |y| @memcpy(mid[y * mid_w * 3 ..][0 .. mid_w * 3], src[y * stride ..][0 .. mid_w * 3]); + } else { + const c = try coeffs(a, w, ow); + for (0..h) |y| for (0..ow) |x| { + const xmin = c.bounds[x * 2]; + const xmax = c.bounds[x * 2 + 1]; + const k = c.kk[x * c.ksize ..]; + const row = src[y * stride ..]; + for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..xmax) |j| ss += @as(i64, row[(xmin + j) * 3 + ch]) * k[j]; + mid[(y * mid_w + x) * 3 + ch] = clip8(ss, c.precision); + } + }; + } + if (oh == h) { + if (mid.ptr != dst.ptr) @memcpy(dst[0 .. mid_w * h * 3], mid[0 .. mid_w * h * 3]); + return; + } + const c = try coeffs(a, h, oh); + for (0..oh) |y| { + const ymin = c.bounds[y * 2]; + const ymax = c.bounds[y * 2 + 1]; + const k = c.kk[y * c.ksize ..]; + for (0..mid_w) |x| for (0..3) |ch| { + var ss: i64 = @as(i64, 1) << (c.precision - 1); + for (0..ymax) |j| ss += @as(i64, mid[((ymin + j) * mid_w + x) * 3 + ch]) * k[j]; + dst[(y * mid_w + x) * 3 + ch] = clip8(ss, c.precision); + }; + } +} + +/// Python's f"{t:.1f} seconds": one decimal, an exact tie rounded to even. +pub fn seconds(buf: []u8, t: f64) ![]const u8 { + // t * 10 exactly from the double's bits, rounded once, ties to even (no 128-bit floats: no compiler-rt is linked) + const bits: u64 = @bitCast(@abs(t)); + const biased: i32 = @intCast(bits >> 52); + const m: u64 = (bits & ((@as(u64, 1) << 52) - 1)) | (if (biased == 0) 0 else @as(u64, 1) << 52); + const e: i32 = (if (biased == 0) 1 else biased) - 1075; + const x = m * 10; // under 2^57 + var r: u64 = 0; + if (e >= 0) { + r = x << @intCast(@min(e, 6)); // times past 2^56 seconds are not times + } else if (e > -64) { + const s: u6 = @intCast(-e); + const q = x >> s; + const rem = x & ((@as(u64, 1) << s) - 1); + const half = @as(u64, 1) << (s - 1); + r = q + @intFromBool(rem > half or (rem == half and q & 1 == 1)); + } + return std.fmt.bufPrint(buf, "{d}.{d} seconds", .{ r / 10, r % 10 }); +} + +test "sampling as Glm5NextVideoProcessor.sample_frames, then the loader's distinct frames" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var dummy: u8 = 0; + const src = struct { + fn of(frames: u64, num: u64, den: u64, ctx: *anyopaque) Source { + return .{ .frames = frames, .rate_num = num, .rate_den = den, .width = 64, .height = 48, .ctx = ctx, .decode = undefined }; + } + }.of; + // 80 frames at 25: 3.2 s, 6 frames at 0, 0.5, 1, 1.5, 2, 2.5 s (the reference run's indices) + try std.testing.expectEqualSlices(u32, &.{ 0, 13, 25, 38, 50, 63 }, try sample(a, src(80, 25, 1, &dummy), 256)); + // 10 frames at 25 (0.4 s): nothing to sample + try std.testing.expectError(error.VideoTooShort, sample(a, src(10, 25, 1, &dummy), 256)); + // a long video: 256 frames over the whole length + const long = try sample(a, src(1500, 10, 1, &dummy), 256); + try std.testing.expectEqual(@as(usize, 256), long.len); + try std.testing.expectEqual(@as(u32, 1499), long[255]); + try std.testing.expectError(error.BadVideo, sample(a, src(0, 25, 1, &dummy), 256)); +} + +test "seconds as Python's f\"{t:.1f}\"" { + var buf: [32]u8 = undefined; + try std.testing.expectEqualStrings("0.2 seconds", try seconds(&buf, 0.25)); + try std.testing.expectEqualStrings("0.8 seconds", try seconds(&buf, 0.75)); + try std.testing.expectEqualStrings("0.5 seconds", try seconds(&buf, 0.52)); + try std.testing.expectEqualStrings("12.0 seconds", try seconds(&buf, 11.96)); +} diff --git a/zig/src/families/glm/vision.zig b/zig/src/families/glm/vision.zig new file mode 100644 index 0000000000..72bb525075 --- /dev/null +++ b/zig/src/families/glm/vision.zig @@ -0,0 +1,537 @@ +//! GLM-5.3-Flash's image tower on the GPU (mlx-vlm's glm5_next/vision.py in bf16): patches to `<|image|>` rows. +const std = @import("std"); +const mtl = @import("metal"); +const sources = @import("kernel_sources"); +const frags = @import("../../core/frags.zig"); +const st = @import("../../core/safetensors.zig"); +const wts = @import("weights.zig"); +const fwd = @import("forward.zig"); +const readAll = @import("weights_fix.zig").readAll; +const Engine = @import("engine.zig").Engine; +const Ref = wts.Ref; + +// The tower this file reads; `Config.parse` refuses a checkpoint whose vision_config differs. +pub const hidden = 1024; +pub const heads = 16; +pub const head_dim = 64; +pub const depth = 24; +pub const mlp_inter = 4096; +pub const out_hidden = 4096; // the language model's width +pub const merger_inter = 10240; +pub const patch = 14; +pub const temporal = 2; +pub const merge = 2; +pub const patch_width = 3 * temporal * patch * patch; // 1176: a patch's values, channel by frame by row by column +pub const patches_per_token = merge * merge; + +/// --vision's limits: image tokens a prompt may hold in all (scratch and rows are sized for them), and images. +pub const Limits = struct { image_tokens: u32 = 4096, max_images: u32 = 4 }; + +pub const Config = struct { + eps: f32 = 1e-5, + limit: f32 = 10, // swiglu_limit + image_token: u32, + image_start: u32, + image_end: u32, + video: ?[3]u32 = null, // video_start_token_id, video_token_id, video_end_token_id: the checkpoint reads videos + mean: [3]f32, + std: [3]f32, + min_tokens: u32, + max_tokens: u32, + + /// config.json's vision_config and image token ids, processor_config.json's image processor; null: no tower. + pub fn parse(arena: std.mem.Allocator, config_json: []const u8, processor_json: ?[]const u8) !?Config { + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, config_json, .{}); + const vc = doc.object.get("vision_config") orelse return null; + if (vc != .object) return null; + const want = .{ .{ "depth", depth }, .{ "hidden_size", hidden }, .{ "num_heads", heads }, .{ "intermediate_size", mlp_inter }, .{ "out_hidden_size", out_hidden }, .{ "projection_intermediate_size", merger_inter }, .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "spatial_merge_size", merge }, .{ "in_channels", 3 } }; + inline for (want) |kv| { + const v = vc.object.get(kv[0]) orelse return error.VisionConfig; + if (v != .integer or v.integer != kv[1]) { + std.log.err("glm vision: vision_config.{s} is {f}; this tower reads {d}", .{ kv[0], std.json.fmt(v, .{}), kv[1] }); + return error.VisionConfig; + } + } + const act = vc.object.get("hidden_act") orelse return error.VisionConfig; + if (act != .string or !std.mem.eql(u8, act.string, "silu")) return error.VisionConfig; + var c: Config = .{ .image_token = 0, .image_start = 0, .image_end = 0, .mean = .{ 0.48145466, 0.4578275, 0.40821073 }, .std = .{ 0.26862954, 0.26130258, 0.27577711 }, .min_tokens = 16, .max_tokens = 8000 }; + if (vc.object.get("rms_norm_eps")) |v| c.eps = @floatCast(number(v) orelse return error.VisionConfig); + if (vc.object.get("swiglu_limit")) |v| c.limit = @floatCast(number(v) orelse return error.VisionConfig); + inline for (.{ .{ "image_token_id", "image_token" }, .{ "image_start_token_id", "image_start" }, .{ "image_end_token_id", "image_end" } }) |kv| { + const v = doc.object.get(kv[0]) orelse return error.VisionConfig; + if (v != .integer) return error.VisionConfig; + @field(c, kv[1]) = @intCast(v.integer); + } + if (doc.object.get("video_start_token_id")) |s| if (doc.object.get("video_token_id")) |t| if (doc.object.get("video_end_token_id")) |e| { + if (s == .integer and t == .integer and e == .integer) c.video = .{ @intCast(s.integer), @intCast(t.integer), @intCast(e.integer) }; + }; + const pj = processor_json orelse return c; + const pdoc = try std.json.parseFromSliceLeaky(std.json.Value, arena, pj, .{}); + const ip = pdoc.object.get("image_processor") orelse return c; + if (ip != .object) return error.VisionConfig; + inline for (.{ .{ "patch_size", patch }, .{ "temporal_patch_size", temporal }, .{ "merge_size", merge } }) |kv| { + if (ip.object.get(kv[0])) |v| if (v != .integer or v.integer != kv[1]) return error.VisionConfig; + } + if (ip.object.get("patch_expand_factor")) |v| if (v != .integer or v.integer != 1) return error.VisionConfig; + inline for (.{ "image_mean", "image_std" }, .{ &c.mean, &c.std }) |key, dst| if (ip.object.get(key)) |v| { + if (v != .array or v.array.items.len != 3) return error.VisionConfig; + for (v.array.items, 0..) |x, i| dst[i] = @floatCast(number(x) orelse return error.VisionConfig); + }; + if (ip.object.get("min_image_tokens")) |v| c.min_tokens = @intCast(v.integer); + if (ip.object.get("max_image_tokens")) |v| c.max_tokens = @intCast(v.integer); + return c; + } + + fn number(v: std.json.Value) ?f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => null, + }; + } +}; + +const Kernels = struct { + lib: mtl.Library, + gemm: [3][2][2]mtl.Pipeline, // [bf16 nt, f32 nt, f32 nn][M % 64 == 0][N % 128 == 0] + bias: mtl.Pipeline, + add: mtl.Pipeline, + rms: mtl.Pipeline, + layer_norm: mtl.Pipeline, + qkv: mtl.Pipeline, + softmax: mtl.Pipeline, + heads_: mtl.Pipeline, + swiglu: mtl.Pipeline, + gelu: mtl.Pipeline, + + fn load(gpa: std.mem.Allocator, device: mtl.Device) !Kernels { + const src = try frags.source(device, gpa, sources.glm_vision); + defer gpa.free(src); + const lib = try mtl.Library.fromSource(device, src, mtl.CompileOptions.mlx()); + var k: Kernels = undefined; + k.lib = lib; + const kinds = [_][]const u8{ "bf16_nt", "f32_nt", "f32_nn" }; + for (kinds, 0..) |kind, ki| for (0..2) |m| for (0..2) |n| { + var name: [64]u8 = undefined; + const full = try std.fmt.bufPrint(&name, "glmv_gemm_{s}_{d}{d}", .{ kind, m, n }); + k.gemm[ki][m][n] = try mtl.Pipeline.init(device, lib, full, false); + }; + inline for (.{ .{ "bias", "glmv_bias" }, .{ "add", "glmv_add" }, .{ "rms", "glmv_rms" }, .{ "layer_norm", "glmv_layer_norm" }, .{ "qkv", "glmv_qkv" }, .{ "softmax", "glmv_softmax" }, .{ "heads_", "glmv_heads" }, .{ "swiglu", "glmv_swiglu" }, .{ "gelu", "glmv_gelu" } }) |kv| { + @field(k, kv[0]) = try mtl.Pipeline.init(device, lib, kv[1], false); + } + return k; + } + + fn deinit(k: *Kernels) void { + for (&k.gemm) |*a| for (a) |*b| for (b) |p| p.deinit(); + inline for (.{ "bias", "add", "rms", "layer_norm", "qkv", "softmax", "heads_", "swiglu", "gelu" }) |f| @field(k, f).deinit(); + k.lib.deinit(); + } +}; + +const Block = struct { norm1: Ref, norm2: Ref, qkv_w: Ref, qkv_b: Ref, proj_w: Ref, proj_b: Ref, q_norm: Ref, k_norm: Ref, gate_w: Ref, gate_b: Ref, up_w: Ref, up_b: Ref, down_w: Ref, down_b: Ref }; + +const Weights = struct { + buf: mtl.Buffer, + patch_w: Ref, + patch_b: Ref, + blocks: [depth]Block, + post_norm: Ref, + down_w: Ref, // [4096, 2, 2, 1024]: the merge's kernel with its taps outermost, as four neighbouring patch rows read it + down_b: Ref, + m_proj: Ref, + m_norm_w: Ref, + m_norm_b: Ref, + m_gate: Ref, + m_up: Ref, + m_down: Ref, +}; + +/// The tower: kernels, weights (bf16, about 0.9 GiB) and scratch for `max_tokens` image tokens a request. +pub const Vision = struct { + gpa: std.mem.Allocator, + inject: ?struct { slot: u32, v: fwd.Inject } = null, // an image prompt's rows while its pass runs (slot_images.zig) + spans: std.ArrayList(fwd.Span) = .empty, + c: Config, + k: Kernels, + w: Weights, + scratch: mtl.Buffer, + max_patches: u32, + chunk: u32, // query rows an attention pass scores at once (bounds the fp32 score buffer) + // scratch views + patches: Ref, // [P, 1176] bf16 + cos: Ref, // [P, 64] f32 + sin: Ref, + x: Ref, // [P, 1024] bf16: the residual stream + xn: Ref, // [P, 1024]: a norm's output; the merger's first [T, 4096] + att: Ref, // [P, 1024]: heads side by side; the merger's second [T, 4096] + branch: Ref, // [P, 1024] + qkv: Ref, // [P, 3072] + q: Ref, // [16, P, 64] f32 + k_: Ref, + v: Ref, + s: Ref, // [16, chunk, P] f32 + o: Ref, // [16, chunk, 64] f32 + g: Ref, // [P, 4096] bf16 (the merger's [T, 10240] fits: T * 10240 < P * 4096) + u: Ref, + act: Ref, + bytes: usize, + + pub fn load(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8, c: Config, max_tokens: u32) !*Vision { + const v = try gpa.create(Vision); + errdefer gpa.destroy(v); + v.gpa = gpa; + v.inject = null; + v.spans = .empty; + v.c = c; + v.k = try Kernels.load(gpa, device); + errdefer v.k.deinit(); + v.w = try loadWeights(gpa, device, dir); + errdefer v.w.buf.deinit(); + const P: usize = @as(usize, max_tokens) * patches_per_token; + v.max_patches = @intCast(P); + // the scores take at most 256 MiB: 16 heads of `chunk` rows by P keys in fp32 + v.chunk = @intCast(std.math.clamp((256 << 20) / (heads * P * 4) / 64 * 64, 64, 4096)); + const parts = [_]usize{ P * patch_width * 2, P * 64 * 4, P * 64 * 4, P * hidden * 2, P * hidden * 2, P * hidden * 2, P * hidden * 2, P * 3 * hidden * 2, P * hidden * 4, P * hidden * 4, P * hidden * 4, @as(usize, heads) * v.chunk * P * 4, @as(usize, heads) * v.chunk * head_dim * 4, P * mlp_inter * 2, P * mlp_inter * 2, P * mlp_inter * 2 }; + var total: usize = 0; + for (parts) |n| total += std.mem.alignForward(usize, n, 256); + v.scratch = try device.buffer(total, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer v.scratch.deinit(); + var at: usize = 0; + const views = [_]*Ref{ &v.patches, &v.cos, &v.sin, &v.x, &v.xn, &v.att, &v.branch, &v.qkv, &v.q, &v.k_, &v.v, &v.s, &v.o, &v.g, &v.u, &v.act }; + for (views, parts) |view, n| { + view.* = .{ .buf = v.scratch, .off = at }; + at += std.mem.alignForward(usize, n, 256); + } + v.bytes = total + v.w.buf.length(); + return v; + } + + pub fn deinit(v: *Vision) void { + v.spans.deinit(v.gpa); + v.scratch.deinit(); + v.w.buf.deinit(); + v.k.deinit(); + v.gpa.destroy(v); + } + + /// Slot `i`'s image rows for its prompt pass, while its images' pass runs; null for any other slot. + pub fn injected(v: *const Vision, i: u32) ?*const fwd.Inject { + if (v.inject) |*inj| if (inj.slot == i) return &inj.v; + return null; + } + + /// Buffers to keep resident with the engine's. + pub fn buffers(v: *const Vision) [2]mtl.Buffer { + return .{ v.w.buf, v.scratch }; + } + + /// One image's patches ([gh * gw, 1176] fp32) into its rows `out`; `stop`: end after that block, residual in `x`. + pub fn encode(v: *Vision, enc: mtl.ComputeEncoder, pixels: []const f32, gh: u32, gw: u32, out: Ref, stop: ?u32) !void { + const P: u32 = gh * gw; + if (P == 0 or gh % merge != 0 or gw % merge != 0) return error.ImageGrid; + if (P > v.max_patches) return error.ImageTooLarge; + if (pixels.len != @as(usize, P) * patch_width) return error.ImageGrid; + const T: u32 = P / patches_per_token; + // host: the patches in bf16 (MLX's astype), and each patch's rotary angles + const dst: [*]u16 = @ptrCast(@alignCast(v.patches.addr())); + for (pixels, 0..) |p, i| dst[i] = bf16(p); + rotary(v, gh, gw); + + v.gemmLinear(enc, v.patches, v.w.patch_w, v.x, P, hidden, patch_width); + v.biasAdd(enc, v.x, v.w.patch_b, P, hidden); + if (stop) |s| if (s == 0) return; + for (v.w.blocks, 0..) |b, i| { + v.block(enc, b, P); + if (stop) |s| if (s == i + 1) return; + } + v.rms(enc, v.x, v.w.post_norm, v.xn, P, hidden); + // the merge: four neighbouring patches' rows are one row of xn read as [T, 4096] + const m0 = v.xn; + const m1 = v.att; + v.gemmLinear(enc, m0, v.w.down_w, m1, T, out_hidden, patches_per_token * hidden); + v.biasAdd(enc, m1, v.w.down_b, T, out_hidden); + v.gemmLinear(enc, m1, v.w.m_proj, m0, T, out_hidden, out_hidden); + v.layerNorm(enc, m0, v.w.m_norm_w, v.w.m_norm_b, m1, T, out_hidden); + enc.setPipeline(v.k.gelu); + enc.setBuffer(m1.buf, m1.off, 0); + enc.dispatchThreads(mtl.Size.of(@as(usize, T) * out_hidden, 1, 1), mtl.Size.of(256, 1, 1)); + v.gemmLinear(enc, m1, v.w.m_gate, v.g, T, merger_inter, out_hidden); + v.gemmLinear(enc, m1, v.w.m_up, v.u, T, merger_inter, out_hidden); + v.swiglu(enc, @as(usize, T) * merger_inter); + v.gemmLinear(enc, v.act, v.w.m_down, out, T, out_hidden, merger_inter); + } + + fn block(v: *Vision, enc: mtl.ComputeEncoder, b: Block, P: u32) void { + v.rms(enc, v.x, b.norm1, v.xn, P, hidden); + v.gemmLinear(enc, v.xn, b.qkv_w, v.qkv, P, 3 * hidden, hidden); + v.biasAdd(enc, v.qkv, b.qkv_b, P, 3 * hidden); + enc.setPipeline(v.k.qkv); + const bufs = [_]Ref{ v.qkv, b.q_norm, b.k_norm, v.cos, v.sin, v.q, v.k_, v.v }; + for (bufs, 0..) |r, i| enc.setBuffer(r.buf, r.off, i); + enc.setValue(extern struct { L: u32, eps: f32 }{ .L = P, .eps = v.c.eps }, 8); + enc.dispatchGroups(mtl.Size.of(@as(usize, P) * heads * 3, 1, 1), mtl.Size.of(32, 1, 1)); + var q0: u32 = 0; + while (q0 < P) : (q0 += v.chunk) { + const lc = @min(v.chunk, P - q0); + // scores of this chunk's queries against every key, each head a batch + v.gemm(enc, 1, v.q.at(@as(usize, q0) * head_dim * 4), v.k_, v.s, lc, P, head_dim, head_dim, head_dim, P, heads, @as(usize, P) * head_dim, @as(usize, P) * head_dim, @as(usize, lc) * P); + enc.setPipeline(v.k.softmax); + enc.setBuffer(v.s.buf, v.s.off, 0); + enc.setValue(P, 1); + enc.setValue(@as(f32, 0.125), 2); // head_dim ** -0.5 + enc.dispatchGroups(mtl.Size.of(@as(usize, heads) * lc, 1, 1), mtl.Size.of(256, 1, 1)); + v.gemm(enc, 2, v.s, v.v, v.o, lc, head_dim, P, P, head_dim, head_dim, heads, @as(usize, lc) * P, @as(usize, P) * head_dim, @as(usize, lc) * head_dim); + enc.setPipeline(v.k.heads_); + enc.setBuffer(v.o.buf, v.o.off, 0); + enc.setBuffer(v.att.buf, v.att.off, 1); + enc.setValue([2]u32{ lc, q0 }, 2); + enc.dispatchThreads(mtl.Size.of(hidden, lc, 1), mtl.Size.of(256, 1, 1)); + } + v.gemmLinear(enc, v.att, b.proj_w, v.branch, P, hidden, hidden); + v.biasAdd(enc, v.branch, b.proj_b, P, hidden); + v.add(enc, v.x, v.branch, @as(usize, P) * hidden); + v.rms(enc, v.x, b.norm2, v.xn, P, hidden); + v.gemmLinear(enc, v.xn, b.gate_w, v.g, P, mlp_inter, hidden); + v.biasAdd(enc, v.g, b.gate_b, P, mlp_inter); + v.gemmLinear(enc, v.xn, b.up_w, v.u, P, mlp_inter, hidden); + v.biasAdd(enc, v.u, b.up_b, P, mlp_inter); + v.swiglu(enc, @as(usize, P) * mlp_inter); + v.gemmLinear(enc, v.act, b.down_w, v.branch, P, hidden, mlp_inter); + v.biasAdd(enc, v.branch, b.down_b, P, hidden); + v.add(enc, v.x, v.branch, @as(usize, P) * hidden); + } + + /// y [M, N] = x [M, K] w^T in bf16 (w as stored, [N, K]). + fn gemmLinear(v: *Vision, enc: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, M: u32, N: u32, K: u32) void { + v.gemm(enc, 0, x, w, y, M, N, K, K, K, N, 1, 0, 0, 0); + } + + /// tfp::gemm_nax: d [batch, M, N] = a x b (kind 0: bf16 a b^T, 1: fp32 a b^T, 2: fp32 a b), batches `s*` apart. + fn gemm(v: *Vision, enc: mtl.ComputeEncoder, kind: usize, a: Ref, b: Ref, d: Ref, M: u32, N: u32, K: u32, lda: u32, ldb: u32, ldd: u32, batch: u32, sa: usize, sb: usize, sd: usize) void { + const tn = (N + 127) / 128; + const tm = (M + 63) / 64; + enc.setPipeline(v.k.gemm[kind][@intFromBool(M % 64 == 0)][@intFromBool(N % 128 == 0)]); + enc.setBuffer(a.buf, a.off, 0); + enc.setBuffer(b.buf, b.off, 1); + var p: [16]i32 = @splat(0); + const vals = [_]usize{ M, N, K, lda, ldb, ldd, tn, tm, 2, K / 256, sa, sb, sd, 0, 0 }; + for (vals, 0..) |x, i| p[i] = @intCast(x); + enc.setBytes(std.mem.asBytes(&p), 2); + enc.setBuffer(d.buf, d.off, 3); + enc.dispatchThreads(mtl.Size.of(@as(usize, tn << 2) * 32, (tm + 3) / 4 * 4, @as(usize, batch) * 2), mtl.Size.of(32, 4, 2)); + } + + fn biasAdd(v: *Vision, enc: mtl.ComputeEncoder, y: Ref, b: Ref, rows: u32, width: u32) void { + enc.setPipeline(v.k.bias); + enc.setBuffer(y.buf, y.off, 0); + enc.setBuffer(b.buf, b.off, 1); + enc.setValue(width, 2); + enc.dispatchThreads(mtl.Size.of(width, rows, 1), mtl.Size.of(256, 1, 1)); + } + + fn add(v: *Vision, enc: mtl.ComputeEncoder, x: Ref, y: Ref, n: usize) void { + enc.setPipeline(v.k.add); + enc.setBuffer(x.buf, x.off, 0); + enc.setBuffer(y.buf, y.off, 1); + enc.dispatchThreads(mtl.Size.of(n, 1, 1), mtl.Size.of(256, 1, 1)); + } + + fn rms(v: *Vision, enc: mtl.ComputeEncoder, x: Ref, w: Ref, y: Ref, rows: u32, width: u32) void { + enc.setPipeline(v.k.rms); + enc.setBuffer(x.buf, x.off, 0); + enc.setBuffer(w.buf, w.off, 1); + enc.setBuffer(y.buf, y.off, 2); + enc.setValue(extern struct { rows: u32, width: u32, eps: f32 }{ .rows = rows, .width = width, .eps = v.c.eps }, 3); + enc.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + fn layerNorm(v: *Vision, enc: mtl.ComputeEncoder, x: Ref, w: Ref, b: Ref, y: Ref, rows: u32, width: u32) void { + enc.setPipeline(v.k.layer_norm); + const bufs = [_]Ref{ x, w, b, y }; + for (bufs, 0..) |r, i| enc.setBuffer(r.buf, r.off, i); + enc.setValue(extern struct { rows: u32, width: u32, eps: f32 }{ .rows = rows, .width = width, .eps = 1e-5 }, 4); + enc.dispatchGroups(mtl.Size.of(rows, 1, 1), mtl.Size.of(256, 1, 1)); + } + + fn swiglu(v: *Vision, enc: mtl.ComputeEncoder, n: usize) void { + enc.setPipeline(v.k.swiglu); + enc.setBuffer(v.g.buf, v.g.off, 0); + enc.setBuffer(v.u.buf, v.u.off, 1); + enc.setBuffer(v.act.buf, v.act.off, 2); + enc.setValue(v.c.limit, 3); + enc.dispatchThreads(mtl.Size.of(n, 1, 1), mtl.Size.of(256, 1, 1)); + } +}; + +/// fp32 to bf16, round to nearest even (MLX's astype). +pub fn bf16(x: f32) u16 { + const b: u32 = @bitCast(x); + if (b & 0x7fffffff > 0x7f800000) return @truncate((b >> 16) | 0x40); // NaN stays NaN + return @truncate((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// Each patch's 2-D rotary cos and sin: its row and column in the patch grid times the 16 frequencies 10000^(-k/16). +fn rotary(v: *Vision, gh: u32, gw: u32) void { + var inv: [16]f32 = undefined; + for (&inv, 0..) |*f, k| f.* = 1.0 / @as(f32, @floatCast(std.math.pow(f64, 10000.0, @as(f64, @floatFromInt(2 * k)) / 32.0))); + const cs: [*]f32 = @ptrCast(@alignCast(v.cos.addr())); + const sn: [*]f32 = @ptrCast(@alignCast(v.sin.addr())); + const wx_n = gw / merge; + var i: usize = 0; + var wy: u32 = 0; + while (wy < gh / merge) : (wy += 1) { + var wx: u32 = 0; + while (wx < wx_n) : (wx += 1) { + for (0..merge) |my| for (0..merge) |mx| { + const h: f32 = @floatFromInt(wy * merge + my); + const w: f32 = @floatFromInt(wx * merge + mx); + for (0..32) |j| { + const angle: f32 = if (j < 16) h * inv[j] else w * inv[j - 16]; + const c: f32 = @floatCast(@cos(@as(f64, angle))); + const s: f32 = @floatCast(@sin(@as(f64, angle))); + cs[i * 64 + j] = c; + cs[i * 64 + j + 32] = c; + sn[i * 64 + j] = s; + sn[i * 64 + j + 32] = s; + } + i += 1; + }; + } + } +} + +/// model.visual.* from the checkpoint's shards into one buffer, the merge kernel's taps moved outermost. +fn loadWeights(gpa: std.mem.Allocator, device: mtl.Device, dir: []const u8) !Weights { + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const arena = arena_state.allocator(); + const Src = struct { path: [:0]const u8, off: u64, len: u64, dtype: st.DType, shape: []const usize }; + var names: std.StringHashMapUnmanaged(Src) = .empty; + { + const path = try std.fmt.allocPrintSentinel(arena, "{s}/model.safetensors.index.json", .{dir}, 0); + const f = try mtl.MappedFile.open(path); + defer f.deinit(); + const doc = try std.json.parseFromSliceLeaky(std.json.Value, arena, f.bytes[0..f.size], .{ .allocate = .alloc_always }); + var files: std.StringArrayHashMapUnmanaged(void) = .empty; + var it = doc.object.get("weight_map").?.object.iterator(); + while (it.next()) |kv| if (std.mem.startsWith(u8, kv.key_ptr.*, "model.visual.")) try files.put(arena, kv.value_ptr.string, {}); + if (files.count() == 0) return error.NoVisionTower; + for (files.keys()) |name| { + const full = try std.fmt.allocPrintSentinel(arena, "{s}/{s}", .{ dir, name }, 0); + const fd = std.c.open(full, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + var head: [8]u8 = undefined; + try readAll(fd, &head, 0); + const hlen = std.mem.readInt(u64, &head, .little); + const header = try arena.alloc(u8, hlen); + try readAll(fd, header, 8); + const end: u64 = @intCast(std.c.lseek(fd, 0, std.c.SEEK.END)); + // read here, not by core/safetensors.zig: the patch kernel is 5-D, past its max_rank + const doc2 = try std.json.parseFromSliceLeaky(std.json.Value, arena, header, .{}); + var eit = doc2.object.iterator(); + while (eit.next()) |kv| { + if (!std.mem.startsWith(u8, kv.key_ptr.*, "model.visual.")) continue; + const o = kv.value_ptr.object; + const dtype = st.DType.parse(o.get("dtype").?.string) orelse return error.UnsupportedDType; + const dims = o.get("shape").?.array.items; + const shape = try arena.alloc(usize, dims.len); + var n: u64 = dtype.size(); + for (dims, shape) |d, *s| { + s.* = @intCast(d.integer); + n *= s.*; + } + const offs = o.get("data_offsets").?.array.items; + const b: u64 = @intCast(offs[0].integer); + const e: u64 = @intCast(offs[1].integer); + if (e < b or e - b != n or 8 + hlen + e > end) return error.BadSafetensors; + try names.put(arena, kv.key_ptr.*, .{ .path = full, .off = 8 + hlen + b, .len = e - b, .dtype = dtype, .shape = shape }); + } + } + } + // every tensor's place in one buffer, 256-aligned + const Want = struct { name: []const u8, shape: []const usize, ref: *Ref }; + var w: Weights = undefined; + var wants: std.ArrayList(Want) = .empty; + try wants.append(arena, .{ .name = "patch_embed.proj.weight", .shape = &.{ hidden, 3, temporal, patch, patch }, .ref = &w.patch_w }); + try wants.append(arena, .{ .name = "patch_embed.proj.bias", .shape = &.{hidden}, .ref = &w.patch_b }); + for (&w.blocks, 0..) |*b, i| { + const fields = .{ .{ "norm1.weight", &b.norm1, &[_]usize{hidden} }, .{ "norm2.weight", &b.norm2, &[_]usize{hidden} }, .{ "attn.qkv.weight", &b.qkv_w, &[_]usize{ 3 * hidden, hidden } }, .{ "attn.qkv.bias", &b.qkv_b, &[_]usize{3 * hidden} }, .{ "attn.proj.weight", &b.proj_w, &[_]usize{ hidden, hidden } }, .{ "attn.proj.bias", &b.proj_b, &[_]usize{hidden} }, .{ "attn.q_norm.weight", &b.q_norm, &[_]usize{head_dim} }, .{ "attn.k_norm.weight", &b.k_norm, &[_]usize{head_dim} }, .{ "mlp.gate_proj.weight", &b.gate_w, &[_]usize{ mlp_inter, hidden } }, .{ "mlp.gate_proj.bias", &b.gate_b, &[_]usize{mlp_inter} }, .{ "mlp.up_proj.weight", &b.up_w, &[_]usize{ mlp_inter, hidden } }, .{ "mlp.up_proj.bias", &b.up_b, &[_]usize{mlp_inter} }, .{ "mlp.down_proj.weight", &b.down_w, &[_]usize{ hidden, mlp_inter } }, .{ "mlp.down_proj.bias", &b.down_b, &[_]usize{hidden} } }; + inline for (fields) |f| try wants.append(arena, .{ .name = try std.fmt.allocPrint(arena, "blocks.{d}.{s}", .{ i, f[0] }), .shape = f[2], .ref = f[1] }); + } + try wants.append(arena, .{ .name = "post_layernorm.weight", .shape = &.{hidden}, .ref = &w.post_norm }); + try wants.append(arena, .{ .name = "downsample.weight", .shape = &.{ out_hidden, hidden, merge, merge }, .ref = &w.down_w }); + try wants.append(arena, .{ .name = "downsample.bias", .shape = &.{out_hidden}, .ref = &w.down_b }); + try wants.append(arena, .{ .name = "merger.proj.weight", .shape = &.{ out_hidden, out_hidden }, .ref = &w.m_proj }); + try wants.append(arena, .{ .name = "merger.post_projection_norm.weight", .shape = &.{out_hidden}, .ref = &w.m_norm_w }); + try wants.append(arena, .{ .name = "merger.post_projection_norm.bias", .shape = &.{out_hidden}, .ref = &w.m_norm_b }); + try wants.append(arena, .{ .name = "merger.gate_proj.weight", .shape = &.{ merger_inter, out_hidden }, .ref = &w.m_gate }); + try wants.append(arena, .{ .name = "merger.up_proj.weight", .shape = &.{ merger_inter, out_hidden }, .ref = &w.m_up }); + try wants.append(arena, .{ .name = "merger.down_proj.weight", .shape = &.{ out_hidden, merger_inter }, .ref = &w.m_down }); + const srcs = try arena.alloc(Src, wants.items.len); + var total: usize = 0; + for (wants.items, srcs) |want, *s| { + const full = try std.fmt.allocPrint(arena, "model.visual.{s}", .{want.name}); + s.* = names.get(full) orelse { + std.log.err("glm vision: the checkpoint has no tensor {s}", .{full}); + return error.MissingTensor; + }; + if (s.dtype != .bf16 or !std.mem.eql(usize, s.shape, want.shape)) { + std.log.err("glm vision: {s} is {t} {any}; the tower reads bf16 {any}", .{ full, s.dtype, s.shape, want.shape }); + return error.UnexpectedTensor; + } + total = std.mem.alignForward(usize, total, 256) + s.len; + } + w.buf = try device.buffer(total, mtl.ResourceOptions.shared | mtl.ResourceOptions.untracked); + errdefer w.buf.deinit(); + var at: usize = 0; + for (wants.items, srcs) |want, s| { + at = std.mem.alignForward(usize, at, 256); + want.ref.* = .{ .buf = w.buf, .off = at }; + const fd = std.c.open(s.path, .{ .ACCMODE = .RDONLY }); + if (fd < 0) return error.OpenFailed; + defer _ = std.c.close(fd); + _ = std.c.fcntl(fd, 48, @as(c_int, 1)); // F_NOCACHE, as the language model's shards + const dst = w.buf.contents()[at..][0..s.len]; + try readAll(fd, dst, s.off); + at += s.len; + } + // downsample [O, C, 2, 2] -> [O, 2, 2, C]: a merge window's four patch rows, tap by tap + const dw: [*]u16 = @ptrCast(@alignCast(w.down_w.addr())); + const tmp = try gpa.alloc(u16, hidden * merge * merge); + defer gpa.free(tmp); + for (0..out_hidden) |o| { + const row = dw[o * hidden * 4 ..][0 .. hidden * 4]; + @memcpy(tmp, row); + for (0..hidden) |ch| for (0..4) |tap| { + row[tap * hidden + ch] = tmp[ch * 4 + tap]; + }; + } + return w; +} + +/// `e` reads images: the tower loaded beside its weights (about 0.9 GiB with scratch) and the rows images fill. +pub fn attach(e: *Engine, gpa: std.mem.Allocator, dir: []const u8, config: []const u8, lim: Limits) !void { + if (e.ep != null) return error.VisionOnPair; + const pj: ?mtl.MappedFile = blk: { + const pp = try std.fmt.allocPrintSentinel(gpa, "{s}/processor_config.json", .{dir}, 0); + defer gpa.free(pp); + break :blk mtl.MappedFile.open(pp) catch null; + }; + defer if (pj) |m| m.deinit(); + var cfg_arena = std.heap.ArenaAllocator.init(gpa); + defer cfg_arena.deinit(); + const vc = (try Config.parse(cfg_arena.allocator(), config, if (pj) |m| m.bytes[0..m.size] else null)) orelse return error.NoVisionTower; + e.vision = try Vision.load(gpa, e.device, dir, vc, lim.image_tokens); + e.vision_rows = try e.arena.buffer(@as(usize, lim.image_tokens) * e.c.hidden * 2); + e.vision_limits = lim; + std.log.info("glm: image tower loaded ({d:.2} GiB with scratch for {d} image tokens a prompt)", .{ @as(f64, @floatFromInt(e.vision.?.bytes)) / (1 << 30), lim.image_tokens }); +} + +test "bf16 rounds to nearest even" { + try std.testing.expectEqual(@as(u16, 0x3f80), bf16(1.0)); + try std.testing.expectEqual(@as(u16, 0x3f80), bf16(@bitCast(@as(u32, 0x3f808000)))); // a tie to even (down) + try std.testing.expectEqual(@as(u16, 0x3f82), bf16(@bitCast(@as(u32, 0x3f818000)))); // a tie to even (up) + try std.testing.expectEqual(@as(u16, 0x3f81), bf16(@bitCast(@as(u32, 0x3f808001)))); +} diff --git a/zig/src/native/glm_host.zig b/zig/src/native/glm_host.zig index dab379dbf3..4c68e3fe59 100644 --- a/zig/src/native/glm_host.zig +++ b/zig/src/native/glm_host.zig @@ -16,6 +16,9 @@ pub const Host = struct { eng: *ge.Engine, slots: glm.slots.Slots, back: glm.backend.Backend, + markers: [3]u32 = undefined, // the image wrapper's tokens (--vision): begin, image, end + video_markers: [3]u32 = undefined, // the video wrapper's: begin, video, end + video_open: api.VisionOpen = .{}, // the video limits the server was started with cfg: lanes.Config, wall: lanes.backend.WallClock, core: lanes.Engine, @@ -79,8 +82,10 @@ fn learnedStates(gpa: Allocator, eng: *const ge.Engine, sl: *glm.slots.Slots, ro } /// A GLM-5.3-Flash checkpoint served: `window` tokens of cache a stream, warmed first (the pair: `speed_up`). -pub fn open(gpa: Allocator, io: std.Io, dir: []const u8, window: u32, speed_up: ?[]const u8, streams: u32, fixed: bool, cache_gib: ?f64, learn: ?[]const u8, learn_cap: u64, learn_floor: u64) !*Host { - const eng = try ge.Engine.loadWith(gpa, dir, window + 64, speed_up, learn != null); +pub fn open(gpa: Allocator, io: std.Io, dir: []const u8, window: u32, speed_up: ?[]const u8, streams: u32, fixed: bool, cache_gib: ?f64, learn: ?[]const u8, learn_cap: u64, learn_floor: u64, vision: ?api.VisionOpen) !*Host { + // a prompt's tower rows and images: the images' budget, and every video's frame groups beside it + const limits: ?glm.vision.Limits = if (vision) |v| .{ .image_tokens = v.image_tokens + v.video_tokens, .max_images = v.max_images + v.max_videos * ((v.max_frames + 1) / 2) } else null; + const eng = try ge.Engine.loadWith(gpa, dir, window + 64, speed_up, learn != null, limits); errdefer eng.deinit(); var toks: [96]u32 = undefined; for (&toks, 0..) |*t, i| t.* = @intCast(1000 + i); @@ -104,7 +109,18 @@ pub fn open(gpa: Allocator, io: std.Io, dir: []const u8, window: u32, speed_up: h.core = lanes.Engine.init(gpa, &h.cfg, h.back.backend(), h.wall.clock()); errdefer h.core.deinit(); const serves = @min(window, glm.pool.capacity(&h.slots) -| (glm.state.max_rows + 1)); // prompt + reply a stream's caches hold beside its last window (a pool keeps whole blocks) - h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "glm-zig", .lanes = n, .context_window = serves, .prefill_step = eng.chunk_rows, .greedy_only = true, .admit = if (h.slots.pool != null) .{ .ctx = &h.back, .fits = glm.backend.Backend.fits } else null }); + var offer: ?api.Vision = null; + if (eng.vision) |v| { + const o = vision.?; + h.markers = .{ v.c.image_start, v.c.image_token, v.c.image_end }; + h.video_open = o; + offer = .{ .ctx = h, .prepare = prepareImage, .image_token = v.c.image_token, .markers = &h.markers, .image_tokens = o.image_tokens, .max_images = o.max_images }; + if (v.c.video) |vm| { + h.video_markers = vm; + offer.?.video = .{ .prepare = prepareVideo, .video_token = vm[1], .group_open = v.c.image_start, .group_close = v.c.image_end, .markers = &h.video_markers, .video_tokens = o.video_tokens, .max_videos = o.max_videos, .max_frames = o.max_frames }; + } + } + h.host = api.LaneHost.init(gpa, io, &h.core, .{ .name = "glm-zig", .lanes = n, .context_window = serves, .prefill_step = eng.chunk_rows, .greedy_only = true, .admit = if (h.slots.pool != null) .{ .ctx = &h.back, .fits = glm.backend.Backend.fits } else null, .vision = offer }); h.cache = null; const budget = cacheBudget(eng, cache_gib); if (!eng.followsPeer() and budget > 0) { @@ -141,6 +157,37 @@ fn growSnap(ctx: *anyopaque, saved: *anyopaque, extra: u64) bool { return if (h.cache) |*store| store.grow(saved, extra) else false; } +/// Info.vision's prepare: an image's bytes to the tower's patches (families/glm/image.zig), on the HTTP thread. +fn prepareImage(ctx: *anyopaque, a: Allocator, bytes: []const u8, max_tokens: u32) anyerror!api.PreparedImage { + const h: *Host = @ptrCast(@alignCast(ctx)); + const p = try glm.image.prepare(a, bytes, h.eng.vision.?.c, max_tokens); + return .{ .pixels = p.pixels, .gh = p.gh, .gw = p.gw, .tokens = p.tokens, .hash = p.hash }; +} + +/// Info.vision's video offer: a decoded video's frames to the tower's frame groups (glm/video.zig), on the HTTP thread. +fn prepareVideo(ctx: *anyopaque, a: Allocator, src: api.VideoSource, max_tokens: u32, max_frames: u32) anyerror!api.PreparedVideo { + const h: *Host = @ptrCast(@alignCast(ctx)); + const Adapter = struct { + src: api.VideoSource, + fn decode(c: *anyopaque, indices: []const u32, sink: *glm.video.Sink) anyerror!void { + const self: *@This() = @ptrCast(@alignCast(c)); + const into = struct { + fn take(s: *anyopaque, k: usize, rgb: []const u8, w: u32, hh: u32, stride: usize) anyerror!void { + const g: *glm.video.Sink = @ptrCast(@alignCast(s)); + try g.take(k, rgb, w, hh, stride); + } + }.take; + try self.src.decode(self.src.ctx, indices, .{ .ctx = sink, .take = into }); + } + }; + var ad: Adapter = .{ .src = src }; + const s: glm.video.Source = .{ .frames = src.frames, .rate_num = src.rate_num, .rate_den = src.rate_den, .width = src.width, .height = src.height, .ctx = &ad, .decode = Adapter.decode }; + const p = try glm.video.prepare(a, s, h.eng.vision.?.c, max_tokens, max_frames); + const groups = try a.alloc(api.PreparedImage, p.groups.len); + for (groups, p.groups) |*g, q| g.* = .{ .pixels = q.pixels, .gh = q.gh, .gw = q.gw, .tokens = q.tokens, .hash = q.hash }; + return .{ .groups = groups, .times = p.times }; +} + fn follow(h: *Host) void { glm.mirror.follow(h.eng, &h.slots) catch |err| std.log.err("speed-up mode: following rank 0 ended: {s}", .{@errorName(err)}); } diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 39fe531408..98c7bfcd54 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -113,6 +113,10 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c if (dense) return openQwen27(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "qwen4_exp")) return openFlashNext(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "glm5_next")) return openGlm(a, gpa, io, o, problem); + if (o.vision != null) { + problem.* = try std.fmt.allocPrint(a, "--vision: the native engine reads images for GLM-5.3-Flash only, not {s} checkpoints", .{o.model_type}); + return null; + } if (std.mem.eql(u8, o.model_type, "glm_moe_dsa")) return openGlmDsa(a, gpa, io, o, problem); if (std.mem.eql(u8, o.model_type, "qwen3_5")) return @import("qwen35.zig").open(a, gpa, io, o, problem); if (!std.mem.eql(u8, o.model_type, "nemotron_h")) { @@ -353,7 +357,11 @@ fn openGlm(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]co } const pool = mtl.objc.Pool.push(); defer pool.pop(); - const h = glm.open(gpa, io, o.dir, @intCast(window), o.speed_up, o.lanes, o.lanes_fixed, o.prompt_cache_gib, o.learn, @intFromFloat(o.learn_gib * (1 << 30)), @intFromFloat(o.learn_min_free_gib * (1 << 30))) catch |e| { + if (o.vision != null and o.speed_up != null) { + problem.* = try std.fmt.allocPrint(a, "--vision is served on one Mac; the --speed-up pair reads text only", .{}); + return null; + } + const h = glm.open(gpa, io, o.dir, @intCast(window), o.speed_up, o.lanes, o.lanes_fixed, o.prompt_cache_gib, o.learn, @intFromFloat(o.learn_gib * (1 << 30)), @intFromFloat(o.learn_min_free_gib * (1 << 30)), o.vision) catch |e| { if (e == error.BadLoadLimit) problem.* = "GLM_LOAD_LIMIT_GB is not a load limit: give a number of GB above 0, or unset it for 70% of RAM" else problem.* = if (linkProblem(e)) |link| try std.fmt.allocPrint(a, "the native GLM-5.3-Flash engine cannot load {s}: {s}", .{ o.dir, link }) else try std.fmt.allocPrint(a, "the native GLM-5.3-Flash engine cannot load {s} ({s})", .{ o.dir, @errorName(e) }); return null; }; diff --git a/zig/src/server/anthropic.zig b/zig/src/server/anthropic.zig index bf465e7ac2..6401ebcca3 100644 --- a/zig/src/server/anthropic.zig +++ b/zig/src/server/anthropic.zig @@ -51,7 +51,7 @@ fn countTokens(srv: *Server, cx: *Cx, chat: Value) errors.Refused!usize { const thinking = try fields.thinkingFields(cx, chat, srv.effort_levels); const on = thinking.enable orelse srv.config.enable_thinking; const tools = try tool_specs.active(cx, chat.get("tools"), chat.get("tool_choice")); - const msgs = try messages_mod.normalize(cx, chat.get("messages"), "system", srv.needs_user_after_tool); + const msgs = try messages_mod.normalizeWith(cx, chat.get("messages"), "system", srv.needs_user_after_tool, .of(srv.info.vision)); const ids_ = prompt.renderIds(srv, cx, msgs, tools, on, srv.effortFor(thinking.effort), true) catch |e| switch (e) { error.Refused => return cx.other(cx.message), else => |x| return x, diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 7c702ef432..d3427ff4c0 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -141,6 +141,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa .decode_spans = try prompt_mod.replySpans(srv, cx, rendered.ids), // from the tokens alone, for every request: cached states are keyed by tokens .tools_json = if (input.tools.len > 0) try json.stringify(a, .{ .array = @constCast(input.tools) }, .{ .ascii = false }) else "", .decide = if (input.prompt == null and !background) prompt_mod.decideOf(srv, rendered.ids) else null, + .images = rendered.images, + .cache_key = rendered.cache_key, }; try srv.checkFeatures(cx, f, input.tools.len > 0, thinking, rendered.ids, input.tools, &request); if (thinking) { @@ -230,35 +232,6 @@ pub fn release(srv: *Server, preparing: bool) void { if (preparing) _ = srv.preparing.fetchSub(1, .acq_rel); } -/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. -const Shown = struct { - text: std.ArrayList(u8) = .empty, - chars: usize = 0, - at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds - extends: bool = false, - - /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. - fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { - const sent = s.text.items; - s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); - if (!s.extends) return reply_text.afterChars(now, s.chars); - if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again - return now[sent.len..]; - } - - fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { - if (s.extends) { - try s.text.appendSlice(a, delta); - s.chars += reply_text.charCount(delta); - } else { - s.text = .empty; - try s.text.appendSlice(a, now); - s.chars = reply_text.charCount(now); - } - s.at = now.ptr; - } -}; - /// The token loop and the text it streams. const Generation = struct { srv: *Server, @@ -276,8 +249,8 @@ const Generation = struct { collected: std.ArrayList(u32) = .empty, visible: reply_text.Incremental = .{}, hidden: std.ArrayList(u8) = .empty, // reused for the answer without its call blocks - streamed: Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices - streamed_reasoning: Shown = .{}, + streamed: reply_text.Shown = .{}, // what content streamed, kept apart from the decode buffer that grows under slices + streamed_reasoning: reply_text.Shown = .{}, streaming_done: bool = false, first_ns: ?i96 = null, last_ns: ?i96 = null, diff --git a/zig/src/server/cli.zig b/zig/src/server/cli.zig index 76b205eeaa..21d4b0b212 100644 --- a/zig/src/server/cli.zig +++ b/zig/src/server/cli.zig @@ -49,11 +49,12 @@ pub const flags = [_]Flag{ .{ .name = "--api-key-file", .native = true }, .{ .name = "--metrics-open", .kind = .store_true, .native = true }, .{ .name = "--dashboard", .kind = .store_true, .native = true }, - .{ .name = "--vision", .kind = .store_true }, - .{ .name = "--vision-urls", .kind = .store_true }, + .{ .name = "--vision", .kind = .store_true, .native = true }, + .{ .name = "--vision-urls", .kind = .store_true, .native = true }, .{ .name = "--vision-offload", .kind = .store_true }, - .{ .name = "--vision-max-images" }, - .{ .name = "--vision-image-tokens" }, + .{ .name = "--vision-max-images", .native = true }, + .{ .name = "--vision-image-tokens", .native = true }, + .{ .name = "--vision-video-tokens", .native = true }, .{ .name = "--context", .native = true }, .{ .name = "--speed-up", .native = !cuda_build }, .{ .name = "--max-tokens", .native = true }, @@ -170,6 +171,11 @@ pub const Args = struct { master_port: u16 = 29551, policy: []const u8 = "", // every --policy joined by commas cortheon: bool = false, // Cortheon's program, research before building, and cortheon_search/cortheon_fetch run by the server + vision: bool = false, // image input (GLM-5.3-Flash): data-URL image_url parts in user messages + vision_urls: bool = false, // with --vision, public HTTPS image URLs too (media_fetch.zig) + vision_max_images: u32 = 4, + vision_image_tokens: u32 = 4096, + vision_video_tokens: u32 = 16384, // the visual tokens a request's videos share (the processor's max_image_tokens) }; /// A usage error's message (argparse's ``error:`` line); the caller exits 2. @@ -221,6 +227,7 @@ pub fn parse(a: Allocator, argv: []const []const u8, u: *Usage) error{ Usage, Ou } out.model = model orelse return fail(u, a, "the following arguments are required: model", .{}); try ranks(u, a, out); + if (out.vision_urls and !out.vision) return fail(u, a, "--vision-urls needs --vision", .{}); out.alias = alias.items; out.api_key = keys.items; return out; @@ -275,7 +282,19 @@ fn apply(a: Allocator, out: *Args, name: []const u8, value: ?[]const u8, u: *Usa const n = try int(u, a, name, v); if (n < 0) return fail(u, a, "argument --compact-keep: expected a token count from 0: '{s}'", .{v}); out.compact_keep = @intCast(n); - } else if (is(name, "--compact-memory")) out.compact_memory = v else if (is(name, "--slide")) out.slide = true else if (is(name, "--slide-graph")) out.slide_graph = v else if (is(name, "--parallel")) out.parallel = v else if (is(name, "--backend")) out.backend = v else if (is(name, "--cortheon")) out.cortheon = true; + } else if (is(name, "--compact-memory")) out.compact_memory = v else if (is(name, "--slide")) out.slide = true else if (is(name, "--slide-graph")) out.slide_graph = v else if (is(name, "--parallel")) out.parallel = v else if (is(name, "--backend")) out.backend = v else if (is(name, "--cortheon")) out.cortheon = true else if (is(name, "--vision")) out.vision = true else if (is(name, "--vision-urls")) out.vision_urls = true else if (is(name, "--vision-max-images")) { + const n = try int(u, a, name, v); + if (n < 1 or n > 64) return fail(u, a, "argument --vision-max-images: expected 1 to 64 images: '{s}'", .{v}); + out.vision_max_images = @intCast(n); + } else if (is(name, "--vision-video-tokens")) { + const n = try int(u, a, name, v); + if (n < 64 or n > 262144) return fail(u, a, "argument --vision-video-tokens: expected 64 to 262144 tokens: '{s}'", .{v}); + out.vision_video_tokens = @intCast(n); + } else if (is(name, "--vision-image-tokens")) { + const n = try int(u, a, name, v); + if (n < 16 or n > 16384) return fail(u, a, "argument --vision-image-tokens: expected 16 to 16384 tokens: '{s}'", .{v}); + out.vision_image_tokens = @intCast(n); + } } /// The CUDA build's --device and --segments; false for any other flag. @@ -367,6 +386,7 @@ pub fn request(a: Allocator, dir: []const u8, model_type: []const u8, args: Args .master = args.master, .master_port = args.master_port, .policy = args.policy, + .vision = if (args.vision) .{ .image_tokens = args.vision_image_tokens, .max_images = args.vision_max_images, .video_tokens = args.vision_video_tokens } else null, }; } @@ -451,6 +471,10 @@ test "parse and capabilities share the table" { try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--compact-at", "0" }, &u)); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--compact-at", "2" }, &u)); try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--compact-keep", "-1" }, &u)); + const urls = try parse(a, &.{ "m", "--vision", "--vision-urls" }, &u); + try std.testing.expect(urls.vision and urls.vision_urls); + try std.testing.expectError(error.Usage, parse(a, &.{ "m", "--vision-urls" }, &u)); + try std.testing.expectEqualStrings("--vision-urls needs --vision", u.message); } test "--device and --segments: CUDA builds serve them, values checked; --speed-up is the Mac build's" { diff --git a/zig/src/server/hf_text.zig b/zig/src/server/hf_text.zig index 65c3a474d5..27ee33fde5 100644 --- a/zig/src/server/hf_text.zig +++ b/zig/src/server/hf_text.zig @@ -211,6 +211,7 @@ pub const HfText = struct { .template_source = sourceFn, .special = specialFn, .token_bytes = tokenBytesFn, + .encode_plain = encodePlainFn, } }; } @@ -231,6 +232,10 @@ pub const HfText = struct { return std.mem.concat(a, u32, &.{ t.prefix, ids, t.suffix }); } + fn encodePlainFn(ctx: *anyopaque, a: Allocator, s: []const u8) model_text.Error![]u32 { + return self(ctx).tok.encodePlain(a, s) catch |e| if (e == error.OutOfMemory) error.OutOfMemory else error.Template; + } + fn decodeFn(ctx: *anyopaque, a: Allocator, ids: []const u32) Allocator.Error![]u8 { return self(ctx).tok.decode(a, ids, false) catch |e| if (e == error.OutOfMemory) error.OutOfMemory else a.dupe(u8, "\u{fffd}"); } diff --git a/zig/src/server/images.zig b/zig/src/server/images.zig new file mode 100644 index 0000000000..b2e69843aa --- /dev/null +++ b/zig/src/server/images.zig @@ -0,0 +1,517 @@ +//! Image input (--vision): image_url parts to the engine's tower input, the prompt's placeholders to image tokens. +const std = @import("std"); +const api = @import("engine_api"); +const json = @import("json.zig"); +const errors = @import("errors.zig"); +const messages_mod = @import("messages.zig"); +const model_text = @import("model_text.zig"); +const media_fetch = @import("media_fetch.zig"); +const videos = @import("videos.zig"); +const Value = json.Value; +const Cx = errors.Cx; + +/// images.py's limits: an image's and a request's encoded bytes, and the visual tokens of a low-detail image. +pub const max_image_bytes = 10 << 20; +pub const max_request_bytes = 20 << 20; +const low_detail_tokens = 256; + +pub const Expanded = struct { + ids: []const u32, + history_len: usize, + images: []const api.Image, + cache_key: []const u32, +}; + +const Part = struct { video: bool = false, url: []const u8, low: bool = false }; + +/// The image and video parts of normalized `messages`, in the order the template meets them. +fn parts(cx: *Cx, messages: Value) errors.Refused![]const Part { + var out: std.ArrayList(Part) = .empty; + if (messages != .array) return out.items; + for (messages.array) |m| { + const content = m.get("content") orelse continue; + if (content != .array) continue; + for (content.array) |p| { + const kind = p.get("type") orelse continue; + if (kind != .string) continue; + if (std.mem.eql(u8, kind.string, "image_url")) { + const iu = p.get("image_url").?; + const detail = iu.get("detail"); + try out.append(cx.a, .{ .url = iu.get("url").?.string, .low = detail != null and std.mem.eql(u8, detail.?.string, "low") }); + } else if (std.mem.eql(u8, kind.string, "video_url")) { + try out.append(cx.a, .{ .video = true, .url = p.get("video_url").?.get("url").?.string }); + } + } + } + return out.items; +} + +/// --vision-urls: how a public HTTPS image's bytes are fetched (media_fetch.zig, or a test's fixtures). +pub const Fetch = struct { + ctx: ?*anyopaque = null, + io: std.Io, + get: *const fn (ctx: ?*anyopaque, a: std.mem.Allocator, io: std.Io, url: []const u8, max_bytes: usize, deadline: std.Io.Timestamp, media: []const []const u8, f: *media_fetch.Failure) media_fetch.Error![]const u8, + + /// The network, as media_fetch checks it. + pub fn network(io: std.Io) Fetch { + return .{ .io = io, .get = struct { + fn get(_: ?*anyopaque, a: std.mem.Allocator, i: std.Io, url: []const u8, max_bytes: usize, deadline: std.Io.Timestamp, media: []const []const u8, f: *media_fetch.Failure) media_fetch.Error![]const u8 { + return (try media_fetch.fetch(a, i, url, max_bytes, deadline, media, f)).data; + } + }.get }; + } +}; + +/// An image part's bytes, at most `limit`: a data URL's, or with --vision-urls a public HTTPS URL's before `deadline`. +pub fn sourceBytes(cx: *Cx, fetch: ?Fetch, url: []const u8, limit: usize, deadline: std.Io.Timestamp) errors.Refused![]const u8 { + if (std.ascii.startsWithIgnoreCase(url, "data:")) { + const bytes = try dataBytes(cx, url); + if (bytes.len > limit) return cx.refuse("the request's images exceed 20 MB encoded"); + return bytes; + } + if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse("images require data URLs or public HTTPS URLs"); + const f = fetch orelse return cx.refuse("image URLs are off on this server; send the image as a data URL, or start the server with --vision-urls"); + if (url.len > media_fetch.max_url_chars) return cx.refuse("image URL is too long"); + var failure: media_fetch.Failure = .{}; + const bytes = f.get(f.ctx, cx.a, f.io, url, limit, media_fetch.within(f.io, deadline, media_fetch.timeout_s), &media_fetch.image_media, &failure) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Media => return cx.refuse(failure.text), + }; + if (bytes.len == 0) return cx.refuse("image is empty or exceeds the encoded byte limit"); + return bytes; +} + +/// A data URL's bytes (base64 or percent-encoded). +pub fn dataBytes(cx: *Cx, url: []const u8) errors.Refused![]const u8 { + return dataBytesOf(cx, url, max_image_bytes, "image"); +} + +/// A data URL's bytes, at most `max` of them; `noun` names the medium in the refusals. +pub fn dataBytesOf(cx: *Cx, url: []const u8, max: usize, comptime noun: []const u8) errors.Refused![]const u8 { + if (!std.ascii.startsWithIgnoreCase(url, "data:")) return cx.refuse("invalid " ++ noun ++ " data URL"); + const comma = std.mem.indexOfScalar(u8, url, ',') orelse return cx.refuse("invalid " ++ noun ++ " data URL"); + const header = url[5..comma]; + const payload = url[comma + 1 ..]; + if (header.len > 256) return cx.refuse("invalid " ++ noun ++ " data URL"); + var it = std.mem.splitScalar(u8, header, ';'); + _ = it.next(); + const enc = it.next(); + if (it.next() != null or (enc != null and !std.ascii.eqlIgnoreCase(enc.?, "base64"))) return cx.refuse(noun ++ " data URL supports only optional base64 encoding"); + if (enc != null) { + if (payload.len > 4 * ((max + 2) / 3)) return cx.refuse(noun ++ " data URL exceeds the encoded byte limit"); + const d = std.base64.standard.Decoder; + const n = d.calcSizeForSlice(payload) catch return cx.refuse("invalid " ++ noun ++ " data URL encoding"); + const out = try cx.a.alloc(u8, n); + d.decode(out, payload) catch return cx.refuse("invalid " ++ noun ++ " data URL encoding"); + if (out.len == 0 or out.len > max) return cx.refuse(noun ++ " is empty or exceeds the encoded byte limit"); + return out; + } + if (payload.len > max * 3) return cx.refuse(noun ++ " data URL exceeds the encoded byte limit"); + var out: std.ArrayList(u8) = .empty; + var i: usize = 0; + while (i < payload.len) : (i += 1) { + if (payload[i] == '%' and i + 2 < payload.len) { + const b = std.fmt.parseInt(u8, payload[i + 1 .. i + 3], 16) catch return cx.refuse("invalid " ++ noun ++ " data URL encoding"); + try out.append(cx.a, b); + i += 2; + } else try out.append(cx.a, payload[i]); + } + if (out.items.len == 0 or out.items.len > max) return cx.refuse(noun ++ " is empty or exceeds the encoded byte limit"); + return out.items; +} + +pub const Format = enum { png, jpeg, gif, webp, other }; + +/// The image's format from its first bytes, before any decoder sees it. +pub fn sniff(bytes: []const u8) Format { + if (std.mem.startsWith(u8, bytes, "\x89PNG\r\n\x1a\n")) return .png; + if (std.mem.startsWith(u8, bytes, "\xff\xd8\xff")) return .jpeg; + if (std.mem.startsWith(u8, bytes, "GIF87a") or std.mem.startsWith(u8, bytes, "GIF89a")) return .gif; + if (bytes.len >= 12 and std.mem.eql(u8, bytes[0..4], "RIFF") and std.mem.eql(u8, bytes[8..12], "WEBP")) return .webp; + return .other; +} + +fn refusal(cx: *Cx, i: usize, err: anyerror) errors.Refused { + return switch (err) { + error.OutOfMemory => error.OutOfMemory, + error.ImageTooLarge => cx.refuse(try std.fmt.allocPrint(cx.a, "image {d} exceeds the limits: 10 MB encoded, 8192 pixels a side, 16 megapixels", .{i + 1})), + error.Animated => cx.refuse(try std.fmt.allocPrint(cx.a, "image {d}: animated and multipage images are unsupported; send a single frame", .{i + 1})), + error.ImageTooSmall => cx.refuse(try std.fmt.allocPrint(cx.a, "image {d}: the visual-token budget is too small for an image", .{i + 1})), + else => cx.refuse(try std.fmt.allocPrint(cx.a, "image {d} could not be decoded; send PNG or JPEG", .{i + 1})), + }; +} + +/// A quoted marker's stand-in while the template renders: two private-use code points, the second naming the marker. +const sentinel = "\u{F0000}"; + +/// The image markers' text (Vision.markers), longest first so one never matches inside another. +pub fn markerText(text: model_text.Text, a: std.mem.Allocator, v: api.Vision) ![]const []const u8 { + var out: std.ArrayList([]const u8) = .empty; + for (v.markers) |id| { + const s = try text.tokenString(a, id); + if (s.len > 0) try out.append(a, s); + } + std.mem.sort([]const u8, out.items, {}, struct { + fn longer(_: void, x: []const u8, y: []const u8) bool { + return x.len > y.len; + } + }.longer); + return out.items; +} + +/// `messages` with each marker the text spells replaced by its sentinel, so only the template's become images (#511). +pub fn shield(a: std.mem.Allocator, v: Value, markers: []const []const u8) !Value { + return switch (v) { + .string => |s| .{ .string = try shieldText(a, s, markers) }, + .array => |items| blk: { + const out = try a.alloc(Value, items.len); + for (items, out) |x, *o| o.* = try shield(a, x, markers); + break :blk .{ .array = out }; + }, + .object => |o| blk: { + const out = try json.newObject(a); + var it = o.iterator(); + while (it.next()) |kv| try out.put(a, kv.key_ptr.*, try shield(a, kv.value_ptr.*, markers)); + break :blk .{ .object = out }; + }, + else => v, + }; +} + +fn shieldText(a: std.mem.Allocator, s: []const u8, markers: []const []const u8) ![]const u8 { + var out: std.ArrayList(u8) = .empty; + var i: usize = 0; + var changed = false; + outer: while (i < s.len) { + for (markers, 0..) |m, k| if (std.mem.startsWith(u8, s[i..], m)) { + try out.appendSlice(a, sentinel); + try out.appendSlice(a, try std.fmt.allocPrint(a, "{u}", .{@as(u21, @intCast(0xF0001 + k))})); + i += m.len; + changed = true; + continue :outer; + }; + try out.append(a, s[i]); + i += 1; + } + return if (changed) out.items else s; +} + +/// The rendered prompt's ids, each sentinel's marker as ordinary text tokens; otherwise exactly `encode`'s ids. +pub fn encodeShielded(text: model_text.Text, a: std.mem.Allocator, rendered: []const u8, markers: []const []const u8) model_text.Error![]u32 { + var ids: std.ArrayList(u32) = .empty; + var rest = rendered; + while (std.mem.indexOf(u8, rest, sentinel)) |at| { + const after = rest[at + sentinel.len ..]; + const cp: u21 = if (after.len >= 4) std.unicode.utf8Decode(after[0..4]) catch 0 else 0; + const k: usize = if (cp >= 0xF0001) cp - 0xF0001 else markers.len; + if (k >= markers.len) { // not one of ours: the code point is ordinary text + try ids.appendSlice(a, try text.encode(a, rest[0 .. at + sentinel.len], false)); + rest = after; + continue; + } + try ids.appendSlice(a, try text.encode(a, rest[0..at], false)); + try ids.appendSlice(a, try text.encodePlain(a, markers[k])); + rest = after[4..]; + } + try ids.appendSlice(a, try text.encode(a, rest, false)); + return ids.items; +} + +/// The prompt with image placeholders expanded to tokens and video ones to frame groups, its images and cache ids. +pub fn expand(vision: ?api.Vision, fetch: ?Fetch, text: model_text.Text, cx: *Cx, messages: Value, ids: []const u32, history_len: usize) errors.Refused!?Expanded { + const v = vision orelse return null; + const ps = try parts(cx, messages); + var n_images: usize = 0; + for (ps) |p| n_images += @intFromBool(!p.video); + const n_videos = ps.len - n_images; + const video_token: ?u32 = if (v.video) |o| o.video_token else null; + // the template's placeholders alone, one per part (text that spells a marker was shielded) + var marks: usize = 0; + var video_marks: usize = 0; + for (ids) |t| { + marks += @intFromBool(t == v.image_token); + video_marks += @intFromBool(video_token != null and t == video_token.?); + } + if (marks != n_images) return cx.fail(.server, "the chat template wrote {d} image placeholders for {d} image parts", .{ marks, n_images }); + if (video_marks != n_videos) return cx.fail(.server, "the chat template wrote {d} video placeholders for {d} video parts", .{ video_marks, n_videos }); + if (ps.len == 0) return null; + if (n_images > v.max_images) return cx.refuse(try std.fmt.allocPrint(cx.a, "a request may hold at most {d} images (--vision-max-images); it has {d}", .{ v.max_images, n_images })); + if (n_videos > videos.max_videos) return cx.refuse(try std.fmt.allocPrint(cx.a, "a request supports at most {d} videos", .{videos.max_videos})); + const budget: u32 = if (n_images > 0) v.image_tokens / @as(u32, @intCast(n_images)) else 0; + if (n_images > 0 and budget < 1) return cx.refuse("the image count exceeds the visual-token budget"); + const prepared = try cx.a.alloc(api.PreparedImage, n_images); + const clips = try cx.a.alloc(api.PreparedVideo, n_videos); + var bytes_total: usize = 0; + // every URL's download inside one request's seconds (images.py total_timeout_seconds, videos.py's) + const deadline: std.Io.Timestamp = if (fetch) |f| std.Io.Clock.awake.now(f.io).addDuration(.fromSeconds(media_fetch.total_s)) else .{ .nanoseconds = 0 }; + var video_budget: videos.Budget = .{ .deadline = if (fetch) |f| std.Io.Clock.awake.now(f.io).addDuration(.fromSeconds(videos.total_s)) else .{ .nanoseconds = 0 } }; + var ii: usize = 0; + var vi: usize = 0; + for (ps) |p| { + if (p.video) { + const o = v.video.?; + clips[vi] = try videos.prepare(cx, o, v.ctx, fetch, p.url, vi, o.video_tokens / @as(u32, @intCast(n_videos)), &video_budget); + vi += 1; + continue; + } + const i = ii; + ii += 1; + const bytes = try sourceBytes(cx, fetch, p.url, @min(max_image_bytes, max_request_bytes - bytes_total), deadline); + bytes_total += bytes.len; + switch (sniff(bytes)) { // the Python frontend's formats (images.py: JPEG, PNG, WebP), WebP not yet decoded here + .png, .jpeg => {}, + .webp => return cx.refuse(try std.fmt.allocPrint(cx.a, "image {d} is WebP, which this server does not decode yet; send PNG or JPEG", .{i + 1})), + .gif, .other => return cx.refuse(try std.fmt.allocPrint(cx.a, "image {d} is not PNG or JPEG; send one of those", .{i + 1})), + } + const cap = if (p.low) @min(budget, low_detail_tokens) else budget; + prepared[i] = v.prepare(v.ctx, cx.a, bytes, cap) catch |e| return refusal(cx, i, e); + if (prepared[i].tokens == 0 or prepared[i].tokens > cap) return cx.refuse("the image processor exceeded the per-image visual-token budget"); + } + var out_ids: std.ArrayList(u32) = .empty; + var key: std.ArrayList(u32) = .empty; + var imgs: std.ArrayList(api.Image) = .empty; + var hist = history_len; + var k: usize = 0; + var kv: usize = 0; + for (ids, 0..) |t, src| { + const before = out_ids.items.len; + if (t == v.image_token) { + const img = prepared[k]; + k += 1; + try placeImage(cx.a, &out_ids, &key, &imgs, t, img); + } else if (video_token != null and t == video_token.?) { + const o = v.video.?; + const clip = clips[kv]; + kv += 1; + for (clip.groups, clip.times) |g, secs| { + try out_ids.append(cx.a, o.group_open); + try key.append(cx.a, o.group_open); + try placeImage(cx.a, &out_ids, &key, &imgs, v.image_token, g); + try out_ids.append(cx.a, o.group_close); + try key.append(cx.a, o.group_close); + var buf: [64]u8 = undefined; + const words = videos.seconds(&buf, secs) catch unreachable; // a time within the hour's limit + const time = text.encodePlain(cx.a, words) catch |e| return if (e == error.OutOfMemory) error.OutOfMemory else cx.fail(.server, "the tokenizer could not encode a video group's time", .{}); + try out_ids.appendSlice(cx.a, time); + try key.appendSlice(cx.a, time); + } + } else { + try out_ids.append(cx.a, t); + try key.append(cx.a, t); + } + if (src < history_len) hist += out_ids.items.len - before - 1; + } + return .{ .ids = out_ids.items, .history_len = hist, .images = imgs.items, .cache_key = key.items }; +} + +/// Appends one image's placeholder rows, its engine entry and its cache ids (past the vocabulary, by content and row). +fn placeImage(a: std.mem.Allocator, ids: *std.ArrayList(u32), key: *std.ArrayList(u32), imgs: *std.ArrayList(api.Image), token: u32, img: api.PreparedImage) !void { + try imgs.append(a, .{ .pixels = img.pixels, .gh = img.gh, .gw = img.gw, .at = @intCast(ids.items.len), .tokens = img.tokens }); + for (0..img.tokens) |j| { + try ids.append(a, token); + try key.append(a, 0x8000_0000 | @as(u32, @truncate(std.hash.Wyhash.hash(img.hash, std.mem.asBytes(&j)) & 0x7fff_ffff))); + } +} + +test "data URLs decode; remote URLs are refused" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var cx: Cx = .{ .a = arena.allocator() }; + try std.testing.expectEqualStrings("hi!", try dataBytes(&cx, "data:image/png;base64,aGkh")); + try std.testing.expectEqualStrings("a b", try dataBytes(&cx, "data:text/plain,a%20b")); + try std.testing.expectError(error.Refused, dataBytes(&cx, "data:image/png;base64,***")); + // without --vision-urls an https URL is refused with the flag's name; other schemes always + const never: std.Io.Timestamp = .{ .nanoseconds = 0 }; + try std.testing.expectError(error.Refused, sourceBytes(&cx, null, "https://example.com/x.png", max_image_bytes, never)); + try std.testing.expectEqualStrings("image URLs are off on this server; send the image as a data URL, or start the server with --vision-urls", cx.message); + for ([_][]const u8{ "http://example.com/x.png", "ftp://example.com/x.png", "HTTPS://example.com/x.png", "file:///etc/passwd" }) |u| { + try std.testing.expectError(error.Refused, sourceBytes(&cx, null, u, max_image_bytes, never)); + try std.testing.expectEqualStrings("images require data URLs or public HTTPS URLs", cx.message); + } + try std.testing.expectError(error.Refused, sourceBytes(&cx, null, "data:,abc", 2, never)); +} + +fn fakePrepare(ctx: *anyopaque, a: std.mem.Allocator, bytes: []const u8, max_tokens: u32) anyerror!api.PreparedImage { + _ = ctx; + _ = max_tokens; + const tokens: u32 = @intCast(bytes.len - 3); // the test's images: a JPEG's first bytes, then one token a byte + return .{ .pixels = try a.alloc(f32, 0), .gh = 2, .gw = 2 * tokens, .tokens = tokens, .hash = bytes[3] }; +} + +test "placeholders expand to each image's tokens, the history and cache key with them" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var cx: Cx = .{ .a = a }; + var dummy: u8 = 0; + const v: api.Vision = .{ .ctx = &dummy, .prepare = fakePrepare, .image_token = 9, .image_tokens = 64, .max_images = 4 }; + const text = + \\[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:,%FF%D8%FFab"}}, {"type": "text", "text": "x"}]}, + \\ {"role": "assistant", "content": "y"}, + \\ {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:,%FF%D8%FFcde"}}]}] + ; + const msgs = switch (try json.parseText(a, text)) { + .ok => |m| m, + .err => return error.TestUnexpectedResult, + }; + const ids = [_]u32{ 1, 9, 2, 3, 9, 4 }; + const x = (try expand(v, null, MarkText.text(), &cx, msgs, &ids, 4)).?; + try std.testing.expectEqualSlices(u32, &.{ 1, 9, 9, 2, 3, 9, 9, 9, 4 }, x.ids); + try std.testing.expectEqual(@as(usize, 5), x.history_len); // the first image is in the history, the second not + try std.testing.expectEqual(@as(u32, 1), x.images[0].at); + try std.testing.expectEqual(@as(u32, 5), x.images[1].at); + try std.testing.expect(x.cache_key[1] >= 0x8000_0000 and x.cache_key[1] != x.cache_key[2]); + try std.testing.expectEqual(@as(u32, 2), x.cache_key[3]); + // a placeholder the images do not account for is refused + try std.testing.expectError(error.Refused, expand(v, null, MarkText.text(), &cx, msgs, &[_]u32{ 9, 9, 9 }, 0)); + _ = messages_mod; +} + +test "--vision-urls: an https part's bytes come from the fetch, within what the request has left" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var cx: Cx = .{ .a = a }; + var dummy: u8 = 0; + const v: api.Vision = .{ .ctx = &dummy, .prepare = fakePrepare, .image_token = 9, .image_tokens = 64, .max_images = 4 }; + const Pages = struct { + var limits: [4]usize = undefined; + var n: usize = 0; + fn get(_: ?*anyopaque, al: std.mem.Allocator, _: std.Io, url: []const u8, max_bytes: usize, _: std.Io.Timestamp, _: []const []const u8, f: *media_fetch.Failure) media_fetch.Error![]const u8 { + limits[n] = max_bytes; + n += 1; + if (std.mem.eql(u8, url, "https://e.com/a.jpg")) return al.dupe(u8, "\xff\xd8\xffxyz"); + f.text = "image download returned HTTP 404"; + return error.Media; + } + }; + const fetch: Fetch = .{ .io = std.testing.io, .get = Pages.get }; + const text = + \\[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:,%FF%D8%FFab"}}, + \\ {"type": "image_url", "image_url": {"url": "https://e.com/a.jpg"}}]}] + ; + const msgs = switch (try json.parseText(a, text)) { + .ok => |m| m, + .err => return error.TestUnexpectedResult, + }; + const x = (try expand(v, fetch, MarkText.text(), &cx, msgs, &[_]u32{ 1, 9, 9, 2 }, 0)).?; + try std.testing.expectEqualSlices(u32, &.{ 1, 9, 9, 9, 9, 9, 2 }, x.ids); // 2 tokens, then the fetched image's 3 + try std.testing.expectEqual(@as(u32, 3), x.images[1].at); + try std.testing.expectEqual(@as(usize, max_image_bytes), Pages.limits[0]); // 5 bytes used of 20 MB: 10 MB left + // the fetch's refusal is the request's; without the fetch the URL is refused by the flag's name + const missing = switch (try json.parseText(a, "[{\"role\": \"user\", \"content\": [{\"type\": \"image_url\", \"image_url\": {\"url\": \"https://e.com/gone.png\"}}]}]")) { + .ok => |m| m, + .err => return error.TestUnexpectedResult, + }; + try std.testing.expectError(error.Refused, expand(v, fetch, MarkText.text(), &cx, missing, &[_]u32{9}, 0)); + try std.testing.expectEqualStrings("image download returned HTTP 404", cx.message); + try std.testing.expectError(error.Refused, expand(v, null, MarkText.text(), &cx, missing, &[_]u32{9}, 0)); + try std.testing.expectEqualStrings("image URLs are off on this server; send the image as a data URL, or start the server with --vision-urls", cx.message); +} + +/// A test tokenizer: "<|image|>" is token 9 unless read plainly, every other byte its own id plus 100. +const MarkText = struct { + fn text() model_text.Text { + return .{ .ctx = undefined, .vtable = &.{ .encode = encode, .decode = undefined, .token_id = undefined, .token_string = tokenString, .vocab_size = undefined, .eos_ids = undefined, .render = undefined, .template_source = undefined, .encode_plain = plain } }; + } + fn encode(_: *anyopaque, a: std.mem.Allocator, s: []const u8, _: bool) model_text.Error![]u32 { + var out: std.ArrayList(u32) = .empty; + var i: usize = 0; + while (i < s.len) { + if (std.mem.startsWith(u8, s[i..], "<|image|>")) { + try out.append(a, 9); + i += 9; + } else { + try out.append(a, @as(u32, s[i]) + 100); + i += 1; + } + } + return out.items; + } + fn plain(_: *anyopaque, a: std.mem.Allocator, s: []const u8) model_text.Error![]u32 { + const out = try a.alloc(u32, s.len); + for (out, s) |*o, c| o.* = @as(u32, c) + 100; + return out; + } + fn tokenString(_: *anyopaque, a: std.mem.Allocator, id: u32) std.mem.Allocator.Error![]u8 { + return a.dupe(u8, if (id == 9) "<|image|>" else ""); + } +}; + +test "a marker the text spells stays text: only the template's markers are image tokens (#511)" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var dummy: u8 = 0; + const v: api.Vision = .{ .ctx = &dummy, .prepare = fakePrepare, .image_token = 9, .markers = &.{9}, .image_tokens = 64, .max_images = 4 }; + const markers = try markerText(MarkText.text(), a, v); + const conversation = + \\[{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:,ab"}}, {"type": "text", "text": "what is <|image|>?"}]}, + \\ {"role": "tool", "content": "source: x = '<|image|>'"}] + ; + const msgs = switch (try json.parseText(a, conversation)) { + .ok => |m| m, + .err => return error.TestUnexpectedResult, + }; + const shielded = try shield(a, msgs, markers); + // the "template": one marker per image part, then each message's text + var rendered: std.ArrayList(u8) = .empty; + try rendered.appendSlice(a, "<|image|>"); + try rendered.appendSlice(a, shielded.array[0].get("content").?.array[1].get("text").?.string); + try rendered.appendSlice(a, shielded.array[1].get("content").?.string); + const ids = try encodeShielded(MarkText.text(), a, rendered.items, markers); + var marks: usize = 0; + for (ids) |t| marks += @intFromBool(t == 9); + try std.testing.expectEqual(@as(usize, 1), marks); + // the quoted markers read as their characters, and text without one is exactly encode's + try std.testing.expectEqualSlices(u32, try MarkText.encode(undefined, a, "<|image|>what is ", false), ids[0..9]); + try std.testing.expectEqual(@as(u32, '<' + 100), ids[9]); + const plain_text = "no markers here"; + try std.testing.expectEqualSlices(u32, try MarkText.encode(undefined, a, plain_text, false), try encodeShielded(MarkText.text(), a, plain_text, markers)); +} + +test "formats from their first bytes: WebP is refused by name" { + try std.testing.expectEqual(Format.png, sniff("\x89PNG\r\n\x1a\nrest")); + try std.testing.expectEqual(Format.jpeg, sniff("\xff\xd8\xff\xe0")); + try std.testing.expectEqual(Format.gif, sniff("GIF89a...")); + try std.testing.expectEqual(Format.webp, sniff("RIFF\x10\x00\x00\x00WEBPVP8 ")); + try std.testing.expectEqual(Format.other, sniff("\x00\x00\x00\x18ftypheic")); +} + +fn fakeVideo(ctx: *anyopaque, a: std.mem.Allocator, src: api.VideoSource, max_tokens: u32, max_frames: u32) anyerror!api.PreparedVideo { + _ = ctx; + _ = max_frames; + _ = max_tokens; + // the test's videos: one group a 10 frames, 2 tokens a group, each group 0.5 s after the last + const n: usize = @intCast(src.frames / 10); + const groups = try a.alloc(api.PreparedImage, n); + const times = try a.alloc(f64, n); + for (groups, times, 0..) |*g, *t, i| { + g.* = .{ .pixels = try a.alloc(f32, 0), .gh = 2, .gw = 4, .tokens = 2, .hash = 100 + i }; + t.* = @as(f64, @floatFromInt(i)) * 0.5; + } + return .{ .groups = groups, .times = times }; +} + +test "video parts: counted from the request, refused past two, and the template's placeholders checked" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var cx: Cx = .{ .a = a }; + var dummy: u8 = 0; + const offer: api.VideoOffer = .{ .prepare = fakeVideo, .video_token = 8, .group_open = 6, .group_close = 7, .video_tokens = 64, .max_videos = 2, .max_frames = 256 }; + const v: api.Vision = .{ .ctx = &dummy, .prepare = fakePrepare, .image_token = 9, .image_tokens = 64, .max_images = 4, .video = offer }; + const three = switch (try json.parseText(a, + \\[{"role": "user", "content": [{"type": "video_url", "video_url": {"url": "data:video/mp4;base64,YQ=="}}, + \\ {"type": "video_url", "video_url": {"url": "data:video/mp4;base64,YQ=="}}, + \\ {"type": "video_url", "video_url": {"url": "data:video/mp4;base64,YQ=="}}]}] + )) { + .ok => |m| m, + .err => return error.TestUnexpectedResult, + }; + try std.testing.expectError(error.Refused, expand(v, null, MarkText.text(), &cx, three, &[_]u32{ 8, 8, 8 }, 0)); + try std.testing.expectEqualStrings("a request supports at most 2 videos", cx.message); + // a template that wrote fewer video placeholders than the request has parts is the server's fault + try std.testing.expectError(error.Refused, expand(v, null, MarkText.text(), &cx, three, &[_]u32{8}, 0)); + try std.testing.expectEqual(errors.Kind.server, cx.kind); +} diff --git a/zig/src/server/media_fetch.zig b/zig/src/server/media_fetch.zig new file mode 100644 index 0000000000..b9386bcec4 --- /dev/null +++ b/zig/src/server/media_fetch.zig @@ -0,0 +1,491 @@ +//! Public HTTPS images for --vision-urls as TensorFold 0.6.6 fetched them: every address and redirect checked public. +const std = @import("std"); +const Io = std.Io; +const net = std.Io.net; + +pub const Error = error{ OutOfMemory, Media }; + +/// The failure's words for the client (the request is refused with them). +pub const Failure = struct { text: []const u8 = "" }; + +/// images.py's ImageLimits: a URL's characters, the redirects followed, one download's seconds and a request's. +pub const max_url_chars = 4096; +pub const max_redirects = 3; +pub const timeout_s = 10; +pub const total_s = 30; + +/// images_http.MEDIA_TYPES: the declared types an image URL may answer with (WebP is refused after, by its bytes). +pub const image_media = [_][]const u8{ "image/jpeg", "image/png", "image/webp" }; +/// videos.py VIDEO_MEDIA_TYPES. +pub const video_media = [_][]const u8{ "video/mp4", "video/webm", "video/quicktime", "video/x-matroska" }; + +fn fail(f: *Failure, text: []const u8) Error { + f.text = text; + return error.Media; +} + +/// Python ipaddress is_global, less multicast, reserved, 192.0.0.0/24, 168.63.129.16 and IPv6 embedding an IPv4 one. +pub fn publicIp(a: net.IpAddress) bool { + switch (a) { + .ip4 => |v| { + const x = std.mem.readInt(u32, &v.bytes, .big); + const blocked = [_]struct { [4]u8, u5 }{ + .{ .{ 0, 0, 0, 0 }, 8 }, .{ .{ 10, 0, 0, 0 }, 8 }, .{ .{ 100, 64, 0, 0 }, 10 }, .{ .{ 127, 0, 0, 0 }, 8 }, + .{ .{ 169, 254, 0, 0 }, 16 }, .{ .{ 172, 16, 0, 0 }, 12 }, .{ .{ 192, 0, 0, 0 }, 24 }, .{ .{ 192, 0, 2, 0 }, 24 }, + .{ .{ 192, 168, 0, 0 }, 16 }, .{ .{ 198, 18, 0, 0 }, 15 }, .{ .{ 198, 51, 100, 0 }, 24 }, .{ .{ 203, 0, 113, 0 }, 24 }, + .{ .{ 224, 0, 0, 0 }, 4 }, .{ .{ 240, 0, 0, 0 }, 4 }, + }; + for (blocked) |blk| { + const mask = ~@as(u32, 0) << @intCast(32 - @as(u6, blk[1])); + if (x & mask == std.mem.readInt(u32, &blk[0], .big) & mask) return false; + } + return !std.mem.eql(u8, &v.bytes, &.{ 168, 63, 129, 16 }); + }, + .ip6 => |v| { + const b = v.bytes; + // global unicast 2000::/3 less documentation, the 2001::/23 IETF block (Teredo among it) and 6to4 + if (b[0] & 0xe0 != 0x20) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] == 0x0d and b[3] == 0xb8) return false; + if (b[0] == 0x20 and b[1] == 0x01 and b[2] < 0x02) return false; + return !(b[0] == 0x20 and b[1] == 0x02); + }, + } +} + +/// A URL as a request takes it: the lowercased host and the percent-quoted path and query. +pub const Url = struct { host: []const u8, target: []const u8 }; + +/// images_http._url: https, port 443, no credentials, fragment, backslash, whitespace, controls, IDNA or IPv6 host. +pub fn parseUrl(a: std.mem.Allocator, value: []const u8, f: *Failure) Error!Url { + if (value.len > max_url_chars) return fail(f, "image URL is too long or contains whitespace/control characters"); + for (value) |c| if (c <= 32 or c == 127) return fail(f, "image URL is too long or contains whitespace/control characters"); + const bad = "image URL must be HTTPS on port 443, without credentials or a fragment"; + if (!std.ascii.startsWithIgnoreCase(value, "https://")) return fail(f, bad); + if (std.mem.indexOfAny(u8, value, "\\#") != null) return fail(f, bad); + const rest = value["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const authority = rest[0..end]; + if (std.mem.indexOfScalar(u8, authority, '@') != null or std.mem.startsWith(u8, authority, "[")) return fail(f, bad); + var host = authority; + if (std.mem.lastIndexOfScalar(u8, authority, ':')) |c| { + const port = authority[c + 1 ..]; + if (port.len > 0 and !std.mem.eql(u8, port, "443")) return fail(f, bad); + host = authority[0..c]; + } + if (host.len == 0 or std.mem.indexOfScalar(u8, host, '%') != null) return fail(f, bad); + for (host) |c| if (c >= 128) return fail(f, bad); + const lower = try std.ascii.allocLowerString(a, host); + const name = std.mem.trimEnd(u8, lower, "."); + for ([_][]const u8{ "localhost", "metadata.google.internal", "instance-data" }) |n| if (std.mem.eql(u8, name, n)) return fail(f, "image URLs must use public internet hosts"); + // the path and query, percent-quoted as urllib.parse.quote(safe=...) leaves them + var t: std.ArrayList(u8) = .empty; + const tail = rest[end..]; + if (tail.len == 0 or tail[0] == '?') try t.append(a, '/'); + var in_query = false; + var query_at: usize = 0; + for (tail) |c| { + if (c == '?' and !in_query) { + in_query = true; + try t.append(a, '?'); + query_at = t.items.len; + continue; + } + const safe = std.ascii.isAlphanumeric(c) or std.mem.indexOfScalar(u8, "/%:@!$&'()*+,;=-._~", c) != null or (in_query and c == '?'); + if (safe) try t.append(a, c) else try t.print(a, "%{X:0>2}", .{c}); + } + if (in_query and t.items.len == query_at) t.items.len -= 1; // an empty query, as urlsplit drops it + return .{ .host = lower, .target = t.items }; +} + +/// The location a redirect names, resolved against `base` (urljoin's absolute, scheme-relative and relative forms). +pub fn join(a: std.mem.Allocator, base: []const u8, location: []const u8) ![]const u8 { + if (std.ascii.startsWithIgnoreCase(location, "https://") or std.ascii.startsWithIgnoreCase(location, "http://")) return a.dupe(u8, location); + if (std.mem.startsWith(u8, location, "//")) return std.fmt.allocPrint(a, "https:{s}", .{location}); + const rest = base["https://".len..]; + const end = std.mem.indexOfAny(u8, rest, "/?") orelse rest.len; + const origin = base[0 .. "https://".len + end]; + if (std.mem.startsWith(u8, location, "/")) return std.fmt.allocPrint(a, "{s}{s}", .{ origin, location }); + const path = rest[end..][0 .. std.mem.indexOfScalar(u8, rest[end..], '?') orelse rest.len - end]; + const dir = if (std.mem.lastIndexOfScalar(u8, path, '/')) |s| path[0 .. s + 1] else "/"; + return std.fmt.allocPrint(a, "{s}{s}{s}", .{ origin, dir, location }); +} + +/// One download's deadline: `seconds` from now, or the request's `deadline` when that comes first. +pub fn within(io: Io, deadline: Io.Timestamp, seconds: i64) Io.Timestamp { + const one = Io.Clock.awake.now(io).addDuration(.fromSeconds(seconds)); + return if (one.nanoseconds < deadline.nanoseconds) one else deadline; +} + +fn remaining(io: Io, deadline: Io.Timestamp, f: *Failure) Error!i64 { + const left = Io.Clock.awake.now(io).durationTo(deadline).toMilliseconds(); + if (left <= 0) return fail(f, "image download timed out"); + return left; +} + +/// A fetched body and its declared media type. +pub const Body = struct { data: []u8, media: []const u8 }; + +/// What one request gave: the body, or a redirect's location. +pub const Got = union(enum) { body: Body, redirect: []const u8 }; + +/// `url`'s bytes and media type (of `media_types`), at most `max_bytes`, by `deadline`, each redirect checked again. +pub fn fetch(a: std.mem.Allocator, io: Io, url: []const u8, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Body { + return fetchVia(a, io, url, max_bytes, deadline, media_types, f, Network{}); +} + +/// fetch over `via`, whose request(a, io, url, max_bytes, deadline, media_types, f) answers one checked URL. +pub fn fetchVia(a: std.mem.Allocator, io: Io, url_in: []const u8, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure, via: anytype) Error!Body { + var url = url_in; + var hop: usize = 0; + while (true) : (hop += 1) { + const u = try parseUrl(a, url, f); + switch (try via.request(a, io, u, max_bytes, deadline, media_types, f)) { + .body => |b| return b, + .redirect => |location| { + if (hop == max_redirects) return fail(f, "image download has too many redirects"); + url = try join(a, url, location); + }, + } + } +} + +/// The GET a checked URL is asked with. +pub fn ask(w: *Io.Writer, u: Url, accept: []const u8) Io.Writer.Error!void { + try w.print("GET {s} HTTP/1.1\r\nHost: {s}\r\nAccept: {s}\r\nAccept-Encoding: identity\r\nUser-Agent: TensorFold-native\r\nConnection: close\r\n\r\n", .{ u.target, u.host, accept }); + try w.flush(); +} + +/// The response from `r`: a redirect's location, or a 200's body when its encoding, declared type and length pass. +pub fn answer(a: std.mem.Allocator, io: Io, r: *Io.Reader, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Got { + const broke = "image download failed or timed out"; + const too_big = "image response exceeds the encoded byte limit"; + const status_line = r.takeDelimiterInclusive('\n') catch return fail(f, broke); + if (status_line.len < 12 or !std.mem.startsWith(u8, status_line, "HTTP/1.")) return fail(f, broke); + const status = std.fmt.parseInt(u16, status_line[9..12], 10) catch return fail(f, broke); + var location: ?[]const u8 = null; + var media: []const u8 = ""; + var length: ?usize = null; + var chunked = false; + var encoded = false; + var header_bytes: usize = 0; + while (true) { + const line = r.takeDelimiterInclusive('\n') catch return fail(f, broke); + header_bytes += line.len; + if (header_bytes > 64 * 1024) return fail(f, broke); + const l = std.mem.trimEnd(u8, line, "\r\n"); + if (l.len == 0) break; + const colon = std.mem.indexOfScalar(u8, l, ':') orelse continue; + const key = std.mem.trim(u8, l[0..colon], " \t"); + const val = std.mem.trim(u8, l[colon + 1 ..], " \t"); + if (std.ascii.eqlIgnoreCase(key, "location")) { + location = try a.dupe(u8, val); + } else if (std.ascii.eqlIgnoreCase(key, "content-type")) { + media = try std.ascii.allocLowerString(a, std.mem.trim(u8, val[0 .. std.mem.indexOfScalar(u8, val, ';') orelse val.len], " \t")); + } else if (std.ascii.eqlIgnoreCase(key, "content-length")) { + if (val.len == 0) return fail(f, too_big); + for (val) |c| if (!std.ascii.isDigit(c)) return fail(f, too_big); + length = std.fmt.parseInt(usize, val, 10) catch return fail(f, too_big); + } else if (std.ascii.eqlIgnoreCase(key, "transfer-encoding")) { + chunked = std.ascii.findIgnoreCase(val, "chunked") != null; + } else if (std.ascii.eqlIgnoreCase(key, "content-encoding") and !std.ascii.eqlIgnoreCase(val, "identity")) encoded = true; + } + _ = try remaining(io, deadline, f); + switch (status) { + 301, 302, 303, 307, 308 => return .{ .redirect = location orelse return fail(f, "image redirect has no destination") }, + 200 => {}, + else => return fail(f, try std.fmt.allocPrint(a, "image download returned HTTP {d}", .{status})), + } + if (encoded) return fail(f, "compressed HTTP image responses are unsupported"); + for (media_types) |m| { + if (std.mem.eql(u8, m, media)) break; + } else return fail(f, if (media_types.ptr == &image_media) "image URL content type must be JPEG, PNG or WebP" else try std.fmt.allocPrint(a, "media URL content type must be one of {s}", .{try std.mem.join(a, ", ", media_types)})); + if (length) |n| if (n > max_bytes) return fail(f, too_big); + var data: std.ArrayList(u8) = .empty; + if (chunked) { + while (true) { + _ = try remaining(io, deadline, f); + const size_line = r.takeDelimiterInclusive('\n') catch return fail(f, broke); + const hex = std.mem.trim(u8, size_line[0 .. std.mem.indexOfScalar(u8, size_line, ';') orelse size_line.len], " \t\r\n"); + const n = std.fmt.parseInt(usize, hex, 16) catch return fail(f, broke); + if (n == 0) break; + if (n > max_bytes - data.items.len) return fail(f, too_big); + try readInto(a, io, r, &data, n, deadline, f); + const crlf = r.takeDelimiterInclusive('\n') catch return fail(f, broke); + if (std.mem.trimEnd(u8, crlf, "\r\n").len != 0) return fail(f, broke); + } + } else if (length) |n| { + try readInto(a, io, r, &data, n, deadline, f); + } else while (true) { // no length: to the end of the stream (Connection: close) + _ = try remaining(io, deadline, f); + const piece = r.peekGreedy(1) catch |e| switch (e) { + error.EndOfStream => break, + else => return fail(f, broke), + }; + if (piece.len > max_bytes - data.items.len) return fail(f, too_big); + try data.appendSlice(a, piece); + r.toss(piece.len); + } + _ = try remaining(io, deadline, f); + return .{ .body = .{ .data = data.items, .media = media } }; +} + +/// Exactly `n` more bytes from `r` onto `data`, in whatever pieces the reader's buffer holds. +fn readInto(a: std.mem.Allocator, io: Io, r: *Io.Reader, data: *std.ArrayList(u8), n: usize, deadline: Io.Timestamp, f: *Failure) Error!void { + var left = n; + while (left > 0) { + _ = try remaining(io, deadline, f); + const piece = r.peekGreedy(1) catch return fail(f, "image download failed or timed out"); + const take = @min(piece.len, left); + try data.appendSlice(a, piece[0..take]); + r.toss(take); + left -= take; + } +} + +/// A TCP connection to the checked address on port 443, its connect bounded by `ms` (std's connect has no timeout). +fn connect(addr: net.IpAddress, ms: i64, f: *Failure) Error!net.Stream { + const c = std.c; + const failed = "image download failed or timed out"; + var storage: c.sockaddr.storage = undefined; + var len: c.socklen_t = undefined; + switch (addr) { + .ip4 => |v| { + const sa: *c.sockaddr.in = @ptrCast(@alignCast(&storage)); + sa.* = std.mem.zeroes(c.sockaddr.in); + sa.family = c.AF.INET; + sa.port = std.mem.nativeToBig(u16, 443); + sa.addr = @bitCast(v.bytes); + len = @sizeOf(c.sockaddr.in); + }, + .ip6 => |v| { + const sa: *c.sockaddr.in6 = @ptrCast(@alignCast(&storage)); + sa.* = std.mem.zeroes(c.sockaddr.in6); + sa.family = c.AF.INET6; + sa.port = std.mem.nativeToBig(u16, 443); + sa.addr = v.bytes; + len = @sizeOf(c.sockaddr.in6); + }, + } + const fd = c.socket(@as(c_uint, storage.family), c.SOCK.STREAM, 0); + if (fd < 0) return fail(f, failed); + errdefer _ = c.close(fd); + _ = c.fcntl(fd, c.F.SETFD, @as(c_int, c.FD_CLOEXEC)); + const flags = c.fcntl(fd, c.F.GETFL); + const nonblock: c_int = @bitCast(@as(u32, @bitCast(c.O{ .NONBLOCK = true }))); + if (flags < 0 or c.fcntl(fd, c.F.SETFL, flags | nonblock) < 0) return fail(f, failed); + if (c.connect(fd, @ptrCast(&storage), len) != 0) { + if (c.errno(@as(c_int, -1)) != .INPROGRESS) return fail(f, failed); + var pfd = [_]c.pollfd{.{ .fd = fd, .events = c.POLL.OUT, .revents = 0 }}; + if (c.poll(&pfd, 1, @intCast(@min(ms, std.math.maxInt(c_int)))) != 1) return fail(f, failed); + var err: c_int = 0; + var err_len: c.socklen_t = @sizeOf(c_int); + if (c.getsockopt(fd, c.SOL.SOCKET, c.SO.ERROR, &err, &err_len) != 0 or err != 0) return fail(f, failed); + } + if (c.fcntl(fd, c.F.SETFL, flags) < 0) return fail(f, failed); + const tv: c.timeval = .{ .sec = @intCast(@divTrunc(ms, 1000)), .usec = @intCast(@mod(ms, 1000) * 1000) }; + _ = c.setsockopt(fd, c.SOL.SOCKET, c.SO.RCVTIMEO, &tv, @sizeOf(c.timeval)); + _ = c.setsockopt(fd, c.SOL.SOCKET, c.SO.SNDTIMEO, &tv, @sizeOf(c.timeval)); + return .{ .socket = .{ .handle = fd, .address = addr } }; +} + +var bundle_lock: Io.RwLock = .init; +var bundle: std.crypto.Certificate.Bundle = .empty; +var bundle_loaded = false; + +/// The network: every address the host resolves to checked public, then TLS verified for the host on the first. +const Network = struct { + fn request(_: Network, a: std.mem.Allocator, io: Io, u: Url, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Got { + const unresolved = "image host could not be resolved"; + var first: ?net.IpAddress = null; + if (net.IpAddress.parseIp4(u.host, 443)) |literal| { + if (!publicIp(literal)) return fail(f, "image URLs must resolve only to public internet addresses"); + first = literal; + } else |_| { + const name = net.HostName.init(std.mem.trimEnd(u8, u.host, ".")) catch return fail(f, unresolved); + var results: [32]net.HostName.LookupResult = undefined; + var q: Io.Queue(net.HostName.LookupResult) = .init(&results); + name.lookup(io, &q, .{ .port = 443 }) catch return fail(f, unresolved); + while (q.getOneUncancelable(io)) |r| switch (r) { + .address => |addr| { + if (!publicIp(addr)) return fail(f, "image URLs must resolve only to public internet addresses"); + if (first == null) first = addr; + }, + .canonical_name => {}, + } else |_| {} + } + const addr = first orelse return fail(f, unresolved); + { + bundle_lock.lockUncancelable(io); + defer bundle_lock.unlock(io); + if (!bundle_loaded) { + bundle.rescan(std.heap.page_allocator, io, Io.Clock.real.now(io)) catch return fail(f, "image download failed: no CA certificates on the server"); + bundle_loaded = true; + } + } + const stream = try connect(addr, try remaining(io, deadline, f), f); + defer stream.close(io); + const tls_len = std.crypto.tls.Client.min_buffer_len; + const bufs = try a.alloc(u8, 4 * tls_len); + var sr = stream.reader(io, bufs[0..tls_len]); + var sw = stream.writer(io, bufs[tls_len .. 2 * tls_len]); + var entropy: [std.crypto.tls.Client.Options.entropy_len]u8 = undefined; + io.random(&entropy); + var tls = std.crypto.tls.Client.init(&sr.interface, &sw.interface, .{ + .host = .{ .explicit = std.mem.trimEnd(u8, u.host, ".") }, + .ca = .{ .bundle = .{ .gpa = std.heap.page_allocator, .io = io, .lock = &bundle_lock, .bundle = &bundle } }, + .read_buffer = bufs[2 * tls_len .. 3 * tls_len], + .write_buffer = bufs[3 * tls_len ..], + .entropy = &entropy, + .realtime_now = Io.Clock.real.now(io), + .allow_truncation_attacks = true, // HTTP's own length and chunk framing end the body + }) catch return fail(f, "image download failed or timed out"); // images_http.py: an ssl.SSLError is an OSError + const accept = try std.mem.join(a, ", ", media_types); + ask(&tls.writer, u, accept) catch return fail(f, "image download failed or timed out"); + sw.interface.flush() catch return fail(f, "image download failed or timed out"); + return answer(a, io, &tls.reader, max_bytes, deadline, media_types, f); + } +}; + +test "public addresses as Python's ipaddress checks them" { + const ip4 = struct { + fn of(b: [4]u8) net.IpAddress { + return .{ .ip4 = .{ .bytes = b, .port = 443 } }; + } + }.of; + const ip6 = struct { + fn of(text: []const u8) net.IpAddress { + return net.IpAddress.parseIp6(text, 443) catch unreachable; + } + }.of; + for ([_][4]u8{ .{ 8, 8, 8, 8 }, .{ 151, 101, 1, 69 }, .{ 1, 1, 1, 1 } }) |b| try std.testing.expect(publicIp(ip4(b))); + for ([_][4]u8{ .{ 10, 1, 2, 3 }, .{ 127, 0, 0, 1 }, .{ 169, 254, 169, 254 }, .{ 172, 20, 0, 1 }, .{ 192, 168, 1, 1 }, .{ 100, 100, 0, 1 }, .{ 0, 0, 0, 0 }, .{ 224, 0, 0, 1 }, .{ 255, 255, 255, 255 }, .{ 168, 63, 129, 16 }, .{ 192, 0, 0, 9 }, .{ 198, 18, 0, 1 }, .{ 203, 0, 113, 5 } }) |b| + try std.testing.expect(!publicIp(ip4(b))); + try std.testing.expect(publicIp(ip6("2606:4700::1111"))); + for ([_][]const u8{ "::1", "::", "fe80::1", "fc00::1", "ff02::1", "::ffff:8.8.8.8", "2002:808:808::1", "2001::1", "2001:db8::1", "64:ff9b::808:808" }) |t| + try std.testing.expect(!publicIp(ip6(t))); +} + +test "URLs as images_http._url takes them" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var f: Failure = .{}; + const ok = try parseUrl(a, "https://Example.COM:443/img/c.png?x=%20&y=\xc3\xa4?z", &f); + try std.testing.expectEqualStrings("example.com", ok.host); + try std.testing.expectEqualStrings("/img/c.png?x=%20&y=%C3%A4?z", ok.target); + try std.testing.expectEqualStrings("/", (try parseUrl(a, "https://example.com", &f)).target); + try std.testing.expectEqualStrings("/", (try parseUrl(a, "https://example.com?", &f)).target); + try std.testing.expectEqualStrings("/a??", (try parseUrl(a, "https://example.com/a??", &f)).target); + try std.testing.expectEqualStrings("/%5Ba%5D%22%3C", (try parseUrl(a, "https://example.com/[a]\"<", &f)).target); + const cases = [_]struct { []const u8, []const u8 }{ + .{ "https://example.com/a b.png", "image URL is too long or contains whitespace/control characters" }, + .{ "http://example.com/x.png", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https://user:pw@example.com/x", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https://example.com:8443/x", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https://example.com/x#frag", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https://example.com\\@evil/x", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https:///x", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https://[::1]/x", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https://localhost/x", "image URLs must use public internet hosts" }, + .{ "https://LOCALHOST./x", "image URLs must use public internet hosts" }, + .{ "https://metadata.google.internal/x", "image URLs must use public internet hosts" }, + }; + for (cases) |c| { + try std.testing.expectError(error.Media, parseUrl(a, c[0], &f)); + try std.testing.expectEqualStrings(c[1], f.text); + } + const long = try a.alloc(u8, max_url_chars + 1); + @memcpy(long[0.."https://e.com/".len], "https://e.com/"); + @memset(long["https://e.com/".len..], 'a'); + try std.testing.expectError(error.Media, parseUrl(a, long, &f)); + try std.testing.expectEqualStrings("https://e.com/b/c.png", try join(a, "https://e.com/b/a.png?q", "c.png")); + try std.testing.expectEqualStrings("https://e.com/z", try join(a, "https://e.com/b/a.png", "/z")); + try std.testing.expectEqualStrings("https://cdn.e.com/z", try join(a, "https://e.com/b/a.png", "//cdn.e.com/z")); + try std.testing.expectEqualStrings("http://e.com/z", try join(a, "https://e.com/b", "http://e.com/z")); +} + +/// Canned responses by host and path, for the HTTP rules without a network. +const Fixtures = struct { + pages: []const struct { []const u8, []const u8 }, + asked: *std.ArrayList(u8), + + fn request(self: Fixtures, a: std.mem.Allocator, io: Io, u: Url, max_bytes: usize, deadline: Io.Timestamp, media_types: []const []const u8, f: *Failure) Error!Got { + var w: Io.Writer.Allocating = .init(a); + ask(&w.writer, u, try std.mem.join(a, ", ", media_types)) catch return error.OutOfMemory; + try self.asked.appendSlice(a, w.written()); + const key = try std.fmt.allocPrint(a, "{s}{s}", .{ u.host, u.target }); + for (self.pages) |p| if (std.mem.eql(u8, p[0], key)) { + var r: Io.Reader = .fixed(p[1]); + return answer(a, io, &r, max_bytes, deadline, media_types, f); + }; + return fail(f, "image host could not be resolved"); + } +}; + +test "responses as images_http._request reads them: redirects, types, lengths, encodings" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const io = std.testing.io; + var asked: std.ArrayList(u8) = .empty; + const png = "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: 5\r\n\r\n\x89PNG!"; + const fx: Fixtures = .{ .asked = &asked, .pages = &.{ + .{ "e.com/a.png", png }, + .{ "e.com/chunked", "HTTP/1.1 200 OK\r\ncontent-type: Image/JPEG; q=1\r\nTransfer-Encoding: chunked\r\n\r\n3;x\r\n\xff\xd8\xff\r\n2\r\nab\r\n0\r\n\r\n" }, + .{ "e.com/eof", "HTTP/1.0 200 OK\r\nContent-Type: image/jpeg\r\n\r\n\xff\xd8\xffxyz" }, + .{ "e.com/r1", "HTTP/1.1 302 Found\r\nLocation: /r2\r\n\r\n" }, + .{ "e.com/r2", "HTTP/1.1 301 Moved\r\nLocation: https://cdn.e.com/b/x\r\n\r\n" }, + .{ "cdn.e.com/b/x", "HTTP/1.1 307 Temporary\r\nLocation: y.png\r\n\r\n" }, + .{ "cdn.e.com/b/y.png", png }, + .{ "e.com/loop", "HTTP/1.1 302 Found\r\nLocation: /loop\r\n\r\n" }, + .{ "e.com/plain", "HTTP/1.1 302 Found\r\nLocation: http://e.com/a.png\r\n\r\n" }, + .{ "e.com/inner", "HTTP/1.1 302 Found\r\nLocation: https://localhost/a.png\r\n\r\n" }, + .{ "e.com/nowhere", "HTTP/1.1 302 Found\r\n\r\n" }, + .{ "e.com/404", "HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\n\r\n" }, + .{ "e.com/gz", "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Encoding: gzip\r\nContent-Length: 1\r\n\r\nx" }, + .{ "e.com/html", "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\nContent-Length: 1\r\n\r\nx" }, + .{ "e.com/big", "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: 11\r\n\r\n01234567890" }, + .{ "e.com/bigchunk", "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nTransfer-Encoding: chunked\r\n\r\n6\r\n012345\r\n6\r\n678901\r\n0\r\n\r\n" }, + .{ "e.com/bigeof", "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\n\r\n01234567890" }, + .{ "e.com/badlen", "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: +5\r\n\r\n12345" }, + .{ "e.com/short", "HTTP/1.1 200 OK\r\nContent-Type: image/png\r\nContent-Length: 9\r\n\r\n12345" }, + .{ "e.com/junk", "SSH-2.0-OpenSSH\r\n\r\n" }, + } }; + const far = Io.Clock.awake.now(io).addDuration(.fromSeconds(60)); + var f: Failure = .{}; + const one = try fetchVia(a, io, "https://e.com/a.png", 10, far, &image_media, &f, fx); + try std.testing.expectEqualStrings("\x89PNG!", one.data); + try std.testing.expectEqualStrings("image/png", one.media); + try std.testing.expectEqualStrings("GET /a.png HTTP/1.1\r\nHost: e.com\r\nAccept: image/jpeg, image/png, image/webp\r\nAccept-Encoding: identity\r\nUser-Agent: TensorFold-native\r\nConnection: close\r\n\r\n", asked.items); + const ch = try fetchVia(a, io, "https://e.com/chunked", 10, far, &image_media, &f, fx); + try std.testing.expectEqualStrings("\xff\xd8\xffab", ch.data); + try std.testing.expectEqualStrings("image/jpeg", ch.media); + try std.testing.expectEqualStrings("\xff\xd8\xffxyz", (try fetchVia(a, io, "https://e.com/eof", 10, far, &image_media, &f, fx)).data); + // three redirects, each joined and checked again, reach the image; a fourth is refused + asked.clearRetainingCapacity(); + try std.testing.expectEqualStrings("\x89PNG!", (try fetchVia(a, io, "https://e.com/r1", 10, far, &image_media, &f, fx)).data); + try std.testing.expectEqual(@as(usize, 4), std.mem.count(u8, asked.items, "GET ")); + try std.testing.expect(std.mem.indexOf(u8, asked.items, "GET /b/y.png HTTP/1.1\r\nHost: cdn.e.com\r\n") != null); + const refusals = [_]struct { []const u8, []const u8 }{ + .{ "https://e.com/loop", "image download has too many redirects" }, + .{ "https://e.com/plain", "image URL must be HTTPS on port 443, without credentials or a fragment" }, + .{ "https://e.com/inner", "image URLs must use public internet hosts" }, + .{ "https://e.com/nowhere", "image redirect has no destination" }, + .{ "https://e.com/404", "image download returned HTTP 404" }, + .{ "https://e.com/gz", "compressed HTTP image responses are unsupported" }, + .{ "https://e.com/html", "image URL content type must be JPEG, PNG or WebP" }, + .{ "https://e.com/big", "image response exceeds the encoded byte limit" }, + .{ "https://e.com/bigchunk", "image response exceeds the encoded byte limit" }, + .{ "https://e.com/bigeof", "image response exceeds the encoded byte limit" }, + .{ "https://e.com/badlen", "image response exceeds the encoded byte limit" }, + .{ "https://e.com/short", "image download failed or timed out" }, + .{ "https://e.com/junk", "image download failed or timed out" }, + }; + for (refusals) |c| { + try std.testing.expectError(error.Media, fetchVia(a, io, c[0], 10, far, &image_media, &f, fx)); + try std.testing.expectEqualStrings(c[1], f.text); + } + // a deadline already passed refuses before the body + try std.testing.expectError(error.Media, fetchVia(a, io, "https://e.com/a.png", 10, Io.Clock.awake.now(io), &image_media, &f, fx)); + try std.testing.expectEqualStrings("image download timed out", f.text); +} diff --git a/zig/src/server/messages.zig b/zig/src/server/messages.zig index 9559ed784b..9d7d8c47c9 100644 --- a/zig/src/server/messages.zig +++ b/zig/src/server/messages.zig @@ -3,6 +3,7 @@ const std = @import("std"); const json = @import("json.zig"); const errors = @import("errors.zig"); const fields = @import("fields.zig"); +const api = @import("engine_api"); const Value = json.Value; const Cx = errors.Cx; @@ -23,6 +24,60 @@ fn withField(cx: *Cx, o: *const json.Object, key: []const u8, value: Value) !*js /// ``normalize_messages`` (text only): leading system and developer text merged, later ones as ``late_system``; a template that needs a user query gains one user turn after a trailing tool run. pub fn normalize(cx: *Cx, messages: ?Value, late_system: []const u8, needs_user_after_tool: bool) errors.Refused!Value { + return normalizeWith(cx, messages, late_system, needs_user_after_tool, .{}); +} + +/// The media parts an engine reads: image_url parts (``Info.vision``), and video_url parts (its ``video``). +pub const Media = struct { + images: bool = false, + videos: bool = false, + + pub fn of(v: ?api.Vision) Media { + return .{ .images = v != null, .videos = v != null and v.?.video != null }; + } +}; + +/// One video part as the template and video path read it: {"type": "video_url", "video_url": {"url": string}}. +fn videoPart(cx: *Cx, part: Value) errors.Refused!Value { + const vu = part.get("video_url") orelse return cx.refuse("video_url must contain a non-empty url string"); + const u = if (vu == .object) vu.get("url") else null; + if (u == null or u.? != .string or u.?.string.len == 0) return cx.refuse("video_url must contain a non-empty url string"); + for ([_][]const u8{ "image", "images", "image_url", "input_image", "audio", "input_audio", "video" }) |k| if (json.truthyField(part, k)) return cx.refuse("video_url parts cannot contain other media"); + const inner = try json.newObject(cx.a); + try inner.put(cx.a, "url", u.?); + const out = try json.newObject(cx.a); + try out.put(cx.a, "type", .{ .string = "video_url" }); + try out.put(cx.a, "video_url", .{ .object = inner }); + return .{ .object = out }; +} + +/// One image part as the template and the image path read it: {"type": "image_url", "image_url": {"url", "detail"?}}. +fn imagePart(cx: *Cx, part: Value) errors.Refused!Value { + const iu = part.get("image_url") orelse return cx.refuse("an image_url content part must contain image_url"); + var url: []const u8 = ""; + var detail: ?Value = null; + if (iu == .string) url = iu.string else if (iu == .object) { + const u = iu.get("url") orelse return cx.refuse("image_url must contain a url string"); + if (u != .string) return cx.refuse("image_url must contain a url string"); + url = u.string; + detail = iu.get("detail"); + if (detail) |d| if (d != .string or !(std.mem.eql(u8, d.string, "auto") or std.mem.eql(u8, d.string, "low") or std.mem.eql(u8, d.string, "high"))) + return cx.refuse("image_url detail must be auto, low or high"); + } else return cx.refuse("image_url must be a url string or an object with one"); + if (url.len == 0) return cx.refuse("image_url must contain a url string"); + for ([_][]const u8{ "audio", "input_audio", "video", "video_url" }) |k| if (json.truthyField(part, k)) return cx.refuse("image_url parts cannot contain other media"); + const inner = try json.newObject(cx.a); + try inner.put(cx.a, "url", .{ .string = url }); + if (detail) |d| try inner.put(cx.a, "detail", d); + const out = try json.newObject(cx.a); + try out.put(cx.a, "type", .{ .string = "image_url" }); + try out.put(cx.a, "image_url", .{ .object = inner }); + return .{ .object = out }; +} + +/// `normalize`, keeping image_url (user/tool, `media.images`) and video_url (user, `media.videos`) parts in order. +pub fn normalizeWith(cx: *Cx, messages: ?Value, late_system: []const u8, needs_user_after_tool: bool, media: Media) errors.Refused!Value { + const images = media.images; const list = messages orelse return cx.refuse("messages must be a non-empty list"); if (list != .array or list.array.len == 0) return cx.refuse("messages must be a non-empty list"); var out: std.ArrayList(Value) = .empty; @@ -36,14 +91,35 @@ pub fn normalize(cx: *Cx, messages: ?Value, late_system: []const u8, needs_user_ var item: *json.Object = message.object; if (content != null and content.? == .array) { var text: std.ArrayList(u8) = .empty; + var parts: std.ArrayList(Value) = .empty; // with images: text and image parts in order + var has_image = false; for (content.?.array) |part| { - const typed = part == .object and part.get("type") != null and part.get("type").? == .string and std.mem.eql(u8, part.get("type").?.string, "text"); - if (!typed or fields.hasMedia(part)) return cx.refuse("this server accepts text parts only; image, audio and video inputs are unsupported"); + const kind = if (part == .object and part.get("type") != null and part.get("type").? == .string) part.get("type").?.string else ""; + if (images and std.mem.eql(u8, kind, "image_url")) { + // tool results carry screenshots too; the template renders them in its response + if (!std.mem.eql(u8, role.?.string, "user") and !std.mem.eql(u8, role.?.string, "tool")) return cx.refuse("image parts are accepted in user and tool messages only"); + try parts.append(cx.a, try imagePart(cx, part)); + has_image = true; + continue; + } + if (media.videos and std.mem.eql(u8, kind, "video_url")) { + if (!std.mem.eql(u8, role.?.string, "user")) return cx.refuse("video_url parts are supported only in user messages"); + try parts.append(cx.a, try videoPart(cx, part)); + has_image = true; + continue; + } + if (!std.mem.eql(u8, kind, "text") or fields.hasMedia(part)) return cx.refuse(if (media.videos) "content parts must be text, image_url or video_url; audio is unsupported" else if (images) "content parts must be text or image_url; audio and video inputs are unsupported" else "this server accepts text parts only; image, audio and video inputs are unsupported"); const t = part.get("text") orelse return cx.refuse("a text content part must contain a text string"); if (t != .string) return cx.refuse("a text content part must contain a text string"); try text.appendSlice(cx.a, t.string); + if (images) { + const tp = try json.newObject(cx.a); + try tp.put(cx.a, "type", .{ .string = "text" }); + try tp.put(cx.a, "text", t); + try parts.append(cx.a, .{ .object = tp }); + } } - item = try withField(cx, message.object, "content", .{ .string = text.items }); + item = try withField(cx, message.object, "content", if (has_image) Value{ .array = parts.items } else Value{ .string = text.items }); } else if (content == null or content.? == .null) { item = try withField(cx, message.object, "content", .{ .string = "" }); } else if (content.? != .string) { diff --git a/zig/src/server/model_text.zig b/zig/src/server/model_text.zig index 0b36055a7f..b036ca8342 100644 --- a/zig/src/server/model_text.zig +++ b/zig/src/server/model_text.zig @@ -36,6 +36,8 @@ pub const Text = struct { template_source: *const fn (ctx: *anyopaque) []const u8, /// Whether an id is a special token (an added token flagged special); null: none are known. special: ?*const fn (ctx: *anyopaque, id: u32) bool = null, + /// Ids of ``text`` with no added token matched (a spelled special token is plain text); null: ``encode``. + encode_plain: ?*const fn (ctx: *anyopaque, a: Allocator, text: []const u8) Error![]u32 = null, /// A token's own bytes (a partial UTF-8 sequence kept whole) for logprob rows; null: the tokenizer can't say. token_bytes: ?*const fn (ctx: *anyopaque, a: Allocator, id: u32) Allocator.Error!?[]u8 = null, }; @@ -44,6 +46,11 @@ pub const Text = struct { return t.vtable.encode(t.ctx, a, text, add_special); } + pub fn encodePlain(t: Text, a: Allocator, text: []const u8) Error![]u32 { + const f = t.vtable.encode_plain orelse return t.encode(a, text, false); + return f(t.ctx, a, text); + } + pub fn decode(t: Text, a: Allocator, ids: []const u32) Allocator.Error![]u8 { return t.vtable.decode(t.ctx, a, ids); } diff --git a/zig/src/server/openai.zig b/zig/src/server/openai.zig index b8782aaf96..9600028d4c 100644 --- a/zig/src/server/openai.zig +++ b/zig/src/server/openai.zig @@ -127,10 +127,10 @@ fn plan(srv: *Server, cx: *Cx, is_chat: bool, raw: Value) errors.Refused!Plan { if (srv.config.request_log) |path| request_log.append(cx.a, path, body); var input: chat.Input = .{ .fields = undefined }; if (is_chat) { - input.messages = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); + input.messages = try messages.normalizeWith(cx, body.get("messages"), "system", srv.needs_user_after_tool, .of(srv.info.vision)); input.tools = try tool_specs.active(cx, body.get("tools"), body.get("tool_choice")); } else if (body.get("messages")) |m| if (m == .array and m.array.len > 0) { - input.messages = try messages.normalize(cx, m, "system", srv.needs_user_after_tool); + input.messages = try messages.normalizeWith(cx, m, "system", srv.needs_user_after_tool, .of(srv.info.vision)); }; if (!is_chat and input.messages.array.len == 0) input.prompt = try legacyPrompt(srv, cx, body.get("prompt")); // ``body.get("max_tokens") or body.get("max_completion_tokens")``: both are ints or None by now diff --git a/zig/src/server/prompt.zig b/zig/src/server/prompt.zig index e68363bfe7..9ffd87b758 100644 --- a/zig/src/server/prompt.zig +++ b/zig/src/server/prompt.zig @@ -1,15 +1,18 @@ //! Prompt ids from a request, as ``prepare_prompt`` and ``render_prompt_ids`` make them. const std = @import("std"); +const api = @import("engine_api"); const json = @import("json.zig"); const errors = @import("errors.zig"); const messages_mod = @import("messages.zig"); const model_text = @import("model_text.zig"); const chat = @import("chat.zig"); +const images = @import("images.zig"); const Server = @import("server.zig").Server; const Value = json.Value; const Cx = errors.Cx; -pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0 }; +/// With image input, `ids` holds each image's tokens, `images` the images and `cache_key` the prompt the cache matches. +pub const Rendered = struct { ids: []const u32, history_len: usize = 0, rewind_len: usize = 0, images: []const api.Image = &.{}, cache_key: []const u32 = &.{} }; /// Where lessons decide to act on a reply: the row before the generation prompt's ``, else before its last. pub fn decideOf(srv: *const Server, ids: []const u32) u32 { @@ -27,7 +30,7 @@ pub const isTitle = messages_mod.isTitleRequest; /// ``render_prompt_ids``: messages normalized for this template, rendered, and a dangling ```` closed. pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, thinking: bool, effort: ?[]const u8, generation: bool) errors.Refused![]const u32 { - const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalize(cx, messages, srv.late_system, srv.needs_user_after_tool)); + const normalized = try messages_mod.toolArguments(cx, try messages_mod.normalizeWith(cx, messages, srv.late_system, srv.needs_user_after_tool, .of(srv.info.vision))); var problem: []const u8 = ""; const options: model_text.RenderOptions = .{ .tools = if (tools.len > 0) Value{ .array = @constCast(tools) } else null, @@ -35,7 +38,7 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t .enable_thinking = thinking, .reasoning_effort = if (thinking) effort else null, }; - var ids = srv.text.renderIds(cx.a, normalized, options, &problem) catch |e| switch (e) { + var ids = render(srv, cx.a, normalized, options, &problem) catch |e| switch (e) { error.OutOfMemory => return error.OutOfMemory, error.Template => return cx.fail(.server, "{s}", .{problem}), }; @@ -52,6 +55,14 @@ pub fn renderIds(srv: *Server, cx: *Cx, messages: Value, tools: []const Value, t return ids; } +/// The template's ids for `messages`; with image input, a marker the text spells is ordinary text (images.zig). +fn render(srv: *Server, a: std.mem.Allocator, messages: Value, options: model_text.RenderOptions, problem: *[]const u8) model_text.Error![]u32 { + const v = srv.info.vision orelse return srv.text.renderIds(a, messages, options, problem); + const markers = try images.markerText(srv.text, a, v); + const text = try srv.text.render(a, try images.shield(a, messages, markers), options, problem); + return images.encodeShielded(srv.text, a, text, markers); +} + /// The template's generation prompt (thinking on or off): what a probe conversation renders past its history. fn generationPrompt(srv: *Server, cx: *Cx, thinking: bool) errors.Refused![]const u32 { const user = try json.newObject(cx.a); @@ -101,6 +112,8 @@ pub fn prepare(srv: *Server, cx: *Cx, input: chat.Input, thinking: bool, effort: var history_len: usize = 0; if (history.len > 0 and history.len < prompt.len and std.mem.eql(u32, prompt[0..history.len], history)) history_len = history.len; if (history_len == 0 and prompt.len > 1 and std.mem.eql(u32, history, prompt)) history_len = prompt.len - 1; // a template with no generation suffix + // --vision: placeholders expanded to image tokens before the window, chunks and cache see it; no rewind kept + if (try images.expand(srv.info.vision, if (srv.config.vision_urls) images.Fetch.network(srv.io) else null, srv.text, cx, input.messages, prompt, history_len)) |x| return .{ .ids = x.ids, .history_len = x.history_len, .images = x.images, .cache_key = x.cache_key }; const count = if (input.messages == .array) input.messages.array.len else 0; const rewind_len = if (srv.info.prompt_cache) rewindLen(srv, cx, input.messages, count, input.tools, prompt, thinking, effort) else 0; // one more render return .{ .ids = prompt, .history_len = history_len, .rewind_len = rewind_len }; diff --git a/zig/src/server/reply_text.zig b/zig/src/server/reply_text.zig index 648a1d8cd2..991ae959ce 100644 --- a/zig/src/server/reply_text.zig +++ b/zig/src/server/reply_text.zig @@ -292,6 +292,35 @@ pub const Stops = struct { } }; +/// The text a stream has sent: Python's ``streamed = visible``, grown by its delta while it only grows. +pub const Shown = struct { + text: std.ArrayList(u8) = .empty, + chars: usize = 0, + at: ?[*]const u8 = null, // where the last shown text lay; the decode buffer only grows, so that prefix holds + extends: bool = false, + + /// Python's ``now[len(shown):]``; ``stable``: ``now`` lies in the append-only decode buffer. + pub fn after(s: *Shown, now: []const u8, stable: bool) []const u8 { + const sent = s.text.items; + s.extends = (stable and s.at == now.ptr and now.len >= sent.len) or std.mem.startsWith(u8, now, sent); + if (!s.extends) return afterChars(now, s.chars); + if (stable) s.at = now.ptr; // the buffer moved but kept its bytes: the next check is free again + return now[sent.len..]; + } + + pub fn set(s: *Shown, a: Allocator, now: []const u8, delta: []const u8) Allocator.Error!void { + if (s.extends) { + try s.text.appendSlice(a, delta); + s.chars += charCount(delta); + } else { + s.text = .empty; + try s.text.appendSlice(a, now); + s.chars = charCount(now); + } + s.at = now.ptr; + } +}; + /// Decodes only new tokens with the previous chunk as context, so characters split across tokens stay whole. pub const Incremental = struct { tokens: std.ArrayList(u32) = .empty, diff --git a/zig/src/server/root.zig b/zig/src/server/root.zig index 7630f5217a..165c8d4947 100644 --- a/zig/src/server/root.zig +++ b/zig/src/server/root.zig @@ -39,6 +39,9 @@ test { _ = @import("process_memory.zig"); _ = @import("log.zig"); _ = @import("messages.zig"); + _ = @import("images.zig"); + _ = @import("media_fetch.zig"); + _ = @import("videos.zig"); _ = @import("stream_preflight_test.zig"); _ = @import("stream_dsml_test.zig"); _ = @import("logprobs_test.zig"); diff --git a/zig/src/server/serve.zig b/zig/src/server/serve.zig index da27cb6b0e..7cd93678ad 100644 --- a/zig/src/server/serve.zig +++ b/zig/src/server/serve.zig @@ -116,6 +116,7 @@ pub fn run(gpa: Allocator, io: std.Io, args: cli.Args, s: Setup) u8 { .keep_warm_s = args.keep_warm, .default_sampling = s.sampling, .use_drafts = !args.no_drafts, + .vision_urls = args.vision_urls, .seed_salt = salt, .request_log = env(s, "TENSORFOLD_REQUEST_LOG"), .dashboard = args.dashboard, diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 011734ee9a..6fe0ffdcd3 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -46,6 +46,8 @@ pub const Config = struct { /// Sampling when a request names none (generation_config.json and the serve flags), or null for greedy. default_sampling: ?Value = null, use_drafts: bool = true, + /// --vision-urls: image parts may name public HTTPS URLs, fetched before admission (media_fetch.zig). + vision_urls: bool = false, seed_salt: i64 = 0, /// TENSORFOLD_REQUEST_LOG: where chat and completion bodies are appended; null: nowhere. request_log: ?[]const u8 = null, diff --git a/zig/src/server/tokens.zig b/zig/src/server/tokens.zig index 3a2e351921..770389bd8f 100644 --- a/zig/src/server/tokens.zig +++ b/zig/src/server/tokens.zig @@ -4,6 +4,7 @@ const json = @import("json.zig"); const errors = @import("errors.zig"); const fields = @import("fields.zig"); const messages = @import("messages.zig"); +const images = @import("images.zig"); const prompt = @import("prompt.zig"); const tool_specs = @import("tool_specs.zig"); const routes = @import("routes.zig"); @@ -29,7 +30,7 @@ fn tokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { var ids: []const u32 = undefined; if (body.has("messages")) { const generation = try flag(cx, body, "add_generation_prompt", true); - const msgs = try messages.normalize(cx, body.get("messages"), "system", srv.needs_user_after_tool); + const msgs = try messages.normalizeWith(cx, body.get("messages"), "system", srv.needs_user_after_tool, .of(srv.info.vision)); const tools = tool_specs.active(cx, body.get("tools"), body.get("tool_choice")) catch |e| switch (e) { error.Refused => return cx.refuse(cx.message), // a ValueError here is a RequestError else => |x| return x, @@ -40,6 +41,8 @@ fn tokenize(srv: *Server, cx: *Cx, body: Value) errors.Refused!Value { error.Refused => return cx.other(cx.message), else => |x| return x, }; + // media expanded as the chat route expands them: each image's tokens, each video's frame groups + if (try images.expand(srv.info.vision, if (srv.config.vision_urls) images.Fetch.network(srv.io) else null, srv.text, cx, msgs, ids, 0)) |x| ids = x.ids; } else { const text = body.get("prompt"); if (text == null or text.? != .string) return cx.refuse("prompt must be a string (or send messages)"); diff --git a/zig/src/server/videos.zig b/zig/src/server/videos.zig new file mode 100644 index 0000000000..b3707a561e --- /dev/null +++ b/zig/src/server/videos.zig @@ -0,0 +1,134 @@ +//! Video input: video_url parts bounded as 0.6.6's vision/videos.py, decoded by libtfvideo, prepared by the family. +const std = @import("std"); +const api = @import("engine_api"); +const errors = @import("errors.zig"); +const images = @import("images.zig"); +const media_fetch = @import("media_fetch.zig"); +const Cx = errors.Cx; + +/// videos.py's VideoLimits. +pub const max_videos = 2; +pub const max_bytes = 16 << 20; +pub const max_total_bytes = 20 << 20; +pub const max_side = 8192; +pub const max_seconds = 3600; +pub const timeout_s = 20; +pub const total_s = 120; + +/// A request's videos' shared budgets: bytes left and the downloads' deadline. +pub const Budget = struct { + used: usize = 0, + deadline: std.Io.Timestamp, +}; + +/// A video part's bytes: a data URL's, or with --vision-urls (`fetch`) a public HTTPS URL's. +fn bytesOf(cx: *Cx, fetch: ?images.Fetch, url: []const u8, b: *Budget) errors.Refused![]const u8 { + const limit = @min(max_bytes, max_total_bytes - b.used); + if (limit == 0) return cx.refuse("video request exceeds the total byte or time limit"); + const bytes = if (std.ascii.startsWithIgnoreCase(url, "data:")) try images.dataBytesOf(cx, url, limit, "video") else blk: { + if (!std.mem.startsWith(u8, url, "https://")) return cx.refuse("videos require data URLs or public HTTPS URLs"); + const f = fetch orelse return cx.refuse("video URLs are off on this server; send the video as a data URL, or start the server with --vision-urls"); + if (url.len > media_fetch.max_url_chars) return cx.refuse("video URL is too long"); + var failure: media_fetch.Failure = .{}; + break :blk f.get(f.ctx, cx.a, f.io, url, limit, media_fetch.within(f.io, b.deadline, timeout_s), &media_fetch.video_media, &failure) catch |e| switch (e) { + error.OutOfMemory => return error.OutOfMemory, + error.Media => return cx.refuse(failure.text), + }; + }; + if (bytes.len == 0) return cx.refuse("video is empty or exceeds the encoded byte limit"); + b.used += bytes.len; + return bytes; +} + +/// libtfvideo's frames into the engine's sink. +const Decoder = struct { + v: api.tfvideo.Video, + + const Into = struct { + sink: api.FrameSink, + pub fn take(s: Into, k: usize, rgb: []const u8, w: u32, h: u32, stride: usize) anyerror!void { + try s.sink.take(s.sink.ctx, k, rgb, w, h, stride); + } + }; + + fn decode(ctx: *anyopaque, indices: []const u32, sink: api.FrameSink) anyerror!void { + const d: *Decoder = @ptrCast(@alignCast(ctx)); + try d.v.frames(indices, Into{ .sink = sink }); + } +}; + +/// Video `i` (0-based) of the request at `url`, prepared within `max_tokens` tokens. +pub fn prepare(cx: *Cx, offer: api.VideoOffer, ctx: *anyopaque, fetch: ?images.Fetch, url: []const u8, i: usize, max_tokens: u32, b: *Budget) errors.Refused!api.PreparedVideo { + const bytes = try bytesOf(cx, fetch, url, b); + var d: Decoder = .{ .v = api.tfvideo.Video.open(bytes) catch |e| return switch (e) { + error.NoLibrary => cx.fail(.server, "video input needs libtfvideo (FFmpeg) beside the server: zig build video -Dffmpeg=PREFIX, after tools/zig/build_ffmpeg.sh PREFIX", .{}), + error.NoStream => cx.refuse(try std.fmt.allocPrint(cx.a, "video {d}: the video has no video stream", .{i + 1})), + error.OutOfMemory => error.OutOfMemory, + else => cx.refuse(try std.fmt.allocPrint(cx.a, "video {d}: the video's bytes are invalid or unsupported; use MP4 or WebM (H.264, HEVC, VP8, VP9 or MPEG-4)", .{i + 1})), + } }; + defer d.v.close(); + const info = d.v.info; + if (info.width <= 0 or info.height <= 0 or info.width > max_side or info.height > max_side) + return cx.refuse(try std.fmt.allocPrint(cx.a, "video {d}: video dimensions are missing or exceed the pixel limit", .{i + 1})); + if (info.frames <= 0 or info.rate_num <= 0 or info.rate_den <= 0) + return cx.refuse(try std.fmt.allocPrint(cx.a, "video {d}: the video's frame count or rate is missing", .{i + 1})); + const length = @as(f64, @floatFromInt(info.frames)) * @as(f64, @floatFromInt(info.rate_den)) / @as(f64, @floatFromInt(info.rate_num)); + if (length > max_seconds) return cx.refuse(try std.fmt.allocPrint(cx.a, "video {d}: videos are limited to {d} minutes", .{ i + 1, max_seconds / 60 })); + const src: api.VideoSource = .{ .frames = @intCast(info.frames), .rate_num = @intCast(info.rate_num), .rate_den = @intCast(info.rate_den), .width = @intCast(info.width), .height = @intCast(info.height), .ctx = &d, .decode = Decoder.decode }; + return offer.prepare(ctx, cx.a, src, max_tokens, offer.max_frames) catch |e| switch (e) { + error.OutOfMemory => error.OutOfMemory, + error.VideoTooShort => cx.refuse(try std.fmt.allocPrint(cx.a, "video {d} is too short: sampled at 2 frames a second it gives no frame; send at least half a second", .{i + 1})), + error.VideoTooLarge => cx.refuse(try std.fmt.allocPrint(cx.a, "video {d}: the visual-token budget (--vision-video-tokens) is too small for its frames", .{i + 1})), + else => cx.refuse(try std.fmt.allocPrint(cx.a, "video {d}: the video's bytes are invalid or unsupported; use MP4 or WebM (H.264, HEVC, VP8, VP9 or MPEG-4)", .{i + 1})), + }; +} + +test "video bytes: data URLs within the limits; URLs only with --vision-urls" { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena.deinit(); + var cx: Cx = .{ .a = arena.allocator() }; + var b: Budget = .{ .deadline = .{ .nanoseconds = 0 } }; + try std.testing.expectEqualStrings("abc", try bytesOf(&cx, null, "data:video/mp4;base64,YWJj", &b)); + try std.testing.expectEqual(@as(usize, 3), b.used); + try std.testing.expectError(error.Refused, bytesOf(&cx, null, "https://example.com/a.mp4", &b)); + try std.testing.expectEqualStrings("video URLs are off on this server; send the video as a data URL, or start the server with --vision-urls", cx.message); + try std.testing.expectError(error.Refused, bytesOf(&cx, null, "http://example.com/a.mp4", &b)); + try std.testing.expectEqualStrings("videos require data URLs or public HTTPS URLs", cx.message); + b.used = max_total_bytes; + try std.testing.expectError(error.Refused, bytesOf(&cx, null, "data:video/mp4;base64,YWJj", &b)); + try std.testing.expectEqualStrings("video request exceeds the total byte or time limit", cx.message); +} + +/// A frame group's time as the processor writes it: Python's f"{t:.1f} seconds" (one decimal, an exact tie to even). +pub fn seconds(buf: []u8, t: f64) ![]const u8 { + // t * 10 exactly from the double's bits, rounded once, ties to even (no 128-bit floats: no compiler-rt is linked) + const bits: u64 = @bitCast(@abs(t)); + const biased: i32 = @intCast(bits >> 52); + const m: u64 = (bits & ((@as(u64, 1) << 52) - 1)) | (if (biased == 0) 0 else @as(u64, 1) << 52); + const e: i32 = (if (biased == 0) 1 else biased) - 1075; + const x = m * 10; // under 2^57 + var r: u64 = 0; + if (e >= 0) { + r = x << @intCast(@min(e, 6)); // times past 2^56 seconds are not times + } else if (e > -64) { + const s: u6 = @intCast(-e); + const q = x >> s; + const rem = x & ((@as(u64, 1) << s) - 1); + const half = @as(u64, 1) << (s - 1); + r = q + @intFromBool(rem > half or (rem == half and q & 1 == 1)); + } + return std.fmt.bufPrint(buf, "{d}.{d} seconds", .{ r / 10, r % 10 }); +} + +test "group times as Python's f\"{t:.1f} seconds\"" { + var buf: [32]u8 = undefined; + try std.testing.expectEqualStrings("0.2 seconds", try seconds(&buf, 0.25)); + try std.testing.expectEqualStrings("0.8 seconds", try seconds(&buf, 0.75)); + try std.testing.expectEqualStrings("149.3 seconds", try seconds(&buf, 149.3)); + // the double's exact value decides, as Python's: 0.15 is below a half, 1.25 an exact tie + const cases = [_]struct { f64, []const u8 }{ .{ 0.05, "0.1" }, .{ 0.15, "0.1" }, .{ 0.35, "0.3" }, .{ 0.45, "0.5" }, .{ 1.25, "1.2" }, .{ 0.95, "0.9" }, .{ 9.95, "9.9" }, .{ 99.95, "100.0" }, .{ 0.0, "0.0" }, .{ 3599.95, "3599.9" } }; + for (cases) |c| { + var want: [32]u8 = undefined; + try std.testing.expectEqualStrings(try std.fmt.bufPrint(&want, "{s} seconds", .{c[1]}), try seconds(&buf, c[0])); + } +} diff --git a/zig/src/video/tf_video.c b/zig/src/video/tf_video.c new file mode 100644 index 0000000000..eb626b8c82 --- /dev/null +++ b/zig/src/video/tf_video.c @@ -0,0 +1,192 @@ +// libtfvideo: a video's stream facts and RGB frames as transformers' read_video_pyav reads them, over FFmpeg 9.0.2. +#include +#include +#include +#include +#include +#include +#include + +typedef struct { + int64_t frames; // the stream's nb_frames, or (counted = 1) its packets when the container doesn't say + int rate_num, rate_den; // the stream's average frame rate (0/1 when unknown) + int width, height; + int counted; +} tf_video_info; + +// One decoded frame: `rgb` holds height rows of width * 3 bytes, `stride` bytes apart; nonzero stops the decode. +typedef int (*tf_frame_fn)(void* ctx, int k, int index, const uint8_t* rgb, int width, int height, int stride); + +enum { + TF_OK = 0, + TF_INVALID = -1, // not a container FFmpeg reads here + TF_NO_STREAM = -2, // no video stream + TF_NO_DECODER = -3, // a codec this build doesn't decode + TF_MEMORY = -4, + TF_DECODE = -5, // the stream failed to decode + TF_SHORT = -6, // the stream ended before a wanted index + TF_STOPPED = -7, // the callback stopped it +}; + +typedef struct { const uint8_t* data; size_t len, pos; } mem_t; + +typedef struct tf_video { + mem_t mem; + AVFormatContext* fmt; + AVIOContext* io; + int stream; +} tf_video; + +static int mem_read(void* o, uint8_t* buf, int n) { + mem_t* m = (mem_t*)o; + if (m->pos >= m->len) return AVERROR_EOF; + size_t k = m->len - m->pos < (size_t)n ? m->len - m->pos : (size_t)n; + memcpy(buf, m->data + m->pos, k); + m->pos += k; + return (int)k; +} + +static int64_t mem_seek(void* o, int64_t off, int whence) { + mem_t* m = (mem_t*)o; + if (whence == AVSEEK_SIZE) return (int64_t)m->len; + whence &= ~AVSEEK_FORCE; + int64_t at = whence == SEEK_SET ? off : whence == SEEK_CUR ? (int64_t)m->pos + off : (int64_t)m->len + off; + if (at < 0 || at > (int64_t)m->len) return -1; + m->pos = (size_t)at; + return at; +} + +// PyAV's _set_frame_colorspace(frame, colorspace, UNSPECIFIED): an SWS_CS_* value names the frame's colour space +static void set_colorspace(AVFrame* f, int cs) { + if (cs == SWS_CS_ITU709) f->colorspace = AVCOL_SPC_BT709; + else if (cs == SWS_CS_FCC) f->colorspace = AVCOL_SPC_FCC; + else if (cs == SWS_CS_ITU601) f->colorspace = AVCOL_SPC_SMPTE170M; + else if (cs == SWS_CS_SMPTE240M) f->colorspace = AVCOL_SPC_SMPTE240M; + else if (cs == SWS_CS_BT2020) f->colorspace = AVCOL_SPC_BT2020_NCL; +} + +const char* tf_video_version(void) { return FFMPEG_VERSION; } + +void tf_video_close(tf_video* v) { + if (!v) return; + if (v->fmt) avformat_close_input(&v->fmt); + if (v->io) { av_freep(&v->io->buffer); avio_context_free(&v->io); } + free(v); +} + +// The container in `data` (which must outlive the handle) and its first video stream's facts. +int tf_video_open(const uint8_t* data, size_t len, tf_video** out, tf_video_info* info) { + *out = NULL; + memset(info, 0, sizeof(*info)); + av_log_set_level(AV_LOG_QUIET); + tf_video* v = calloc(1, sizeof(tf_video)); + if (!v) return TF_MEMORY; + v->mem = (mem_t){data, len, 0}; + unsigned char* buf = av_malloc(1 << 16); + if (!buf) { free(v); return TF_MEMORY; } + v->io = avio_alloc_context(buf, 1 << 16, 0, &v->mem, mem_read, NULL, mem_seek); + if (!v->io) { av_free(buf); free(v); return TF_MEMORY; } + v->fmt = avformat_alloc_context(); + if (!v->fmt) { tf_video_close(v); return TF_MEMORY; } + v->fmt->pb = v->io; + if (avformat_open_input(&v->fmt, NULL, NULL, NULL) < 0) { v->fmt = NULL; tf_video_close(v); return TF_INVALID; } + if (avformat_find_stream_info(v->fmt, NULL) < 0) { tf_video_close(v); return TF_INVALID; } + v->stream = -1; // PyAV's streams.video[0]: the first video stream + for (unsigned i = 0; i < v->fmt->nb_streams; ++i) + if (v->fmt->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO) { v->stream = (int)i; break; } + if (v->stream < 0) { tf_video_close(v); return TF_NO_STREAM; } + AVStream* st = v->fmt->streams[v->stream]; + info->frames = st->nb_frames; + info->rate_num = st->avg_frame_rate.num; + info->rate_den = st->avg_frame_rate.den; + info->width = st->codecpar->width; + info->height = st->codecpar->height; + if (!avcodec_find_decoder(st->codecpar->codec_id)) { tf_video_close(v); return TF_NO_DECODER; } + if (info->frames <= 0) { // Matroska and WebM keep no count: the stream's packets, one a shown frame, then back + AVPacket* pkt = av_packet_alloc(); + if (!pkt) { tf_video_close(v); return TF_MEMORY; } + int64_t n = 0; + while (av_read_frame(v->fmt, pkt) >= 0) { + if (pkt->stream_index == v->stream) ++n; + av_packet_unref(pkt); + } + av_packet_free(&pkt); + info->frames = n; + info->counted = 1; + } + *out = v; + return TF_OK; +} + +// The frames at `indices` (ascending, distinct), decoded from the start, each to `fn` as RGB; returns frames delivered. +int tf_video_frames(tf_video* v, const int* indices, int n, tf_frame_fn fn, void* ctx) { + if (n <= 0) return 0; + int rc = TF_DECODE; + AVStream* st = v->fmt->streams[v->stream]; + const AVCodec* codec = avcodec_find_decoder(st->codecpar->codec_id); + AVCodecContext* dec = avcodec_alloc_context3(codec); + AVPacket* pkt = av_packet_alloc(); + AVFrame* frame = av_frame_alloc(); + AVFrame* rgb = av_frame_alloc(); + struct SwsContext* sws = sws_alloc_context(); + if (!dec || !pkt || !frame || !rgb || !sws) { rc = TF_MEMORY; goto done; } + if (avcodec_parameters_to_context(dec, st->codecpar) < 0) goto done; + dec->pkt_timebase = st->time_base; + dec->thread_count = 0; // PyAV's stream codec context: every CPU, slice threads + dec->thread_type = FF_THREAD_SLICE; + if (avcodec_open2(dec, codec, NULL) < 0) goto done; + if (av_seek_frame(v->fmt, v->stream, 0, AVSEEK_FLAG_BACKWARD) < 0) avformat_seek_file(v->fmt, -1, INT64_MIN, 0, 0, 0); + sws->threads = 0; + sws->flags = SWS_BILINEAR; + int at = 0, next = 0, flushing = 0; + while (next < n) { + if (!flushing) { + int e = av_read_frame(v->fmt, pkt); + if (e < 0) { flushing = 1; avcodec_send_packet(dec, NULL); } + else { + if (pkt->stream_index == v->stream) avcodec_send_packet(dec, pkt); + av_packet_unref(pkt); + } + } + for (;;) { + int e = avcodec_receive_frame(dec, frame); + if (e == AVERROR(EAGAIN)) break; + if (e == AVERROR_EOF) { rc = TF_SHORT; goto done; } + if (e < 0) goto done; + if (at == indices[next]) { + // PyAV's _reformat: props copied, colour space and range via _set_frame_colorspace, trc/primaries unset + av_frame_unref(rgb); + av_frame_copy_props(rgb, frame); + set_colorspace(frame, (int)frame->colorspace); + rgb->colorspace = frame->colorspace; + rgb->color_range = frame->color_range; + set_colorspace(rgb, (int)frame->colorspace); + const enum AVColorTransferCharacteristic trc = frame->color_trc; + const enum AVColorPrimaries pri = frame->color_primaries; + frame->color_trc = rgb->color_trc = AVCOL_TRC_UNSPECIFIED; + frame->color_primaries = rgb->color_primaries = AVCOL_PRI_UNSPECIFIED; + rgb->format = AV_PIX_FMT_RGB24; + rgb->width = frame->width; + rgb->height = frame->height; + if (av_frame_get_buffer(rgb, 0) < 0) { rc = TF_MEMORY; goto done; } + int e2 = sws_scale_frame(sws, rgb, frame); + frame->color_trc = trc; + frame->color_primaries = pri; + if (e2 < 0) goto done; + if (fn(ctx, next, at, rgb->data[0], rgb->width, rgb->height, rgb->linesize[0]) != 0) { rc = TF_STOPPED; goto done; } + ++next; + if (next == n) break; + } + ++at; + av_frame_unref(frame); + } + } + rc = next; +done: + sws_freeContext(sws); + av_frame_free(&rgb); + av_frame_free(&frame); + av_packet_free(&pkt); + avcodec_free_context(&dec); + return rc; +} diff --git a/zig/tests/glm_video.zig b/zig/tests/glm_video.zig new file mode 100644 index 0000000000..91f90b9fa2 --- /dev/null +++ b/zig/tests/glm_video.zig @@ -0,0 +1,69 @@ +//! GLM-5.3-Flash's video preparation on one clip, for tools/zig/glm_video_fixtures.py (libtfvideo, families/glm/video). +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const api = @import("engine_api"); +const video = tf.glm.video; +const vision = tf.glm.vision; + +const usage = + \\tf-glm-video MODEL_DIR CLIP OUT_DIR [MAX_TOKENS [MAX_FRAMES]] + \\Writes OUT_DIR/native_patches.f32 (every group's patches in order, [groups * gh * gw, 1176]) and + \\OUT_DIR/native.json: the stream's facts, the decoded indices, the grid [groups, gh, gw] and each group's text + \\("T.T seconds"). MAX_TOKENS defaults to 16384 tokens, MAX_FRAMES to 256 frames. +; + +fn readFile(a: std.mem.Allocator, path: []const u8) ![]u8 { + const f = try mtl.MappedFile.open(try std.fmt.allocPrintSentinel(a, "{s}", .{path}, 0)); + defer f.deinit(); + return a.dupe(u8, f.bytes[0..f.size]); +} + +const Decoder = struct { + v: api.tfvideo.Video, + + fn decode(ctx: *anyopaque, indices: []const u32, sink: *video.Sink) anyerror!void { + const d: *Decoder = @ptrCast(@alignCast(ctx)); + try d.v.frames(indices, sink); + } +}; + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len < 4) { + std.debug.print("{s}\n", .{usage}); + return error.Usage; + } + const io = init.io; + const max_tokens: u32 = if (args.len > 4) try std.fmt.parseInt(u32, args[4], 10) else 16384; + const max_frames: u32 = if (args.len > 5) try std.fmt.parseInt(u32, args[5], 10) else 256; + const cfg_text = try readFile(a, try std.fmt.allocPrint(a, "{s}/config.json", .{args[1]})); + const proc_text = readFile(a, try std.fmt.allocPrint(a, "{s}/processor_config.json", .{args[1]})) catch null; + const c = (try vision.Config.parse(a, cfg_text, proc_text)) orelse return error.NoVisionConfig; + const bytes = try readFile(a, args[2]); + var d: Decoder = .{ .v = try api.tfvideo.Video.open(bytes) }; + defer d.v.close(); + const i = d.v.info; + const src: video.Source = .{ .frames = @intCast(i.frames), .rate_num = @intCast(i.rate_num), .rate_den = @intCast(i.rate_den), .width = @intCast(i.width), .height = @intCast(i.height), .ctx = &d, .decode = Decoder.decode }; + const t0 = std.Io.Clock.awake.now(io); + const p = try video.plan(a, src, c, max_tokens, max_frames); + const prep = try video.prepare(a, src, c, max_tokens, max_frames); + const ms = t0.durationTo(std.Io.Clock.awake.now(io)).toMilliseconds(); + var dir = try std.Io.Dir.cwd().createDirPathOpen(io, args[3], .{}); + defer dir.close(io); + var all: std.ArrayList(u8) = .empty; + for (prep.groups) |g| try all.appendSlice(a, std.mem.sliceAsBytes(g.pixels)); + try dir.writeFile(io, .{ .sub_path = "native_patches.f32", .data = all.items }); + var texts: std.ArrayList([]const u8) = .empty; + for (prep.times) |t| { + var buf: [64]u8 = undefined; + try texts.append(a, try a.dupe(u8, try video.seconds(&buf, t))); + } + const g0 = prep.groups[0]; + const doc = .{ .frames = i.frames, .counted = i.counted != 0, .rate = .{ i.rate_num, i.rate_den }, .width = i.width, .height = i.height, .indices = p.indices, .grid = .{ prep.groups.len, g0.gh, g0.gw }, .tokens_per_group = g0.tokens, .texts = texts.items, .ms = ms }; + var out: std.Io.Writer.Allocating = .init(a); + try std.json.Stringify.value(doc, .{}, &out.writer); + try dir.writeFile(io, .{ .sub_path = "native.json", .data = out.written() }); + std.debug.print("{s}: {d} frames decoded, grid {d}x{d}x{d}, {d} tokens a group, {d} ms\n", .{ args[2], p.indices.len, prep.groups.len, g0.gh, g0.gw, g0.tokens, ms }); +} diff --git a/zig/tests/glm_vision.zig b/zig/tests/glm_vision.zig new file mode 100644 index 0000000000..7df41b512c --- /dev/null +++ b/zig/tests/glm_vision.zig @@ -0,0 +1,148 @@ +//! GLM-5.3-Flash's image tower against a reference dump (tools/zig/glm_vision_fixtures.py: mlx-vlm's tower). +const std = @import("std"); +const mtl = @import("metal"); +const tf = @import("tensorfold"); +const vision = tf.glm.vision; + +const usage = + \\tf-glm-vision MODEL_DIR REF_DIR [STAGE] [IMAGE] + \\REF_DIR holds meta.json (grid_thw), patches.f32 and the stage's reference: embed.f32 (STAGE embed, after the patch + \\projection), block0.f32 (STAGE block0, after the first block), blocks.f32 (STAGE blocks, after the last) or + \\features.f32 (STAGE features, the default: the tower's output rows). Writes native_STAGE.f32 there and prints the + \\largest and mean absolute differences, the reference's mean magnitude and the rows' cosine similarity. + \\IMAGE: prepare the patches from this file natively (decode, resize, pad, normalize), compare them with + \\patches.f32 value by value, and run the tower on them instead (meta.json's max_tokens is the budget). +; + +fn readFile(a: std.mem.Allocator, dir: []const u8, name: []const u8) ![]u8 { + const path = try std.fmt.allocPrintSentinel(a, "{s}/{s}", .{ dir, name }, 0); + const f = try mtl.MappedFile.open(path); + defer f.deinit(); + return a.dupe(u8, f.bytes[0..f.size]); +} + +fn bf16ToF32(x: u16) f32 { + return @bitCast(@as(u32, x) << 16); +} + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const args = try init.minimal.args.toSlice(a); + if (args.len < 3) { + std.debug.print("{s}\n", .{usage}); + return error.Usage; + } + const dir = args[1]; + const ref = args[2]; + const stage = if (args.len > 3) args[3] else "features"; + const stop: ?u32 = if (std.mem.eql(u8, stage, "embed")) 0 else if (std.mem.eql(u8, stage, "block0")) 1 else if (std.mem.eql(u8, stage, "blocks")) vision.depth else if (std.mem.eql(u8, stage, "features")) null else return error.Usage; + + const meta = try std.json.parseFromSliceLeaky(std.json.Value, a, try readFile(a, ref, "meta.json"), .{}); + const grid = meta.object.get("grid_thw").?.array.items; + const gh: u32 = @intCast(grid[1].integer); + const gw: u32 = @intCast(grid[2].integer); + const raw = try readFile(a, ref, "patches.f32"); + var pixels = try a.alloc(f32, raw.len / 4); + @memcpy(std.mem.sliceAsBytes(pixels), raw[0 .. pixels.len * 4]); + const P = gh * gw; + const T = P / vision.patches_per_token; + + const device = try mtl.Device.init(); + const pool = mtl.objc.Pool.push(); + defer pool.pop(); + const queue = try device.queue(); + const cfg_text = try readFile(a, dir, "config.json"); + const proc_text = readFile(a, dir, "processor_config.json") catch null; + const c = (try vision.Config.parse(a, cfg_text, proc_text)) orelse return error.NoVisionConfig; + if (args.len > 4) { + const img_path = try std.fmt.allocPrintSentinel(a, "{s}", .{args[4]}, 0); + const f = try mtl.MappedFile.open(img_path); + defer f.deinit(); + const max_tokens: u32 = @intCast(meta.object.get("max_tokens").?.integer); + const t0i = std.Io.Clock.awake.now(init.io).toNanoseconds(); + const prep = try tf.glm.image.prepare(a, f.bytes[0..f.size], c, max_tokens); + const t1i = std.Io.Clock.awake.now(init.io).toNanoseconds(); + std.debug.print("prepared {s}: grid {d}x{d}, {d} tokens in {d:.1} ms\n", .{ args[4], prep.gh, prep.gw, prep.tokens, @as(f64, @floatFromInt(t1i - t0i)) / 1e6 }); + if (prep.gh != gh or prep.gw != gw) { + std.debug.print("grid differs from the reference's {d}x{d}\n", .{ gh, gw }); + return error.GridMismatch; + } + var differ: usize = 0; + var max_p: f32 = 0; + for (prep.pixels, pixels) |x, y| { + if (x != y) differ += 1; + max_p = @max(max_p, @abs(x - y)); + } + std.debug.print("patches: {d} of {d} values differ from the reference (largest {d:.6})\n", .{ differ, pixels.len, max_p }); + pixels = prep.pixels; + } + var t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + const v = try vision.Vision.load(std.heap.page_allocator, device, dir, c, @max(T, 1024)); + defer v.deinit(); + var t1 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + std.debug.print("tower loaded in {d:.2} s ({d:.2} GiB with scratch for {d} patches, attention chunks of {d} rows)\n", .{ @as(f64, @floatFromInt(t1 - t0)) / 1e9, @as(f64, @floatFromInt(v.bytes)) / (1 << 30), v.max_patches, v.chunk }); + const out_buf = try device.buffer(@as(usize, T) * vision.out_hidden * 2, mtl.ResourceOptions.shared); + defer out_buf.deinit(); + + for (0..2) |run| { // the second run is the timed one (the first compiles pipelines' state on the GPU) + const cb = queue.commandBuffer(); + const enc = cb.compute(.serial); + t0 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + try v.encode(enc, pixels, gh, gw, .{ .buf = out_buf }, stop); + enc.end(); + cb.commit(); + cb.wait(); + t1 = std.Io.Clock.awake.now(init.io).toNanoseconds(); + if (cb.failure()) |msg| { + std.debug.print("command buffer failed: {s}\n", .{msg}); + return error.GpuFailed; + } + if (run == 1) std.debug.print("encode: {d} patches -> {d} tokens in {d:.1} ms\n", .{ P, T, @as(f64, @floatFromInt(t1 - t0)) / 1e6 }); + } + // the stage's rows: the residual stream for a stop, else the output rows + const width: usize = if (stop != null) vision.hidden else vision.out_hidden; + const rows: usize = if (stop != null) P else T; + const src: [*]const u16 = if (stop != null) @ptrCast(@alignCast(v.x.addr())) else @ptrCast(@alignCast(out_buf.contents())); + const got = try a.alloc(f32, rows * width); + for (got, 0..) |*g, i| g.* = bf16ToF32(src[i]); + const name = try std.fmt.allocPrint(a, "native_{s}.f32", .{stage}); + { + const path = try std.fmt.allocPrint(a, "{s}/{s}", .{ ref, name }); + const f = try std.Io.Dir.cwd().createFile(init.io, path, .{}); + defer f.close(init.io); + try f.writeStreamingAll(init.io, std.mem.sliceAsBytes(got)); + } + const want_raw = try readFile(a, ref, try std.fmt.allocPrint(a, "{s}.f32", .{stage})); + const want = try a.alloc(f32, want_raw.len / 4); + @memcpy(std.mem.sliceAsBytes(want), want_raw[0 .. want.len * 4]); + if (want.len != got.len) { + std.debug.print("size: native {d} values, reference {d}\n", .{ got.len, want.len }); + return error.SizeMismatch; + } + var max_d: f64 = 0; + var sum_d: f64 = 0; + var sum_m: f64 = 0; + var worst_cos: f64 = 1; + var mean_cos: f64 = 0; + for (0..rows) |r| { + var dot: f64 = 0; + var ng: f64 = 0; + var nw: f64 = 0; + for (0..width) |j| { + const g: f64 = got[r * width + j]; + const w: f64 = want[r * width + j]; + const d = @abs(g - w); + max_d = @max(max_d, d); + sum_d += d; + sum_m += @abs(w); + dot += g * w; + ng += g * g; + nw += w * w; + } + const cs = dot / (@sqrt(ng * nw) + 1e-30); + worst_cos = @min(worst_cos, cs); + mean_cos += cs; + } + const n: f64 = @floatFromInt(got.len); + std.debug.print("{s}: max |d| {d:.4}, mean |d| {d:.5}, mean |ref| {d:.4}, cosine mean {d:.6} worst {d:.6} ({d} rows of {d})\n", .{ stage, max_d, sum_d / n, sum_m / n, mean_cos / @as(f64, @floatFromInt(rows)), worst_cos, rows, width }); +}