[mm] sglang-mm: server vision pipeline core (fetch/driver/pipeline) + Qwen VL (#32364)
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Kan Wu
@@ -0,0 +1,70 @@
|
||||
name: PR Test (Rust extensions)
|
||||
|
||||
# In-crate checks for the Rust extensions bundled with the `sglang` wheel
|
||||
# (rust/sglang-mm). The pyo3 bindings themselves are exercised by the Python
|
||||
# unit suite (base-a-test-cpu), which builds the extensions from source; this
|
||||
# job covers the pure-Rust core, exactly as sglang-server links it, plus lints.
|
||||
#
|
||||
# Scoped to sglang-mm on purpose: sglang-server and sglang-grpc have no
|
||||
# in-crate test suite yet, so a `rust/**` trigger would spawn a job that does
|
||||
# not actually cover the changed crate. Widen the paths together with the job.
|
||||
|
||||
on:
|
||||
push:
|
||||
branches: [ main ]
|
||||
paths:
|
||||
- "rust/sglang-mm/**"
|
||||
- "rust/Cargo.toml"
|
||||
- ".github/workflows/pr-test-rust-exts.yml"
|
||||
pull_request:
|
||||
branches: [ main ]
|
||||
types: [opened, synchronize, reopened, labeled]
|
||||
paths:
|
||||
- "rust/sglang-mm/**"
|
||||
- "rust/Cargo.toml"
|
||||
- ".github/workflows/pr-test-rust-exts.yml"
|
||||
workflow_dispatch:
|
||||
|
||||
concurrency:
|
||||
group: rust-exts-${{ github.ref }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
sglang-mm-unit:
|
||||
if: |
|
||||
github.event_name != 'pull_request' ||
|
||||
(github.event.action != 'labeled' && contains(github.event.pull_request.labels.*.name, 'run-ci')) ||
|
||||
(github.event.action == 'labeled' && github.event.label.name == 'run-ci')
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 15
|
||||
defaults:
|
||||
run:
|
||||
working-directory: rust/sglang-mm
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: Rust cache
|
||||
uses: Swatinem/rust-cache@v2
|
||||
with:
|
||||
# The workspace root owns Cargo.lock and target/.
|
||||
workspaces: rust
|
||||
|
||||
# rustfmt and non-test clippy for rust/ are already gated by the Lint
|
||||
# workflow's pre-commit hooks; `--all-targets` is the part it misses, so
|
||||
# this lints the test code too rather than duplicating that gate.
|
||||
- name: cargo clippy (test targets)
|
||||
run: |
|
||||
rustup component add clippy
|
||||
cargo clippy --all-targets -- -D warnings
|
||||
|
||||
# Default features are exactly what sglang-server links: no pyo3, no
|
||||
# rayon. `rlib_is_single_threaded` only runs in this shape.
|
||||
- name: cargo test (rlib shape — no pyo3, no rayon)
|
||||
run: cargo test
|
||||
|
||||
# The wheel's shape (minus pyo3, whose extension-module crates cannot
|
||||
# link test binaries). Both shapes must pass: `common::par` has two
|
||||
# implementations and results are required to be identical.
|
||||
- name: cargo test (parallel shape)
|
||||
run: cargo test --features parallel
|
||||
@@ -130,10 +130,12 @@ repos:
|
||||
pass_filenames: false
|
||||
- id: clippy-rust-workspace
|
||||
name: clippy rust/ workspace (auto-fix)
|
||||
# Two steps: apply machine-applicable fixes, then fail on anything left
|
||||
# (combining --fix with -D warnings can abort the fix phase). protoc is
|
||||
# not required: sglang-grpc's build.rs falls back to a vendored binary.
|
||||
entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings'
|
||||
# Three steps: apply machine-applicable fixes, then fail on anything
|
||||
# left (combining --fix with -D warnings can abort the fix phase), then
|
||||
# cover sglang-mm's PyO3 bindings + rayon fan-out — both sit behind
|
||||
# non-default features, so `--workspace` alone never compiles them. protoc
|
||||
# is not required: sglang-grpc's build.rs falls back to a vendored binary.
|
||||
entry: bash -c 'cd rust && cargo clippy --workspace --fix --allow-dirty --allow-staged && cargo clippy --workspace -- -D warnings && cargo clippy -p sglang-mm --features python,parallel --lib -- -D warnings'
|
||||
language: system
|
||||
files: ^rust/.*\.rs$
|
||||
pass_filenames: false
|
||||
|
||||
@@ -6,6 +6,7 @@ crate whose Cargo.toml declares
|
||||
[package.metadata.sglang]
|
||||
python-module = "sglang.srt.<pkg>._core" # import path inside the wheel
|
||||
debug = false # optional RustExtension knob
|
||||
features = ["python"] # optional cargo features to enable
|
||||
|
||||
is built as a PyO3 extension module at that import path. Adding a new extension
|
||||
crate therefore needs no pyproject changes — declare the metadata in the crate.
|
||||
@@ -115,6 +116,9 @@ def _discovered_rust_extensions():
|
||||
path=package["manifest_path"],
|
||||
binding=Binding.PyO3,
|
||||
debug=sglang_meta.get("debug"),
|
||||
# Crates that gate their PyO3 bindings behind a non-default
|
||||
# feature (so the pure-Rust core stays pyo3-free) declare it here.
|
||||
features=sglang_meta.get("features"),
|
||||
)
|
||||
)
|
||||
if not extensions:
|
||||
|
||||
Generated
+3
@@ -3679,6 +3679,9 @@ dependencies = [
|
||||
"numpy",
|
||||
"pyo3",
|
||||
"rayon",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"ureq",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
|
||||
@@ -8,9 +8,12 @@ license.workspace = true
|
||||
# Consumed by python/setup.py: registers this crate as a PyO3 extension module
|
||||
# of the main sglang wheel at the given import path. debug = false keeps
|
||||
# editable installs on release builds (image preprocessing is perf-sensitive).
|
||||
# features: the wheel build needs the PyO3 bindings, which are NOT default (see
|
||||
# [features] below).
|
||||
[package.metadata.sglang]
|
||||
python-module = "sglang.srt.multimodal._core"
|
||||
debug = false
|
||||
features = ["python", "parallel"]
|
||||
|
||||
[lib]
|
||||
# Unique per-crate artifact name (the shared workspace target/ dir cannot hold
|
||||
@@ -18,14 +21,45 @@ debug = false
|
||||
# name comes from the `#[pymodule]` entry point and the setuptools-rust
|
||||
# `target` in python/pyproject.toml, which renames the built artifact.
|
||||
name = "sglang_mm_core"
|
||||
crate-type = ["cdylib"]
|
||||
# cdylib: the PyO3 module (`sglang.srt.multimodal._core`).
|
||||
# rlib: pure-Rust core linked by sglang-server's native MM path.
|
||||
crate-type = ["cdylib", "rlib"]
|
||||
|
||||
[features]
|
||||
# Both features are deliberately NOT default: cargo unifies a package's
|
||||
# features across everything selected in one invocation, so anything default-on
|
||||
# here would leak back into sglang-server's copy of this crate under
|
||||
# `cargo build --workspace` even though the server depends on it with
|
||||
# `default-features = false`. Opting in explicitly keeps the pure-Rust core
|
||||
# pure no matter how the workspace is built; the wheel asks for both via
|
||||
# `[package.metadata.sglang] features` above.
|
||||
default = []
|
||||
# Every PyO3 binding.
|
||||
python = ["dep:pyo3", "dep:numpy"]
|
||||
# Fan work out onto crate-owned rayon pools (see `common::par`). Off for the
|
||||
# rlib: sglang-server provides concurrency across requests and owns its core
|
||||
# budget, so preprocessing runs inline on the calling thread and rayon is not
|
||||
# linked at all. Orthogonal to `python` so both shapes stay testable.
|
||||
parallel = ["dep:rayon"]
|
||||
|
||||
[dependencies]
|
||||
pyo3 = { workspace = true }
|
||||
pyo3 = { workspace = true, optional = true }
|
||||
numpy = { version = "0.29", optional = true }
|
||||
|
||||
base64 = "0.22"
|
||||
blake3 = "1"
|
||||
half = "2.4"
|
||||
image = { version = "0.25", default-features = false, features = ["jpeg", "png"] }
|
||||
numpy = "0.29"
|
||||
rayon = "1.10"
|
||||
# The formats the Python (PIL) path commonly accepts; all pure-Rust decoders.
|
||||
image = { version = "0.25", default-features = false, features = [
|
||||
"jpeg",
|
||||
"png",
|
||||
"webp",
|
||||
"gif",
|
||||
"bmp",
|
||||
] }
|
||||
rayon = { version = "1.10", optional = true }
|
||||
serde = { workspace = true }
|
||||
serde_json = { workspace = true }
|
||||
# Native media fetch (stage 1 of the server MM pipeline); rustls so the
|
||||
# extension never links the system OpenSSL.
|
||||
ureq = { version = "2", default-features = false, features = ["tls", "gzip"] }
|
||||
|
||||
+141
-24
@@ -1,24 +1,117 @@
|
||||
# sglang-mm
|
||||
|
||||
Rust-accelerated multimodal preprocessing for SGLang. Fused image decode,
|
||||
resize, patchify, normalize, and content hash — all parallel and GIL-released.
|
||||
fetch, resize, patchify, normalize, and content hash — all parallel and
|
||||
GIL-released.
|
||||
|
||||
Compiled as `sglang.srt.multimodal._core` via setuptools-rust when installing sglang.
|
||||
Built two ways:
|
||||
|
||||
- **PyO3 extension** `sglang.srt.multimodal._core` (features `python,parallel`,
|
||||
requested by the wheel build) via setuptools-rust when installing sglang —
|
||||
used by Python processors and parity tests.
|
||||
- **Pure-Rust `rlib`** (default features, i.e. neither) linked by
|
||||
`sglang-server`'s MM worker path — that copy needs no pyo3, no libpython, and
|
||||
no rayon: it spawns no threads and runs inline on the calling thread, because
|
||||
the server supplies concurrency across requests and pins its own cores.
|
||||
`tests/rlib_is_single_threaded.rs` guards that from the outside.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
src/
|
||||
├── lib.rs # PyO3 module root (_core)
|
||||
├── registry.rs # ImageProcessorSpec trait + ProcessorRegistry
|
||||
├── lib.rs # module root; PyO3 module (_core) feature-gated
|
||||
├── pipeline.rs # the server-pipeline contract: MmFamilyProcessor
|
||||
│ # trait + the carriers (Tensor, TokenLayout, ...)
|
||||
├── driver.rs # model-independent request driver (fetch →
|
||||
│ # decode → process_item → layout → positions)
|
||||
├── registry.rs # ImageProcessorSpec registry (Python-facing)
|
||||
│ # + pipeline_from_spec (family factory)
|
||||
├── common/
|
||||
│ ├── mod.rs # thread pool, image decode, SHA256 hash, base64
|
||||
│ ├── resize.rs # PIL-exact Lanczos resize
|
||||
│ ├── mod.rs # thread pool, image decode, content hash, base64
|
||||
│ ├── fetch.rs # media source → bytes (data:/base64/file/http)
|
||||
│ ├── par.rs # the only fan-out seam (rayon, or inline)
|
||||
│ ├── resize.rs # PIL-exact Lanczos + Bicubic resize
|
||||
│ ├── token_layout.rs # TokenLayout mechanics (apply_layout + helpers)
|
||||
│ └── transforms.rs # reusable primitives: normalize, pad, extract_patches
|
||||
└── <model>/
|
||||
└── mod.rs # model-specific processor
|
||||
└── mod.rs # model-specific processor (inkling, qwen_vl, ...)
|
||||
```
|
||||
|
||||
## Server pipeline architecture
|
||||
|
||||
`sglang-server`'s MM workers process an image request entirely in Rust.
|
||||
`driver::process` runs the same fixed steps for every model family:
|
||||
|
||||
```
|
||||
MmInput { text?, input_ids?, images }
|
||||
1. per image: fetch_bytes (inline, sequential — see Design notes), then
|
||||
fanned out via common::par:
|
||||
content hash → decode_rgb → family.process_item()
|
||||
→ ProcessedItem { feature, aux, geometry }
|
||||
2. family.layout(input_ids, geometries) → TokenLayout
|
||||
apply_layout: expanded input_ids + per-item (start, end) offsets
|
||||
3. family.positions(len, offsets, geoms) → Rope1D | MRope
|
||||
4. Output { input_ids, items: [{feature, aux, hash}], offsets, positions }
|
||||
```
|
||||
|
||||
The driver owns these steps and their failure semantics — any `Err` at any
|
||||
step rejects the request as a 400 (there is no Python fallback path). A
|
||||
model family fills in only the `family.*` calls, by implementing
|
||||
`MmFamilyProcessor` (`pipeline.rs`): it describes its data, it never runs
|
||||
the request. With qwen as the example:
|
||||
|
||||
- **`process_item`** — one decoded image → `ProcessedItem`:
|
||||
- `feature`: the model's feature tensor. Qwen: `pixel_values`, from
|
||||
smart_resize → bicubic → normalize → patchify. The item identity is the
|
||||
driver's hash of the raw encoded source bytes, taken before decode — the
|
||||
same role as Python's `hash_feature`, but a different algorithm over
|
||||
different input, so never comparable across paths.
|
||||
- `aux`: named tensors for the model runner. Qwen: `image_grid_thw`;
|
||||
other families: `image_sizes`, `tgt_sizes`, ... (Python:
|
||||
`model_specific_data`).
|
||||
- `geometry`: whatever this family's `layout`/`positions` need later.
|
||||
Qwen: the `[t, h, w]` patch grid.
|
||||
- **`layout`** — how the prompt expands, described as a value. Example: the
|
||||
prompt `[A, <pad>, B]` with one 4-token image becomes
|
||||
|
||||
```
|
||||
[Text(0..1), Media { item: 0, Repeat(<pad> × 4) }, Text(2..3)]
|
||||
```
|
||||
|
||||
which the driver expands to `[A, <pad>, <pad>, <pad>, <pad>, B]` with
|
||||
offsets `[(1, 4)]`. Qwen builds this with the `layout_by_placeholder`
|
||||
helper; families that interleave tile markers or row separators
|
||||
(internvl/minicpm-style) use `Explicit` id sequences instead. Expansion,
|
||||
offsets, and position inputs all derive from this one value, so a family
|
||||
cannot get them out of sync.
|
||||
- **`positions`** — `Rope1D` (default: the scheduler needs nothing extra)
|
||||
or `MRope` (qwen's image-only fast path).
|
||||
- **`capabilities`** — which modalities the family accepts; the server
|
||||
rejects everything else per family.
|
||||
|
||||
Why not give each family the whole request, like Python's per-family
|
||||
`process_mm_data_async` override? In the server core, every request must
|
||||
resolve to exactly one accept/reject with its buffers parked in order —
|
||||
that invariant only holds structurally if the driver owns the flow.
|
||||
|
||||
Two things stay in Python permanently: HF config parsing (a family is
|
||||
configured by a spec JSON of already-resolved params, selected via
|
||||
`registry::pipeline_from_spec`) and the thin drain adapter mapping
|
||||
feature/aux tensors to model kwargs. The carriers grow by need, not
|
||||
speculation: `DecodedMedia` gains a variant per modality (video/audio),
|
||||
`Geometry` per family style (tile sets), `TensorData` per dtype.
|
||||
|
||||
Supported families: `qwen_vl` (Qwen2-VL / 2.5-VL / 3-VL / 3.5; images only).
|
||||
Adding one = a `MmFamilyProcessor` impl in `src/<model>/mod.rs` plus a
|
||||
`family` arm in `pipeline_from_spec`.
|
||||
|
||||
`common::fetch` matches the Python `get_image_bytes` semantics
|
||||
(`REQUEST_TIMEOUT` env, `HTTP(S)_PROXY` / `ALL_PROXY` / `NO_PROXY` including
|
||||
IPv4-CIDR and `host:port` entries) with two deliberate differences: every
|
||||
source form is capped at 64 MiB — plus 64 items / 256 MiB per request in the
|
||||
driver — and `file://` URLs actually work (the Python helper passes the
|
||||
un-stripped URL to `open()`).
|
||||
|
||||
## Python API
|
||||
|
||||
```python
|
||||
@@ -28,7 +121,8 @@ from sglang.srt.multimodal._core import common, inkling
|
||||
common.resize_rgb(arr, out_w, out_h)
|
||||
common.scaled_dims(w, h, rescale_frac, rescale_cap)
|
||||
common.image_decode_rgb(bytes) # -> (h, w, ndarray)
|
||||
common.data_hash(bytes) # -> u64 SHA256
|
||||
common.content_hash(bytes) # -> u64 (blake3, truncated)
|
||||
common.fetch_bytes(source) # -> bytes (data:/base64/file/http)
|
||||
common.base64_decode(str) # -> bytes
|
||||
|
||||
# Model-specific
|
||||
@@ -44,8 +138,8 @@ inkling.patchify_rgb(arr, patch_size)
|
||||
|
||||
```rust
|
||||
use crate::common;
|
||||
use crate::common::par;
|
||||
use crate::registry::ImageProcessorSpec;
|
||||
use rayon::prelude::*;
|
||||
|
||||
pub struct MyModelProcessor;
|
||||
|
||||
@@ -61,14 +155,14 @@ impl ImageProcessorSpec for MyModelProcessor {
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> Result<Vec<(usize, usize, Vec<u16>, u64)>, String> {
|
||||
common::pool().install(|| {
|
||||
datas.par_iter().map(|data| {
|
||||
let hash = common::sha256_u64(data);
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
// Use common::transforms::* or model-specific logic
|
||||
let patches = my_patchify(&rgb, h, w, patch_size);
|
||||
Ok((h, w, patches, hash))
|
||||
}).collect()
|
||||
// Always fan out through `par`, never rayon directly: that is what
|
||||
// keeps the rlib build rayon-free (see Design notes).
|
||||
par::try_map(datas, |data| {
|
||||
let hash = common::content_hash_u64(data);
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
// Use common::transforms::* or model-specific logic
|
||||
let patches = my_patchify(&rgb, h, w, patch_size);
|
||||
Ok((h, w, patches, hash))
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -93,9 +187,26 @@ impl ImageProcessorSpec for MyModelProcessor {
|
||||
|
||||
## Design notes
|
||||
|
||||
- Thread pool capped at `min(8, cores)`. Override: `SGL_MM_RS_THREADS`.
|
||||
- PNG decode is bit-exact vs PIL; JPEG may differ by ±1 LSB.
|
||||
- Lanczos resize is a bit-exact clone of PIL's fixed-point implementation.
|
||||
- All fan-out goes through `common::par`, so whether this crate owns threads is
|
||||
decided by the `parallel` feature alone. With it on: CPU pool capped at
|
||||
`min(8, cores)` (override `SGL_MM_RS_THREADS`). With it off: no rayon, no
|
||||
threads, everything inline. Output is bit-identical either way — the fan-outs
|
||||
are order-preserving maps and writes into disjoint slices, never reductions.
|
||||
Note that sizing a pool to 1 is *not* the same as off: `install` blocks the
|
||||
caller and would serialize every concurrent request in the process.
|
||||
- Media fetch is blocking I/O and deliberately never enters the CPU pool; it
|
||||
runs inline and sequentially in `driver::process`. Contract: callers on a
|
||||
fixed worker pool (sglang-server) must resolve I/O-backed string sources —
|
||||
URLs *and* file paths (a network mount can hang far longer than any HTTP
|
||||
timeout) — on their own I/O layer and pass bytes, so workers never block on
|
||||
I/O. `data:`/base64 sources are pure CPU and stay on the worker.
|
||||
- PNG decode is bit-exact vs PIL; JPEG may differ by ±1 LSB. WebP/GIF/BMP also
|
||||
decode (GIF: first frame); their parity is not bit-audited. Samples deeper
|
||||
than 8 bits are rejected rather than rescaled (PIL clips instead).
|
||||
- Lanczos and Bicubic resize are bit-exact clones of PIL's fixed-point
|
||||
implementations.
|
||||
- `common::content_hash_u64` is blake3, *not* Python's SHA-256
|
||||
`mm_utils.data_hash`. Hashes are consistent within one path only.
|
||||
|
||||
## Build
|
||||
|
||||
@@ -104,17 +215,23 @@ Automatically built when installing sglang:
|
||||
pip install -e "python"
|
||||
```
|
||||
|
||||
Or standalone for development:
|
||||
Or standalone for development (the PyO3 bindings are behind a non-default
|
||||
feature — see `[features]` in `Cargo.toml` for why):
|
||||
```bash
|
||||
cd rust/sglang-mm
|
||||
pip install maturin
|
||||
maturin develop --release
|
||||
maturin develop --release --features python
|
||||
```
|
||||
|
||||
## Test
|
||||
|
||||
```bash
|
||||
python bench/generate_golden.py # regenerate fixtures
|
||||
pytest bench/test_golden.py # regression tests
|
||||
cd rust/sglang-mm
|
||||
cargo test --no-default-features # pure-Rust unit tests (CI: pr-test-rust-exts)
|
||||
python tests/generate_golden.py # regenerate fixtures
|
||||
pytest tests/test_golden.py # regression tests
|
||||
python bench/bench_parity.py # parity + benchmark
|
||||
```
|
||||
|
||||
Scheduler-boundary parity tests against the real HF processors live in
|
||||
`test/registered/unit/multimodal/rust/`.
|
||||
|
||||
@@ -5,7 +5,7 @@ import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
import sglang.srt.multimodal._core.inkling
|
||||
from sglang.srt.multimodal._core import inkling as _rs_inkling
|
||||
from sglang.srt.multimodal.inkling.image_processing import (
|
||||
IMAGE_MEAN,
|
||||
IMAGE_STD,
|
||||
@@ -30,12 +30,12 @@ def rs_patchify(arr: np.ndarray) -> torch.Tensor:
|
||||
h, w, _ = arr.shape
|
||||
nph = (h + PS - 1) // PS
|
||||
npw = w // PS + 1
|
||||
bits = sglang.srt.multimodal._core.inkling.patchify_rgb(arr, PS)
|
||||
bits = _rs_inkling.patchify_rgb(arr, PS)
|
||||
return torch.from_numpy(bits).view(torch.bfloat16).reshape(nph * npw, PS, PS, 3)
|
||||
|
||||
|
||||
def rs_decode_patchify(data: bytes) -> torch.Tensor:
|
||||
h, w, bits = sglang.srt.multimodal._core.inkling.decode_patchify(data, PS)
|
||||
h, w, bits = _rs_inkling.decode_patchify(data, PS)
|
||||
nph = (h + PS - 1) // PS
|
||||
npw = w // PS + 1
|
||||
return torch.from_numpy(bits).view(torch.bfloat16).reshape(nph * npw, PS, PS, 3)
|
||||
@@ -114,7 +114,7 @@ def bench():
|
||||
rescale_image_max_upscaled_long_edge=None,
|
||||
)
|
||||
rs_decode_patchify(jpeg)
|
||||
sglang.srt.multimodal._core.inkling.decode_patchify_batch([jpeg] * 5, PS)
|
||||
_rs_inkling.decode_patchify_batch([jpeg] * 5, PS)
|
||||
|
||||
def run(label, fn, iters=n, images_per_call=1):
|
||||
t0, c0 = time.perf_counter(), time.process_time()
|
||||
@@ -137,9 +137,7 @@ def bench():
|
||||
w_rs, c_rs = run("rust decode_patchify", lambda: rs_decode_patchify(jpeg))
|
||||
w_rb, c_rb = run(
|
||||
"rust decode_patchify_batch (5 imgs/call)",
|
||||
lambda: sglang.srt.multimodal._core.inkling.decode_patchify_batch(
|
||||
[jpeg] * 5, PS
|
||||
),
|
||||
lambda: _rs_inkling.decode_patchify_batch([jpeg] * 5, PS),
|
||||
iters=max(n // 5, 5),
|
||||
images_per_call=5,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,331 @@
|
||||
//! Stage 1 of the server MM pipeline: resolve one media source to raw bytes.
|
||||
//!
|
||||
//! Mirrors the Python `get_image_bytes` source handling (and its precedence):
|
||||
//! raw bytes, `http(s)://` (bounded download, `REQUEST_TIMEOUT` and the proxy
|
||||
//! env vars like Python), `file://` / absolute path, `data:` URL, else bare
|
||||
//! base64.
|
||||
|
||||
use std::io::Read;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use base64::Engine;
|
||||
|
||||
/// Cap on any single resolved payload — HTTP, file, or base64 — so no source
|
||||
/// form can exhaust memory (the Python path has no such cap; oversized
|
||||
/// payloads reject the request here).
|
||||
pub const MAX_FETCH_BYTES: u64 = 64 << 20;
|
||||
|
||||
/// Resolve one string-typed image source into raw encoded-image bytes.
|
||||
/// An `Err` rejects the request, matching the Python per-request
|
||||
/// exception → 400.
|
||||
pub fn fetch_bytes(src: &str) -> Result<Vec<u8>, String> {
|
||||
if src.starts_with("http://") || src.starts_with("https://") {
|
||||
return http_get(src);
|
||||
}
|
||||
if let Some(path) = src.strip_prefix("file://") {
|
||||
return read_file(path);
|
||||
}
|
||||
if src.starts_with('/') {
|
||||
return read_file(src);
|
||||
}
|
||||
if let Some(rest) = src.strip_prefix("data:") {
|
||||
let encoded = rest
|
||||
.split_once(',')
|
||||
.ok_or_else(|| "media fetch: malformed data: URL".to_string())?
|
||||
.1;
|
||||
return b64(encoded);
|
||||
}
|
||||
// Python treats any other string as bare base64.
|
||||
b64(src)
|
||||
}
|
||||
|
||||
/// Bounded read: never trusts metadata, so huge and non-regular files
|
||||
/// (`/dev/zero`) hit the cap instead of exhausting memory.
|
||||
fn read_file(path: &str) -> Result<Vec<u8>, String> {
|
||||
let file = std::fs::File::open(path).map_err(|e| format!("media fetch: {path}: {e}"))?;
|
||||
read_capped(file, path)
|
||||
}
|
||||
|
||||
fn read_capped(reader: impl Read, what: &str) -> Result<Vec<u8>, String> {
|
||||
let mut buf = Vec::new();
|
||||
reader
|
||||
.take(MAX_FETCH_BYTES + 1)
|
||||
.read_to_end(&mut buf)
|
||||
.map_err(|e| format!("media fetch: read {what}: {e}"))?;
|
||||
if buf.len() as u64 > MAX_FETCH_BYTES {
|
||||
return Err(format!(
|
||||
"media fetch: {what}: exceeds {MAX_FETCH_BYTES} bytes"
|
||||
));
|
||||
}
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
fn b64(encoded: &str) -> Result<Vec<u8>, String> {
|
||||
// Slightly laxer than Python's `pybase64.b64decode(validate=True)`:
|
||||
// surrounding whitespace (e.g. a trailing newline) is trimmed here.
|
||||
let encoded = encoded.trim();
|
||||
// Reject by encoded length before allocating the decode buffer.
|
||||
if encoded.len() as u64 / 4 * 3 > MAX_FETCH_BYTES {
|
||||
return Err(format!(
|
||||
"media fetch: base64 payload exceeds {MAX_FETCH_BYTES} bytes"
|
||||
));
|
||||
}
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(encoded.as_bytes())
|
||||
.map_err(|e| format!("media fetch: base64 decode: {e}"))
|
||||
}
|
||||
|
||||
/// Shared pooled agent honoring `HTTP_PROXY`/`HTTPS_PROXY`/`ALL_PROXY`, as the
|
||||
/// Python `requests` session does.
|
||||
fn http_agent() -> &'static ureq::Agent {
|
||||
static AGENT: OnceLock<ureq::Agent> = OnceLock::new();
|
||||
AGENT.get_or_init(|| ureq::AgentBuilder::new().try_proxy_from_env(true).build())
|
||||
}
|
||||
|
||||
/// Companion agent that ignores the proxy env vars, for hosts matched by
|
||||
/// `NO_PROXY`. ureq has no `NO_PROXY` support of its own, and silently sending
|
||||
/// an internal image host through a corporate proxy breaks deployments that
|
||||
/// work on the Python path, so the match is applied here.
|
||||
fn direct_agent() -> &'static ureq::Agent {
|
||||
static AGENT: OnceLock<ureq::Agent> = OnceLock::new();
|
||||
AGENT.get_or_init(|| ureq::AgentBuilder::new().build())
|
||||
}
|
||||
|
||||
/// The host component of an `http(s)://` URL (lowercased, without userinfo)
|
||||
/// plus its explicit port, if any — what `NO_PROXY` entries are matched
|
||||
/// against.
|
||||
fn host_port_of(url: &str) -> Option<(String, Option<u16>)> {
|
||||
let rest = url
|
||||
.strip_prefix("http://")
|
||||
.or_else(|| url.strip_prefix("https://"))?;
|
||||
let authority = rest.split(['/', '?', '#']).next()?;
|
||||
let host = authority.rsplit_once('@').map_or(authority, |(_, h)| h);
|
||||
// Bracketed IPv6 literal, else split off a trailing `:port`.
|
||||
let (host, port) = match host.strip_prefix('[') {
|
||||
Some(v6) => match v6.split_once(']') {
|
||||
Some((h, p)) => (h, p.strip_prefix(':')),
|
||||
None => (v6, None),
|
||||
},
|
||||
None => match host.split_once(':') {
|
||||
Some((h, p)) => (h, Some(p)),
|
||||
None => (host, None),
|
||||
},
|
||||
};
|
||||
let port = port.and_then(|p| p.parse().ok());
|
||||
(!host.is_empty()).then(|| (host.to_ascii_lowercase(), port))
|
||||
}
|
||||
|
||||
fn bypasses_proxy(host: &str, port: Option<u16>) -> bool {
|
||||
["no_proxy", "NO_PROXY"]
|
||||
.iter()
|
||||
.find_map(|key| std::env::var(key).ok())
|
||||
.is_some_and(|list| no_proxy_matches(&list, host, port))
|
||||
}
|
||||
|
||||
/// `NO_PROXY` semantics as `requests` implements them: comma-separated
|
||||
/// entries; `*` matches everything; an IPv4 CIDR entry matches an IPv4 host in
|
||||
/// that network (requests supports IPv4 networks only); a `host:port` entry
|
||||
/// matches only that explicit port; otherwise an entry matches a host that
|
||||
/// equals it or is a subdomain of it (leading dots ignored). Kept pure so it
|
||||
/// is testable without mutating process-global env.
|
||||
fn no_proxy_matches(no_proxy: &str, host: &str, port: Option<u16>) -> bool {
|
||||
no_proxy.split(',').any(|entry| {
|
||||
let entry = entry.trim().trim_start_matches('.').to_ascii_lowercase();
|
||||
if entry.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if entry == "*" {
|
||||
return true;
|
||||
}
|
||||
if let Some((net, bits)) = parse_ipv4_cidr(&entry) {
|
||||
return host
|
||||
.parse::<std::net::Ipv4Addr>()
|
||||
.is_ok_and(|ip| in_ipv4_network(ip, net, bits));
|
||||
}
|
||||
let (entry_host, entry_port) = match entry.rsplit_once(':') {
|
||||
Some((h, p)) if p.bytes().all(|b| b.is_ascii_digit()) => (h, p.parse::<u16>().ok()),
|
||||
_ => (entry.as_str(), None),
|
||||
};
|
||||
if entry_port.is_some() && entry_port != port {
|
||||
return false;
|
||||
}
|
||||
host == entry_host || host.ends_with(&format!(".{entry_host}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn parse_ipv4_cidr(entry: &str) -> Option<(std::net::Ipv4Addr, u32)> {
|
||||
let (net, bits) = entry.split_once('/')?;
|
||||
Some((net.parse().ok()?, bits.parse().ok().filter(|b| *b <= 32)?))
|
||||
}
|
||||
|
||||
fn in_ipv4_network(ip: std::net::Ipv4Addr, net: std::net::Ipv4Addr, bits: u32) -> bool {
|
||||
let mask = u32::MAX.checked_shl(32 - bits).unwrap_or(0);
|
||||
u32::from(ip) & mask == u32::from(net) & mask
|
||||
}
|
||||
|
||||
fn http_get(url: &str) -> Result<Vec<u8>, String> {
|
||||
// Python: `int(os.getenv("REQUEST_TIMEOUT", "3"))` seconds per image GET.
|
||||
let timeout = std::env::var("REQUEST_TIMEOUT")
|
||||
.ok()
|
||||
.and_then(|s| s.parse::<u64>().ok())
|
||||
.unwrap_or(3);
|
||||
let agent = match host_port_of(url) {
|
||||
Some((host, port)) if bypasses_proxy(&host, port) => direct_agent(),
|
||||
_ => http_agent(),
|
||||
};
|
||||
let resp = agent
|
||||
.get(url)
|
||||
.timeout(std::time::Duration::from_secs(timeout))
|
||||
.call()
|
||||
.map_err(|e| format!("media fetch: GET {url}: {e}"))?;
|
||||
read_capped(resp.into_reader(), url)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn data_url_and_bare_base64_decode() {
|
||||
let b64 = base64::engine::general_purpose::STANDARD.encode(b"hello");
|
||||
assert_eq!(
|
||||
fetch_bytes(&format!("data:image/png;base64,{b64}")).unwrap(),
|
||||
b"hello"
|
||||
);
|
||||
assert_eq!(fetch_bytes(&b64).unwrap(), b"hello");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn bad_base64_fails() {
|
||||
assert!(fetch_bytes("!!not-base64!!").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_file_fails() {
|
||||
assert!(fetch_bytes("file:///definitely/not/here.jpg").is_err());
|
||||
assert!(fetch_bytes("/definitely/not/here.jpg").is_err());
|
||||
}
|
||||
|
||||
/// A non-regular file must hit the byte cap, not exhaust memory.
|
||||
#[test]
|
||||
fn unbounded_file_capped() {
|
||||
let err = fetch_bytes("/dev/zero").err().unwrap();
|
||||
assert!(err.contains("exceeds"), "{err}");
|
||||
}
|
||||
|
||||
/// Oversized base64 is rejected from its encoded length, before decoding.
|
||||
#[test]
|
||||
fn oversized_base64_rejected() {
|
||||
let encoded = "A".repeat((MAX_FETCH_BYTES / 3 * 4 + 8) as usize);
|
||||
let err = fetch_bytes(&encoded).err().unwrap();
|
||||
assert!(err.contains("exceeds"), "{err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn host_parsing_strips_userinfo_and_path() {
|
||||
assert_eq!(
|
||||
host_port_of("http://Example.COM/a/b.png").unwrap(),
|
||||
("example.com".into(), None)
|
||||
);
|
||||
assert_eq!(
|
||||
host_port_of("https://u:p@images.internal:8443/x").unwrap(),
|
||||
("images.internal".into(), Some(8443))
|
||||
);
|
||||
assert_eq!(
|
||||
host_port_of("http://[::1]:8080/x.png").unwrap(),
|
||||
("::1".into(), Some(8080))
|
||||
);
|
||||
assert_eq!(
|
||||
host_port_of("http://host?q=1").unwrap(),
|
||||
("host".into(), None)
|
||||
);
|
||||
assert!(host_port_of("data:image/png;base64,AAA").is_none());
|
||||
}
|
||||
|
||||
/// `NO_PROXY` must bypass the proxy for exact hosts and subdomains but not
|
||||
/// for lookalike suffixes — sending an internal host to a corporate proxy
|
||||
/// is a silent failure that works fine on the Python path.
|
||||
#[test]
|
||||
fn no_proxy_matches_host_and_subdomains_only() {
|
||||
let list = " .internal ,localhost";
|
||||
assert!(no_proxy_matches(list, "images.internal", None));
|
||||
assert!(no_proxy_matches(list, "internal", None));
|
||||
assert!(no_proxy_matches(list, "localhost", None));
|
||||
assert!(!no_proxy_matches(list, "notinternal", None));
|
||||
assert!(!no_proxy_matches(list, "example.com", None));
|
||||
assert!(no_proxy_matches("*", "anything.example.com", None));
|
||||
// An empty or all-empty list must not bypass everything.
|
||||
assert!(!no_proxy_matches("", "example.com", None));
|
||||
assert!(!no_proxy_matches(" , ", "example.com", None));
|
||||
}
|
||||
|
||||
/// IPv4 CIDR entries match IP-literal hosts, as `requests` does.
|
||||
#[test]
|
||||
fn no_proxy_matches_ipv4_cidr() {
|
||||
assert!(no_proxy_matches("10.0.0.0/8", "10.1.2.3", None));
|
||||
assert!(!no_proxy_matches("10.0.0.0/8", "11.1.2.3", None));
|
||||
assert!(no_proxy_matches("192.168.1.0/24", "192.168.1.77", Some(80)));
|
||||
assert!(!no_proxy_matches("192.168.1.0/24", "192.168.2.1", None));
|
||||
assert!(no_proxy_matches("0.0.0.0/0", "8.8.8.8", None));
|
||||
// CIDR entries never match hostnames.
|
||||
assert!(!no_proxy_matches("10.0.0.0/8", "example.com", None));
|
||||
}
|
||||
|
||||
/// `host:port` entries match only that explicit port.
|
||||
#[test]
|
||||
fn no_proxy_matches_host_with_port() {
|
||||
assert!(no_proxy_matches("internal:8443", "internal", Some(8443)));
|
||||
assert!(!no_proxy_matches("internal:8443", "internal", Some(80)));
|
||||
assert!(!no_proxy_matches("internal:8443", "internal", None));
|
||||
assert!(no_proxy_matches(
|
||||
"internal:8443",
|
||||
"img.internal",
|
||||
Some(8443)
|
||||
));
|
||||
// A port-free entry matches any port.
|
||||
assert!(no_proxy_matches("internal", "internal", Some(8443)));
|
||||
}
|
||||
|
||||
/// End-to-end HTTP download against a local one-shot server, and the
|
||||
/// capped rejection of an oversized response.
|
||||
#[test]
|
||||
fn http_download_and_cap() {
|
||||
let serve = |body: Vec<u8>, content_length: u64| {
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let handle = std::thread::spawn(move || {
|
||||
use std::io::{BufRead, Write};
|
||||
let (stream, _) = listener.accept().unwrap();
|
||||
let mut reader = std::io::BufReader::new(stream);
|
||||
let mut line = String::new();
|
||||
while reader.read_line(&mut line).unwrap() > 2 {
|
||||
line.clear(); // headers until the blank line
|
||||
}
|
||||
let mut stream = reader.into_inner();
|
||||
write!(
|
||||
stream,
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {content_length}\r\n\r\n"
|
||||
)
|
||||
.unwrap();
|
||||
stream.write_all(&body).unwrap();
|
||||
});
|
||||
(addr, handle)
|
||||
};
|
||||
|
||||
let (addr, handle) = serve(b"tiny image".to_vec(), 10);
|
||||
assert_eq!(
|
||||
fetch_bytes(&format!("http://{addr}/img.png")).unwrap(),
|
||||
b"tiny image"
|
||||
);
|
||||
handle.join().unwrap();
|
||||
|
||||
// A response over the cap is rejected without buffering it all.
|
||||
let over = MAX_FETCH_BYTES + 2;
|
||||
let (addr, handle) = serve(vec![0u8; over as usize], over);
|
||||
let err = fetch_bytes(&format!("http://{addr}/big.png"))
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(err.contains("exceeds"), "{err}");
|
||||
drop(handle); // server thread may die on the closed socket; don't join
|
||||
}
|
||||
}
|
||||
@@ -1,12 +1,19 @@
|
||||
pub mod fetch;
|
||||
pub mod par;
|
||||
pub mod resize;
|
||||
pub mod token_layout;
|
||||
pub mod transforms;
|
||||
|
||||
#[cfg(feature = "parallel")]
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
/// CPU pool: decode, resize, patchify, hash. Sized to cores (capped at 8)
|
||||
/// because the work is compute-bound — never run blocking I/O on it, or one
|
||||
/// request's remote fetches stall every other request's preprocessing.
|
||||
///
|
||||
/// Only exists under the `parallel` feature; reach it through [`par`], never
|
||||
/// directly, so the rayon-less build stays compiling.
|
||||
#[cfg(feature = "parallel")]
|
||||
pub fn pool() -> &'static rayon::ThreadPool {
|
||||
static POOL: OnceLock<rayon::ThreadPool> = OnceLock::new();
|
||||
POOL.get_or_init(|| {
|
||||
@@ -23,13 +30,27 @@ pub fn pool() -> &'static rayon::ThreadPool {
|
||||
})
|
||||
}
|
||||
|
||||
pub fn sha256_u64(data: &[u8]) -> u64 {
|
||||
/// Content hash for cache/dedup identity: blake3 truncated to its first 8
|
||||
/// bytes, big-endian. Deliberately *not* Python's `mm_utils.data_hash` (which
|
||||
/// is SHA-256 truncated the same way) — hashes are consistent within a path,
|
||||
/// never comparable across the Rust and Python paths.
|
||||
pub fn content_hash_u64(data: &[u8]) -> u64 {
|
||||
let digest = blake3::hash(data);
|
||||
u64::from_be_bytes(digest.as_bytes()[..8].try_into().unwrap())
|
||||
}
|
||||
|
||||
pub fn decode_rgb(data: &[u8]) -> Result<(Vec<u8>, usize, usize), String> {
|
||||
use image::ColorType;
|
||||
|
||||
let img = image::load_from_memory(data).map_err(|e| format!("image decode: {e}"))?;
|
||||
// >8-bit samples: PIL clips to 255 where `to_rgb8` would rescale. Refuse
|
||||
// rather than silently diverge from the Python (PIL) pipeline.
|
||||
if !matches!(
|
||||
img.color(),
|
||||
ColorType::L8 | ColorType::La8 | ColorType::Rgb8 | ColorType::Rgba8
|
||||
) {
|
||||
return Err(format!("image decode: unsupported color {:?}", img.color()));
|
||||
}
|
||||
let rgb = img.to_rgb8();
|
||||
let (w, h) = rgb.dimensions();
|
||||
Ok((rgb.into_raw(), h as usize, w as usize))
|
||||
@@ -48,87 +69,144 @@ pub fn decode_rescale(
|
||||
Ok((resize::resize_lanczos_rgb(&rgb, h, w, th, tw), th, tw))
|
||||
}
|
||||
|
||||
// --- Python-exposed functions ---
|
||||
// --- Python bindings (feature-gated: absent from the pure-Rust rlib) ---
|
||||
|
||||
#[pyfunction]
|
||||
pub fn resize_rgb<'py>(
|
||||
py: Python<'py>,
|
||||
arr: PyReadonlyArray3<'py, u8>,
|
||||
out_w: usize,
|
||||
out_h: usize,
|
||||
) -> PyResult<Bound<'py, PyArray1<u8>>> {
|
||||
if out_w == 0 || out_h == 0 {
|
||||
return Err(PyValueError::new_err("output size must be positive"));
|
||||
#[cfg(feature = "python")]
|
||||
mod python {
|
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use pyo3::types::PyBytes;
|
||||
|
||||
use super::{decode_rgb, resize};
|
||||
|
||||
#[pyfunction]
|
||||
pub fn resize_rgb<'py>(
|
||||
py: Python<'py>,
|
||||
arr: PyReadonlyArray3<'py, u8>,
|
||||
out_w: usize,
|
||||
out_h: usize,
|
||||
) -> PyResult<Bound<'py, PyArray1<u8>>> {
|
||||
if out_w == 0 || out_h == 0 {
|
||||
return Err(PyValueError::new_err("output size must be positive"));
|
||||
}
|
||||
let shape = arr.shape();
|
||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||
if c != 3 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"expected HWC RGB array with 3 channels, got {c}"
|
||||
)));
|
||||
}
|
||||
let data = arr
|
||||
.as_slice()
|
||||
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
||||
.to_vec();
|
||||
let out = py.detach(move || resize::resize_lanczos_rgb(&data, h, w, out_h, out_w));
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
let shape = arr.shape();
|
||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||
if c != 3 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"expected HWC RGB array with 3 channels, got {c}"
|
||||
)));
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (w, h, rescale_frac=None, rescale_cap=None))]
|
||||
pub fn scaled_dims(
|
||||
w: usize,
|
||||
h: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> (usize, usize) {
|
||||
resize::scaled_dims(w, h, rescale_frac, rescale_cap)
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn image_decode_rgb<'py>(
|
||||
py: Python<'py>,
|
||||
data: Vec<u8>,
|
||||
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u8>>)> {
|
||||
let (rgb, h, w) = py
|
||||
.detach(move || decode_rgb(&data))
|
||||
.map_err(PyValueError::new_err)?;
|
||||
Ok((h, w, rgb.into_pyarray(py)))
|
||||
}
|
||||
|
||||
/// Named `content_hash`, not `data_hash`, so it is not mistaken for
|
||||
/// `sglang.srt.managers.mm_utils.data_hash` (SHA-256); this is blake3.
|
||||
#[pyfunction]
|
||||
pub fn content_hash(py: Python<'_>, data: Vec<u8>) -> u64 {
|
||||
py.detach(move || super::content_hash_u64(&data))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn fetch_bytes<'py>(py: Python<'py>, source: String) -> PyResult<Bound<'py, PyBytes>> {
|
||||
let data = py
|
||||
.detach(move || super::fetch::fetch_bytes(&source))
|
||||
.map_err(|error| PyValueError::new_err(error.to_string()))?;
|
||||
Ok(PyBytes::new(py, &data))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn base64_decode<'py>(
|
||||
py: Python<'py>,
|
||||
encoded: &str,
|
||||
) -> PyResult<Bound<'py, pyo3::types::PyBytes>> {
|
||||
use base64::Engine;
|
||||
let decoded = py
|
||||
.detach(|| {
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(encoded)
|
||||
.map_err(|e| format!("base64 decode error: {e}"))
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
Ok(pyo3::types::PyBytes::new(py, &decoded))
|
||||
}
|
||||
|
||||
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let m = PyModule::new(parent.py(), "common")?;
|
||||
m.add_function(wrap_pyfunction!(resize_rgb, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(scaled_dims, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(image_decode_rgb, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(content_hash, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(fetch_bytes, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(base64_decode, &m)?)?;
|
||||
parent.add_submodule(&m)?;
|
||||
Ok(())
|
||||
}
|
||||
let data = arr
|
||||
.as_slice()
|
||||
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
||||
.to_vec();
|
||||
let out =
|
||||
py.detach(move || pool().install(|| resize::resize_lanczos_rgb(&data, h, w, out_h, out_w)));
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (w, h, rescale_frac=None, rescale_cap=None))]
|
||||
pub fn scaled_dims(
|
||||
w: usize,
|
||||
h: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> (usize, usize) {
|
||||
resize::scaled_dims(w, h, rescale_frac, rescale_cap)
|
||||
}
|
||||
#[cfg(feature = "python")]
|
||||
pub use python::register;
|
||||
|
||||
#[pyfunction]
|
||||
pub fn image_decode_rgb<'py>(
|
||||
py: Python<'py>,
|
||||
data: Vec<u8>,
|
||||
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u8>>)> {
|
||||
let (rgb, h, w) = py
|
||||
.detach(move || decode_rgb(&data))
|
||||
.map_err(PyValueError::new_err)?;
|
||||
Ok((h, w, rgb.into_pyarray(py)))
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::decode_rgb;
|
||||
use image::ImageFormat;
|
||||
|
||||
#[pyfunction]
|
||||
pub fn data_hash(py: Python<'_>, data: Vec<u8>) -> u64 {
|
||||
py.detach(move || {
|
||||
let digest = blake3::hash(&data);
|
||||
u64::from_be_bytes(digest.as_bytes()[..8].try_into().unwrap())
|
||||
})
|
||||
}
|
||||
fn encode(img: &image::DynamicImage, fmt: ImageFormat) -> Vec<u8> {
|
||||
let mut buf = std::io::Cursor::new(Vec::new());
|
||||
img.write_to(&mut buf, fmt).unwrap();
|
||||
buf.into_inner()
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
pub fn base64_decode<'py>(
|
||||
py: Python<'py>,
|
||||
encoded: &str,
|
||||
) -> PyResult<Bound<'py, pyo3::types::PyBytes>> {
|
||||
use base64::Engine;
|
||||
let decoded = py
|
||||
.detach(|| {
|
||||
base64::engine::general_purpose::STANDARD
|
||||
.decode(encoded)
|
||||
.map_err(|e| format!("base64 decode error: {e}"))
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
Ok(pyo3::types::PyBytes::new(py, &decoded))
|
||||
}
|
||||
/// Formats the Python (PIL) path accepts must decode, not reject.
|
||||
#[test]
|
||||
fn decodes_webp_gif_bmp() {
|
||||
let img = image::DynamicImage::ImageRgb8(image::RgbImage::from_fn(6, 4, |x, y| {
|
||||
image::Rgb([x as u8 * 40, y as u8 * 60, 7])
|
||||
}));
|
||||
for fmt in [ImageFormat::WebP, ImageFormat::Gif, ImageFormat::Bmp] {
|
||||
let (rgb, h, w) = decode_rgb(&encode(&img, fmt)).unwrap();
|
||||
assert_eq!((h, w), (4, 6), "{fmt:?}");
|
||||
assert_eq!(rgb.len(), 4 * 6 * 3, "{fmt:?}");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let m = PyModule::new(parent.py(), "common")?;
|
||||
m.add_function(wrap_pyfunction!(resize_rgb, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(scaled_dims, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(image_decode_rgb, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(data_hash, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(base64_decode, &m)?)?;
|
||||
parent.add_submodule(&m)?;
|
||||
Ok(())
|
||||
/// Samples deeper than 8 bits stay rejected (PIL clips; we refuse).
|
||||
#[test]
|
||||
fn deep_png_rejected() {
|
||||
let img = image::DynamicImage::ImageRgb16(image::ImageBuffer::from_pixel(
|
||||
2,
|
||||
2,
|
||||
image::Rgb([65535u16, 0, 0]),
|
||||
));
|
||||
let err = decode_rgb(&encode(&img, ImageFormat::Png)).err().unwrap();
|
||||
assert!(err.contains("unsupported color"), "{err}");
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
//! The crate's only parallelism seam.
|
||||
//!
|
||||
//! Every fan-out in the crate goes through the functions below, so whether
|
||||
//! this crate owns worker threads at all is decided in exactly one place: the
|
||||
//! `parallel` cargo feature.
|
||||
//!
|
||||
//! * **feature on** (the PyO3 extension): work is fanned out on the crate's
|
||||
//! rayon pool. A Python processor calls in from one or two worker threads
|
||||
//! with the GIL released, so intra-call parallelism is the whole point.
|
||||
//! * **feature off** (the pure-Rust `rlib` that `sglang-server` links): rayon
|
||||
//! is not even a dependency, and everything runs inline on the calling
|
||||
//! thread. A server supplies concurrency across requests and owns its own
|
||||
//! core budget (it pins threads via `core_affinity`), so a library that
|
||||
//! silently spawns its own pools would fight it.
|
||||
//!
|
||||
//! Note that "sequential" here means *inline on the caller*, not a one-thread
|
||||
//! pool: `ThreadPool::install` injects work into the pool and blocks the
|
||||
//! caller, so sizing a pool to 1 would serialize every concurrent request in
|
||||
//! the process instead of just declining to fan out.
|
||||
//!
|
||||
//! Results are identical either way — the fan-outs are order-preserving maps
|
||||
//! and writes into disjoint slices, never reductions.
|
||||
|
||||
#[cfg(feature = "parallel")]
|
||||
use rayon::prelude::*;
|
||||
|
||||
/// Map `items`, short-circuiting on the first error. Output order matches input
|
||||
/// order. CPU-bound work: decode, resize, patchify, hash.
|
||||
#[cfg(feature = "parallel")]
|
||||
pub fn try_map<'a, T, R, E>(
|
||||
items: &'a [T],
|
||||
f: impl Fn(&'a T) -> Result<R, E> + Send + Sync,
|
||||
) -> Result<Vec<R>, E>
|
||||
where
|
||||
T: Send + Sync,
|
||||
R: Send,
|
||||
E: Send,
|
||||
{
|
||||
super::pool().install(|| items.par_iter().map(f).collect())
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "parallel"))]
|
||||
pub fn try_map<'a, T, R, E>(
|
||||
items: &'a [T],
|
||||
f: impl Fn(&'a T) -> Result<R, E> + Send + Sync,
|
||||
) -> Result<Vec<R>, E>
|
||||
where
|
||||
T: Send + Sync,
|
||||
R: Send,
|
||||
E: Send,
|
||||
{
|
||||
items.iter().map(f).collect()
|
||||
}
|
||||
|
||||
/// Apply `f(chunk_index, chunk)` over disjoint `chunk_size`-element windows of
|
||||
/// `buf`. The final chunk is short when `chunk_size` does not divide the length.
|
||||
#[cfg(feature = "parallel")]
|
||||
pub fn for_chunks_mut<T: Send>(
|
||||
buf: &mut [T],
|
||||
chunk_size: usize,
|
||||
f: impl Fn(usize, &mut [T]) + Send + Sync,
|
||||
) {
|
||||
super::pool().install(|| {
|
||||
buf.par_chunks_mut(chunk_size)
|
||||
.enumerate()
|
||||
.for_each(|(index, chunk)| f(index, chunk));
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "parallel"))]
|
||||
pub fn for_chunks_mut<T: Send>(
|
||||
buf: &mut [T],
|
||||
chunk_size: usize,
|
||||
f: impl Fn(usize, &mut [T]) + Send + Sync,
|
||||
) {
|
||||
for (index, chunk) in buf.chunks_mut(chunk_size).enumerate() {
|
||||
f(index, chunk);
|
||||
}
|
||||
}
|
||||
|
||||
/// Run `f` with the CPU pool already entered, so nested [`for_chunks_mut`]
|
||||
/// calls inside it reuse this entry instead of injecting a job each. Use it to
|
||||
/// wrap a multi-stage leaf (e.g. the two passes of a separable resize) that
|
||||
/// would otherwise pay per-stage pool entry.
|
||||
#[cfg(feature = "parallel")]
|
||||
pub fn in_pool<R: Send>(f: impl FnOnce() -> R + Send) -> R {
|
||||
super::pool().install(f)
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "parallel"))]
|
||||
pub fn in_pool<R: Send>(f: impl FnOnce() -> R + Send) -> R {
|
||||
f()
|
||||
}
|
||||
@@ -1,7 +1,33 @@
|
||||
use rayon::prelude::*;
|
||||
use super::par;
|
||||
|
||||
const PRECISION_BITS: i32 = 32 - 8 - 2;
|
||||
|
||||
/// Resampling filters, bit-exact clones of PIL's kernels.
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||||
pub enum Filter {
|
||||
/// support 3.0 — PIL `LANCZOS`.
|
||||
Lanczos,
|
||||
/// support 2.0, a = -0.5 — PIL `BICUBIC` (≈ torchvision antialiased
|
||||
/// bicubic, which the HF "fast" image processors use).
|
||||
Bicubic,
|
||||
}
|
||||
|
||||
impl Filter {
|
||||
fn support(self) -> f64 {
|
||||
match self {
|
||||
Filter::Lanczos => 3.0,
|
||||
Filter::Bicubic => 2.0,
|
||||
}
|
||||
}
|
||||
|
||||
fn eval(self, x: f64) -> f64 {
|
||||
match self {
|
||||
Filter::Lanczos => lanczos(x),
|
||||
Filter::Bicubic => bicubic(x),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn sinc(x: f64) -> f64 {
|
||||
if x == 0.0 {
|
||||
return 1.0;
|
||||
@@ -18,16 +44,28 @@ fn lanczos(x: f64) -> f64 {
|
||||
}
|
||||
}
|
||||
|
||||
fn bicubic(x: f64) -> f64 {
|
||||
const A: f64 = -0.5;
|
||||
let x = x.abs();
|
||||
if x < 1.0 {
|
||||
((A + 2.0) * x - (A + 3.0)) * x * x + 1.0
|
||||
} else if x < 2.0 {
|
||||
(((x - 5.0) * x + 8.0) * x - 4.0) * A
|
||||
} else {
|
||||
0.0
|
||||
}
|
||||
}
|
||||
|
||||
struct Coeffs {
|
||||
bounds: Vec<(usize, usize)>,
|
||||
kk: Vec<i32>,
|
||||
ksize: usize,
|
||||
}
|
||||
|
||||
fn precompute_coeffs(in_size: usize, out_size: usize) -> Coeffs {
|
||||
fn precompute_coeffs(in_size: usize, out_size: usize, filter: Filter) -> Coeffs {
|
||||
let scale = in_size as f64 / out_size as f64;
|
||||
let filterscale = if scale < 1.0 { 1.0 } else { scale };
|
||||
let support = 3.0 * filterscale;
|
||||
let support = filter.support() * filterscale;
|
||||
let ksize = support.ceil() as usize * 2 + 1;
|
||||
let ss = 1.0 / filterscale;
|
||||
|
||||
@@ -47,7 +85,7 @@ fn precompute_coeffs(in_size: usize, out_size: usize) -> Coeffs {
|
||||
let k = &mut kkf[xx * ksize..(xx + 1) * ksize];
|
||||
let mut ww = 0.0f64;
|
||||
for (x, kv) in k[..count].iter_mut().enumerate() {
|
||||
let w = lanczos((x as f64 + xmin as f64 - center + 0.5) * ss);
|
||||
let w = filter.eval((x as f64 + xmin as f64 - center + 0.5) * ss);
|
||||
*kv = w;
|
||||
ww += w;
|
||||
}
|
||||
@@ -86,32 +124,30 @@ fn clip8(v: i32) -> u8 {
|
||||
|
||||
fn resample_horizontal(src: &[u8], h: usize, w: usize, out_w: usize, c: &Coeffs) -> Vec<u8> {
|
||||
let mut out = vec![0u8; h * out_w * 3];
|
||||
out.par_chunks_mut(out_w * 3)
|
||||
.enumerate()
|
||||
.for_each(|(y, row)| {
|
||||
let src_row = &src[y * w * 3..(y + 1) * w * 3];
|
||||
for xx in 0..out_w {
|
||||
let (xmin, count) = c.bounds[xx];
|
||||
let k = &c.kk[xx * c.ksize..xx * c.ksize + count];
|
||||
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
|
||||
for (x, &coef) in k.iter().enumerate() {
|
||||
let p = (xmin + x) * 3;
|
||||
s[0] += src_row[p] as i32 * coef;
|
||||
s[1] += src_row[p + 1] as i32 * coef;
|
||||
s[2] += src_row[p + 2] as i32 * coef;
|
||||
}
|
||||
let o = xx * 3;
|
||||
row[o] = clip8(s[0]);
|
||||
row[o + 1] = clip8(s[1]);
|
||||
row[o + 2] = clip8(s[2]);
|
||||
par::for_chunks_mut(&mut out, out_w * 3, |y, row| {
|
||||
let src_row = &src[y * w * 3..(y + 1) * w * 3];
|
||||
for xx in 0..out_w {
|
||||
let (xmin, count) = c.bounds[xx];
|
||||
let k = &c.kk[xx * c.ksize..xx * c.ksize + count];
|
||||
let mut s = [1i32 << (PRECISION_BITS - 1); 3];
|
||||
for (x, &coef) in k.iter().enumerate() {
|
||||
let p = (xmin + x) * 3;
|
||||
s[0] += src_row[p] as i32 * coef;
|
||||
s[1] += src_row[p + 1] as i32 * coef;
|
||||
s[2] += src_row[p + 2] as i32 * coef;
|
||||
}
|
||||
});
|
||||
let o = xx * 3;
|
||||
row[o] = clip8(s[0]);
|
||||
row[o + 1] = clip8(s[1]);
|
||||
row[o + 2] = clip8(s[2]);
|
||||
}
|
||||
});
|
||||
out
|
||||
}
|
||||
|
||||
fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec<u8> {
|
||||
let mut out = vec![0u8; out_h * w * 3];
|
||||
out.par_chunks_mut(w * 3).enumerate().for_each(|(yy, row)| {
|
||||
par::for_chunks_mut(&mut out, w * 3, |yy, row| {
|
||||
let (ymin, count) = c.bounds[yy];
|
||||
let k = &c.kk[yy * c.ksize..yy * c.ksize + count];
|
||||
for x in 0..w {
|
||||
@@ -131,25 +167,51 @@ fn resample_vertical(src: &[u8], w: usize, out_h: usize, c: &Coeffs) -> Vec<u8>
|
||||
out
|
||||
}
|
||||
|
||||
pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec<u8> {
|
||||
/// PIL-exact separable resize of a flat HWC RGB buffer with the given filter.
|
||||
///
|
||||
/// Enters the fan-out pool once for both passes; the per-row `for_chunks_mut`
|
||||
/// calls inside then reuse that entry rather than injecting a job per pass.
|
||||
pub fn resize_rgb_filter(
|
||||
src: &[u8],
|
||||
h: usize,
|
||||
w: usize,
|
||||
out_h: usize,
|
||||
out_w: usize,
|
||||
filter: Filter,
|
||||
) -> Vec<u8> {
|
||||
par::in_pool(move || resize_passes(src, h, w, out_h, out_w, filter))
|
||||
}
|
||||
|
||||
fn resize_passes(
|
||||
src: &[u8],
|
||||
h: usize,
|
||||
w: usize,
|
||||
out_h: usize,
|
||||
out_w: usize,
|
||||
filter: Filter,
|
||||
) -> Vec<u8> {
|
||||
let need_h = out_w != w;
|
||||
let need_v = out_h != h;
|
||||
if need_h && need_v {
|
||||
let ch = precompute_coeffs(w, out_w);
|
||||
let ch = precompute_coeffs(w, out_w, filter);
|
||||
let tmp = resample_horizontal(src, h, w, out_w, &ch);
|
||||
let cv = precompute_coeffs(h, out_h);
|
||||
let cv = precompute_coeffs(h, out_h, filter);
|
||||
resample_vertical(&tmp, out_w, out_h, &cv)
|
||||
} else if need_h {
|
||||
let ch = precompute_coeffs(w, out_w);
|
||||
let ch = precompute_coeffs(w, out_w, filter);
|
||||
resample_horizontal(src, h, w, out_w, &ch)
|
||||
} else if need_v {
|
||||
let cv = precompute_coeffs(h, out_h);
|
||||
let cv = precompute_coeffs(h, out_h, filter);
|
||||
resample_vertical(src, w, out_h, &cv)
|
||||
} else {
|
||||
src.to_vec()
|
||||
}
|
||||
}
|
||||
|
||||
pub fn resize_lanczos_rgb(src: &[u8], h: usize, w: usize, out_h: usize, out_w: usize) -> Vec<u8> {
|
||||
resize_rgb_filter(src, h, w, out_h, out_w, Filter::Lanczos)
|
||||
}
|
||||
|
||||
pub fn scaled_dims(w: usize, h: usize, frac: Option<f64>, cap: Option<i64>) -> (usize, usize) {
|
||||
let Some(frac) = frac else {
|
||||
return (w, h);
|
||||
|
||||
@@ -0,0 +1,239 @@
|
||||
//! Token-layout mechanics for the server MM pipeline.
|
||||
//!
|
||||
//! Families describe their prompt geometry as a [`TokenLayout`] value
|
||||
//! (`pipeline.rs`); [`apply_layout`] applies it mechanically. Expanding the
|
||||
//! already-tokenized prompt means non-media tokens can never drift from a
|
||||
//! retokenize (the `SGLANG_MM_AVOID_RETOKENIZE` idea, unconditional here).
|
||||
|
||||
use crate::pipeline::{Segment, TokenLayout, TokenPattern};
|
||||
|
||||
/// The expanded prompt plus, per media item (indexed as in the layout), the
|
||||
/// inclusive `(start, end)` token range it occupies — the Python
|
||||
/// `get_mm_items_offset` convention.
|
||||
pub struct ExpandedPrompt {
|
||||
pub input_ids: Vec<i32>,
|
||||
pub offsets: Vec<(u32, u32)>,
|
||||
}
|
||||
|
||||
/// Apply a family's [`TokenLayout`] to the original prompt.
|
||||
///
|
||||
/// The point of the layout being data is that a family cannot get expansion,
|
||||
/// offsets, and positions out of sync, so this validates the whole contract
|
||||
/// rather than just indexing safely:
|
||||
/// * text ranges are in bounds, ascending, and non-overlapping;
|
||||
/// * together with the media placeholders they cover every source token
|
||||
/// exactly once — a family that forgets a tail segment must not silently
|
||||
/// truncate the prompt;
|
||||
/// * every one of the `n_items` media items is placed exactly once;
|
||||
/// * no item expands to zero tokens (which would have no representable offset).
|
||||
pub fn apply_layout(
|
||||
src: &[i32],
|
||||
layout: &TokenLayout,
|
||||
n_items: usize,
|
||||
) -> Result<ExpandedPrompt, String> {
|
||||
let mut out = Vec::new();
|
||||
let mut offsets: Vec<Option<(u32, u32)>> = vec![None; n_items];
|
||||
// Source tokens consumed so far: `Text` copies them, a `Media` segment
|
||||
// replaces exactly the one placeholder token sitting at this position.
|
||||
let mut consumed = 0usize;
|
||||
for segment in &layout.segments {
|
||||
match segment {
|
||||
Segment::Text(range) => {
|
||||
let text = src
|
||||
.get(range.clone())
|
||||
.ok_or_else(|| format!("layout: text range {range:?} out of bounds"))?;
|
||||
if range.start != consumed {
|
||||
return Err(format!(
|
||||
"layout: text range {range:?} does not resume at source index {consumed}"
|
||||
));
|
||||
}
|
||||
consumed = range.end;
|
||||
out.extend_from_slice(text);
|
||||
}
|
||||
Segment::Media { item, pattern } => {
|
||||
if consumed >= src.len() {
|
||||
return Err(format!(
|
||||
"layout: media item {item} has no source placeholder at index {consumed}"
|
||||
));
|
||||
}
|
||||
consumed += 1;
|
||||
let start = out.len() as u32;
|
||||
let n = match pattern {
|
||||
TokenPattern::Repeat { id, n } => {
|
||||
out.resize(out.len() + n, *id);
|
||||
*n
|
||||
}
|
||||
TokenPattern::Explicit(ids) => {
|
||||
out.extend_from_slice(ids);
|
||||
ids.len()
|
||||
}
|
||||
};
|
||||
if n == 0 {
|
||||
return Err(format!("layout: media item {item} expands to zero tokens"));
|
||||
}
|
||||
let slot = offsets
|
||||
.get_mut(*item)
|
||||
.ok_or_else(|| format!("layout: media item {item} out of range"))?;
|
||||
if slot.replace((start, start + n as u32 - 1)).is_some() {
|
||||
return Err(format!("layout: media item {item} placed twice"));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if consumed != src.len() {
|
||||
return Err(format!(
|
||||
"layout: covers {consumed} of {} source token(s)",
|
||||
src.len()
|
||||
));
|
||||
}
|
||||
let offsets = offsets
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(i, slot)| slot.ok_or_else(|| format!("layout: media item {i} not placed")))
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(ExpandedPrompt {
|
||||
input_ids: out,
|
||||
offsets,
|
||||
})
|
||||
}
|
||||
|
||||
/// Build the simplest layout: each occurrence of `placeholder_id` in `ids`
|
||||
/// becomes `counts[i]` copies (i-th occurrence ↔ i-th media item). Errs when
|
||||
/// the occurrence count and `counts` disagree.
|
||||
pub fn layout_by_placeholder(
|
||||
ids: &[i32],
|
||||
placeholder_id: i32,
|
||||
counts: &[usize],
|
||||
) -> Result<TokenLayout, String> {
|
||||
let found = ids.iter().filter(|&&id| id == placeholder_id).count();
|
||||
if found != counts.len() {
|
||||
return Err(format!(
|
||||
"prompt has {found} media placeholder(s) but {} media item(s)",
|
||||
counts.len()
|
||||
));
|
||||
}
|
||||
let mut segments = Vec::new();
|
||||
let mut text_start = 0;
|
||||
let mut item = 0;
|
||||
for (pos, &id) in ids.iter().enumerate() {
|
||||
if id == placeholder_id {
|
||||
if text_start < pos {
|
||||
segments.push(Segment::Text(text_start..pos));
|
||||
}
|
||||
segments.push(Segment::Media {
|
||||
item,
|
||||
pattern: TokenPattern::Repeat {
|
||||
id: placeholder_id,
|
||||
n: counts[item],
|
||||
},
|
||||
});
|
||||
item += 1;
|
||||
text_start = pos + 1;
|
||||
}
|
||||
}
|
||||
if text_start < ids.len() {
|
||||
segments.push(Segment::Text(text_start..ids.len()));
|
||||
}
|
||||
Ok(TokenLayout { segments })
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn expand(ids: &[i32], placeholder: i32, counts: &[usize]) -> Result<ExpandedPrompt, String> {
|
||||
apply_layout(
|
||||
ids,
|
||||
&layout_by_placeholder(ids, placeholder, counts)?,
|
||||
counts.len(),
|
||||
)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn expands_in_order_with_inclusive_offsets() {
|
||||
// [7, PAD, 8, PAD, 9] with counts [2, 3]
|
||||
let e = expand(&[7, 1, 8, 1, 9], 1, &[2, 3]).unwrap();
|
||||
assert_eq!(e.input_ids, vec![7, 1, 1, 8, 1, 1, 1, 9]);
|
||||
assert_eq!(e.offsets, vec![(1, 2), (4, 6)]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn count_mismatch_errs() {
|
||||
assert!(expand(&[7, 1, 9], 1, &[2, 3]).is_err());
|
||||
assert!(expand(&[7, 1, 1, 9], 1, &[2]).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn zero_count_errs() {
|
||||
assert!(expand(&[7, 1, 9], 1, &[0]).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_placeholders_no_items_ok() {
|
||||
let e = expand(&[7, 8], 1, &[]).unwrap();
|
||||
assert_eq!(e.input_ids, vec![7, 8]);
|
||||
assert!(e.offsets.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn explicit_patterns_and_placement_validation() {
|
||||
// Structured expansion: marker tokens around the item span.
|
||||
let layout = TokenLayout {
|
||||
segments: vec![
|
||||
Segment::Text(0..1),
|
||||
Segment::Media {
|
||||
item: 0,
|
||||
pattern: TokenPattern::Explicit(vec![90, 5, 5, 91]),
|
||||
},
|
||||
Segment::Text(2..3),
|
||||
],
|
||||
};
|
||||
let e = apply_layout(&[7, 1, 9], &layout, 1).unwrap();
|
||||
assert_eq!(e.input_ids, vec![7, 90, 5, 5, 91, 9]);
|
||||
assert_eq!(e.offsets, vec![(1, 4)]);
|
||||
|
||||
// Every item must be placed exactly once; ranges must be in bounds.
|
||||
let missing = TokenLayout {
|
||||
segments: vec![Segment::Text(0..3)],
|
||||
};
|
||||
assert!(apply_layout(&[7, 1, 9], &missing, 1).is_err());
|
||||
let out_of_bounds = TokenLayout {
|
||||
segments: vec![Segment::Text(0..4)],
|
||||
};
|
||||
assert!(apply_layout(&[7, 1, 9], &out_of_bounds, 0).is_err());
|
||||
}
|
||||
|
||||
/// A family that skips, repeats, or reorders source tokens would silently
|
||||
/// serve a truncated or scrambled prompt; the layout must reject it instead.
|
||||
#[test]
|
||||
fn incomplete_or_disordered_coverage_errs() {
|
||||
let media = || Segment::Media {
|
||||
item: 0,
|
||||
pattern: TokenPattern::Repeat { id: 5, n: 2 },
|
||||
};
|
||||
let cases = [
|
||||
// Dropped tail: [7, PAD, 9] expanded without the trailing 9.
|
||||
vec![Segment::Text(0..1), media()],
|
||||
// Dropped head.
|
||||
vec![media(), Segment::Text(2..3)],
|
||||
// Gap in the middle (source index 1 never consumed).
|
||||
vec![Segment::Text(0..1), Segment::Text(2..3), media()],
|
||||
// Duplicated source span.
|
||||
vec![
|
||||
Segment::Text(0..1),
|
||||
Segment::Text(0..1),
|
||||
media(),
|
||||
Segment::Text(2..3),
|
||||
],
|
||||
// Out of order.
|
||||
vec![Segment::Text(2..3), media(), Segment::Text(0..1)],
|
||||
];
|
||||
for (i, segments) in cases.into_iter().enumerate() {
|
||||
let layout = TokenLayout { segments };
|
||||
assert!(
|
||||
apply_layout(&[7, 1, 9], &layout, 1).is_err(),
|
||||
"case {i} should be rejected"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,280 @@
|
||||
//! Shared multimodal request driver for the server (pure-Rust) pipeline.
|
||||
//!
|
||||
//! Owns the request control flow — parallel fan-out, layout application,
|
||||
//! failure semantics — while every model decision lives behind
|
||||
//! [`MmFamilyProcessor`] (see `pipeline.rs`). Families produce data; they
|
||||
//! cannot alter orchestration.
|
||||
|
||||
use crate::common::{self, fetch, par, token_layout};
|
||||
use crate::pipeline::{DecodedMedia, MmFamilyProcessor, PositionOutput, ProcessedItem};
|
||||
|
||||
/// Per-request bounds: together with [`fetch::MAX_FETCH_BYTES`] they cap what
|
||||
/// one request can make the pipeline buffer.
|
||||
pub const MAX_ITEMS_PER_REQUEST: usize = 64;
|
||||
pub const MAX_REQUEST_BYTES: u64 = 256 << 20;
|
||||
|
||||
/// One raw image source from the request.
|
||||
#[derive(Debug)]
|
||||
pub enum ImageSource {
|
||||
/// `data:`/base64/file/http — resolved by [`fetch::fetch_bytes`].
|
||||
String(String),
|
||||
/// Already-raw encoded image bytes.
|
||||
Bytes(Vec<u8>),
|
||||
}
|
||||
|
||||
/// Typed multimodal request input. The server's message layer owns the wire
|
||||
/// format and parses its payload into this before calling [`process`].
|
||||
pub struct MmInput {
|
||||
pub text: Option<String>,
|
||||
pub input_ids: Option<Vec<i32>>,
|
||||
pub images: Vec<ImageSource>,
|
||||
}
|
||||
|
||||
/// One processed media item at the request boundary.
|
||||
pub struct OutputItem {
|
||||
pub feature: crate::pipeline::Tensor,
|
||||
pub aux: crate::pipeline::NamedTensors,
|
||||
/// [`common::content_hash_u64`] of the raw encoded source bytes — the same
|
||||
/// identity role as the Python path's `hash_feature`, but a different
|
||||
/// algorithm, so hashes are consistent within the server pipeline and
|
||||
/// never comparable across the two paths.
|
||||
pub hash: u64,
|
||||
}
|
||||
|
||||
/// The per-request result parked for the scheduler drain.
|
||||
pub struct Output {
|
||||
pub input_ids: Vec<i32>,
|
||||
/// In prompt order; `offsets[i]` is `items[i]`'s inclusive token range.
|
||||
pub items: Vec<OutputItem>,
|
||||
pub offsets: Vec<(u32, u32)>,
|
||||
pub positions: PositionOutput,
|
||||
}
|
||||
|
||||
/// Resolve one image source to raw encoded bytes, borrowing when the request
|
||||
/// already carries them.
|
||||
fn resolve(source: &ImageSource) -> Result<std::borrow::Cow<'_, [u8]>, String> {
|
||||
match source {
|
||||
ImageSource::String(source) => Ok(fetch::fetch_bytes(source)?.into()),
|
||||
ImageSource::Bytes(bytes) => Ok(bytes.as_slice().into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Run one request through the pipeline. Any `Err` rejects the request back
|
||||
/// to the client — including inputs merely outside the pipeline's scope
|
||||
/// (video/audio, precomputed features, undecodable images), since there is
|
||||
/// no Python fallback path.
|
||||
pub fn process(
|
||||
family: &dyn MmFamilyProcessor,
|
||||
input: MmInput,
|
||||
tokenize: impl FnOnce(&str) -> Result<Vec<i32>, String>,
|
||||
) -> Result<Output, String> {
|
||||
if input.images.is_empty() {
|
||||
return Err("multimodal request without image sources".into());
|
||||
}
|
||||
if input.images.len() > MAX_ITEMS_PER_REQUEST {
|
||||
return Err(format!(
|
||||
"multimodal request exceeds {MAX_ITEMS_PER_REQUEST} media items"
|
||||
));
|
||||
}
|
||||
// Stage 1 (fetch) is blocking I/O and runs inline, sequentially — never on
|
||||
// the CPU pool, where a slow URL or file read would starve decode/resize
|
||||
// for other requests. Contract: callers on a fixed worker pool (the
|
||||
// server) must resolve I/O-backed string sources — URLs and file paths —
|
||||
// on their own I/O layer and pass `Bytes`.
|
||||
let mut fetched: Vec<std::borrow::Cow<'_, [u8]>> = Vec::with_capacity(input.images.len());
|
||||
let mut total: u64 = 0;
|
||||
for source in &input.images {
|
||||
let bytes = resolve(source)?;
|
||||
total += bytes.len() as u64;
|
||||
if total > MAX_REQUEST_BYTES {
|
||||
return Err(format!(
|
||||
"multimodal request exceeds {MAX_REQUEST_BYTES} total media bytes"
|
||||
));
|
||||
}
|
||||
fetched.push(bytes);
|
||||
}
|
||||
let processed: Vec<(ProcessedItem, u64)> =
|
||||
par::try_map(&fetched, |bytes| -> Result<(ProcessedItem, u64), String> {
|
||||
let hash = common::content_hash_u64(bytes);
|
||||
// Inputs PIL accepts but decode_rgb refuses (e.g. 16-bit PNG)
|
||||
// error here and reject the request.
|
||||
let (rgb, height, width) = common::decode_rgb(bytes)?;
|
||||
let item = family.process_item(&DecodedMedia::Image { rgb, height, width })?;
|
||||
Ok((item, hash))
|
||||
})?;
|
||||
|
||||
let input_ids = match input.input_ids {
|
||||
Some(input_ids) if !input_ids.is_empty() => input_ids,
|
||||
_ => {
|
||||
let text = input
|
||||
.text
|
||||
.as_deref()
|
||||
.ok_or("multimodal request without text or input_ids")?;
|
||||
tokenize(text)?
|
||||
}
|
||||
};
|
||||
let geometries = processed
|
||||
.iter()
|
||||
.map(|(item, _)| item.geometry.clone())
|
||||
.collect::<Vec<_>>();
|
||||
let layout = family.layout(&input_ids, &geometries)?;
|
||||
let expanded = token_layout::apply_layout(&input_ids, &layout, processed.len())?;
|
||||
let positions = family.positions(expanded.input_ids.len(), &expanded.offsets, &geometries)?;
|
||||
|
||||
Ok(Output {
|
||||
input_ids: expanded.input_ids,
|
||||
items: processed
|
||||
.into_iter()
|
||||
.map(|(item, hash)| OutputItem {
|
||||
feature: item.feature,
|
||||
aux: item.aux,
|
||||
hash,
|
||||
})
|
||||
.collect(),
|
||||
offsets: expanded.offsets,
|
||||
positions,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::registry::pipeline_from_spec;
|
||||
|
||||
const SPEC: &str = r#"{"family":"qwen_vl","image_token_id":1,"patch_size":2,
|
||||
"merge_size":2,"temporal_patch_size":2,"min_pixels":4,
|
||||
"max_pixels":1073741824,"image_mean":[0.0,0.0,0.0],"image_std":[1.0,1.0,1.0]}"#;
|
||||
|
||||
fn png(w: u32, h: u32) -> Vec<u8> {
|
||||
let img = image::RgbImage::from_fn(w, h, |x, y| image::Rgb([x as u8, y as u8, 7]));
|
||||
let mut buf = std::io::Cursor::new(Vec::new());
|
||||
img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
|
||||
buf.into_inner()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn processes_typed_image_request() {
|
||||
let family = pipeline_from_spec(SPEC).unwrap();
|
||||
let input = MmInput {
|
||||
text: None,
|
||||
input_ids: Some(vec![7, 1, 8]),
|
||||
images: vec![ImageSource::Bytes(png(8, 8))],
|
||||
};
|
||||
let out = process(family.as_ref(), input, |_| Err("no tokenizer".into())).unwrap();
|
||||
// 8x8, factor 4 → grid [1, 4, 4] → 16 patches / merge² = 4 tokens.
|
||||
assert_eq!(out.input_ids, vec![7, 1, 1, 1, 1, 8]);
|
||||
assert_eq!(out.offsets, vec![(1, 4)]);
|
||||
let item = &out.items[0];
|
||||
assert_eq!(item.feature.shape, [16, 3 * 2 * 2 * 2]);
|
||||
assert_eq!(item.aux[0].0, "image_grid_thw");
|
||||
let crate::pipeline::PositionOutput::MRope { positions, .. } = &out.positions else {
|
||||
panic!("qwen emits mrope")
|
||||
};
|
||||
assert_eq!(positions.len(), 3 * out.input_ids.len());
|
||||
}
|
||||
|
||||
/// String sources — HTTP URL, `file://`, and bare path — all resolve to
|
||||
/// the same bytes and flow through the full pipeline.
|
||||
#[test]
|
||||
fn fetches_url_and_file_sources() {
|
||||
let png = png(8, 8);
|
||||
let dir = std::env::temp_dir().join(format!("sglang-mm-driver-{}", std::process::id()));
|
||||
std::fs::create_dir_all(&dir).unwrap();
|
||||
let path = dir.join("img.png");
|
||||
std::fs::write(&path, &png).unwrap();
|
||||
|
||||
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let body = png.clone();
|
||||
let server = std::thread::spawn(move || {
|
||||
use std::io::{BufRead, Write};
|
||||
let (stream, _) = listener.accept().unwrap();
|
||||
let mut reader = std::io::BufReader::new(stream);
|
||||
let mut line = String::new();
|
||||
while reader.read_line(&mut line).unwrap() > 2 {
|
||||
line.clear(); // headers until the blank line
|
||||
}
|
||||
let mut stream = reader.into_inner();
|
||||
write!(
|
||||
stream,
|
||||
"HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n",
|
||||
body.len()
|
||||
)
|
||||
.unwrap();
|
||||
stream.write_all(&body).unwrap();
|
||||
});
|
||||
|
||||
let family = pipeline_from_spec(SPEC).unwrap();
|
||||
let input = MmInput {
|
||||
text: None,
|
||||
input_ids: Some(vec![7, 1, 1, 1, 8]),
|
||||
images: vec![
|
||||
ImageSource::String(format!("http://{addr}/img.png")),
|
||||
ImageSource::String(format!("file://{}", path.display())),
|
||||
ImageSource::String(path.display().to_string()),
|
||||
],
|
||||
};
|
||||
let out = process(family.as_ref(), input, |_| unreachable!()).unwrap();
|
||||
server.join().unwrap();
|
||||
std::fs::remove_dir_all(&dir).ok();
|
||||
|
||||
assert_eq!(out.items.len(), 3);
|
||||
// Identical source bytes → identical content hashes.
|
||||
assert_eq!(out.items[0].hash, out.items[1].hash);
|
||||
assert_eq!(out.items[1].hash, out.items[2].hash);
|
||||
assert_eq!(out.input_ids.len(), 5 + 3 * 3); // each placeholder → 4 tokens
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn per_request_caps_enforced() {
|
||||
let family = pipeline_from_spec(SPEC).unwrap();
|
||||
let too_many = MmInput {
|
||||
text: None,
|
||||
input_ids: Some(vec![1]),
|
||||
images: (0..=MAX_ITEMS_PER_REQUEST)
|
||||
.map(|_| ImageSource::Bytes(vec![]))
|
||||
.collect(),
|
||||
};
|
||||
let err = process(family.as_ref(), too_many, |_| unreachable!())
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(err.contains("media items"), "{err}");
|
||||
|
||||
let chunk = (MAX_REQUEST_BYTES / 2 + 1) as usize;
|
||||
let too_big = MmInput {
|
||||
text: None,
|
||||
input_ids: Some(vec![1, 1]),
|
||||
images: vec![
|
||||
ImageSource::Bytes(vec![0; chunk]),
|
||||
ImageSource::Bytes(vec![0; chunk]),
|
||||
],
|
||||
};
|
||||
let err = process(family.as_ref(), too_big, |_| unreachable!())
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(err.contains("total media bytes"), "{err}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn image_free_and_mismatched_requests_rejected() {
|
||||
let family = pipeline_from_spec(SPEC).unwrap();
|
||||
let no_images = MmInput {
|
||||
text: None,
|
||||
input_ids: Some(vec![7, 1]),
|
||||
images: vec![],
|
||||
};
|
||||
let err = process(family.as_ref(), no_images, |_| unreachable!())
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(err.contains("image sources"));
|
||||
let no_placeholder = MmInput {
|
||||
text: None,
|
||||
input_ids: Some(vec![7, 8]),
|
||||
images: vec![ImageSource::Bytes(png(8, 8))],
|
||||
};
|
||||
let err = process(family.as_ref(), no_placeholder, |_| unreachable!())
|
||||
.err()
|
||||
.unwrap();
|
||||
assert!(err.contains("placeholder"));
|
||||
}
|
||||
}
|
||||
+166
-177
@@ -1,21 +1,8 @@
|
||||
use std::sync::OnceLock;
|
||||
|
||||
use half::bf16;
|
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
use rayon::prelude::*;
|
||||
|
||||
use crate::common;
|
||||
|
||||
/// `(height, width, patches_as_u16_bits)` for one decoded image.
|
||||
type Patches = (usize, usize, Vec<u16>);
|
||||
/// [`Patches`] plus the image content hash.
|
||||
type HashedPatches = (usize, usize, Vec<u16>, u64);
|
||||
/// [`Patches`] with the patch data as a numpy array bound to `'py`.
|
||||
type PyPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>);
|
||||
/// [`PyPatches`] plus the image content hash.
|
||||
type PyHashedPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>, u64);
|
||||
use crate::common::par;
|
||||
use half::bf16;
|
||||
|
||||
const MEAN: [f32; 3] = [
|
||||
0.48145466f64 as f32,
|
||||
@@ -90,9 +77,7 @@ fn patchify_into(arr: &[u8], h: usize, w: usize, ps: usize, out: &mut [u16]) {
|
||||
}
|
||||
};
|
||||
|
||||
common::pool().install(|| {
|
||||
out.par_chunks_mut(row_elems).enumerate().for_each(body);
|
||||
});
|
||||
par::for_chunks_mut(out, row_elems, |index, chunk| body((index, chunk)));
|
||||
}
|
||||
|
||||
fn patchify_alloc(arr: &[u8], h: usize, w: usize, ps: usize) -> Vec<u16> {
|
||||
@@ -102,119 +87,6 @@ fn patchify_alloc(arr: &[u8], h: usize, w: usize, ps: usize) -> Vec<u16> {
|
||||
out
|
||||
}
|
||||
|
||||
fn check_ps(ps: usize) -> PyResult<()> {
|
||||
if ps == 0 {
|
||||
return Err(PyValueError::new_err(
|
||||
"patch_size must be greater than zero",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn patchify_rgb<'py>(
|
||||
py: Python<'py>,
|
||||
arr: PyReadonlyArray3<'py, u8>,
|
||||
patch_size: usize,
|
||||
) -> PyResult<Bound<'py, PyArray1<u16>>> {
|
||||
check_ps(patch_size)?;
|
||||
let shape = arr.shape();
|
||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||
if c != 3 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"expected HWC RGB array with 3 channels, got {c}"
|
||||
)));
|
||||
}
|
||||
let data = arr
|
||||
.as_slice()
|
||||
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
||||
.to_vec();
|
||||
let out = py.detach(move || patchify_alloc(&data, h, w, patch_size));
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (data, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn decode_patchify<'py>(
|
||||
py: Python<'py>,
|
||||
data: Vec<u8>,
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u16>>)> {
|
||||
check_ps(patch_size)?;
|
||||
let (h, w, out) = py
|
||||
.detach(move || {
|
||||
common::pool().install(|| {
|
||||
let (rgb, h, w) = common::decode_rescale(&data, rescale_frac, rescale_cap)?;
|
||||
Ok::<_, String>((h, w, patchify_alloc(&rgb, h, w, patch_size)))
|
||||
})
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
Ok((h, w, out.into_pyarray(py)))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn decode_patchify_batch<'py>(
|
||||
py: Python<'py>,
|
||||
datas: Vec<Vec<u8>>,
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<Vec<PyPatches<'py>>> {
|
||||
check_ps(patch_size)?;
|
||||
let results: Vec<Result<Patches, String>> = py.detach(move || {
|
||||
common::pool().install(|| {
|
||||
datas
|
||||
.par_iter()
|
||||
.map(|data| {
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size)))
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
});
|
||||
results
|
||||
.into_iter()
|
||||
.map(|r| {
|
||||
let (h, w, v) = r.map_err(PyValueError::new_err)?;
|
||||
Ok((h, w, v.into_pyarray(py)))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn preprocess_images<'py>(
|
||||
py: Python<'py>,
|
||||
datas: Vec<Vec<u8>>,
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<Vec<PyHashedPatches<'py>>> {
|
||||
check_ps(patch_size)?;
|
||||
let results: Vec<Result<HashedPatches, String>> = py.detach(move || {
|
||||
common::pool().install(|| {
|
||||
datas
|
||||
.par_iter()
|
||||
.map(|data| {
|
||||
let hash = common::sha256_u64(data);
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
});
|
||||
results
|
||||
.into_iter()
|
||||
.map(|r| {
|
||||
let (h, w, v, hash) = r.map_err(PyValueError::new_err)?;
|
||||
Ok((h, w, v.into_pyarray(py), hash))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Struct implementing ImageProcessorSpec for Inkling.
|
||||
pub struct InklingProcessor;
|
||||
|
||||
@@ -233,44 +105,158 @@ impl crate::registry::ImageProcessorSpec for InklingProcessor {
|
||||
if patch_size == 0 {
|
||||
return Err("patch_size must be greater than zero".into());
|
||||
}
|
||||
common::pool().install(|| {
|
||||
datas
|
||||
.par_iter()
|
||||
.map(|data| {
|
||||
let hash = common::sha256_u64(data);
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
|
||||
})
|
||||
.collect()
|
||||
par::try_map(datas, |data| {
|
||||
let hash = common::content_hash_u64(data);
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (arr, raw_bytes, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn rescale_patchify_hash<'py>(
|
||||
py: Python<'py>,
|
||||
arr: PyReadonlyArray3<'py, u8>,
|
||||
raw_bytes: &[u8],
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u16>>, u64)> {
|
||||
check_ps(patch_size)?;
|
||||
let shape = arr.shape();
|
||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||
if c != 3 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"expected HWC RGB array with 3 channels, got {c}"
|
||||
)));
|
||||
// --- Python bindings (feature-gated: absent from the pure-Rust rlib) ---
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
mod python {
|
||||
use numpy::{IntoPyArray, PyArray1, PyReadonlyArray3, PyUntypedArrayMethods};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::patchify_alloc;
|
||||
use crate::common;
|
||||
use crate::common::par;
|
||||
|
||||
/// `(height, width, patches_as_u16_bits)` for one decoded image.
|
||||
type Patches = (usize, usize, Vec<u16>);
|
||||
/// [`Patches`] plus the image content hash.
|
||||
type HashedPatches = (usize, usize, Vec<u16>, u64);
|
||||
/// [`Patches`] with the patch data as a numpy array bound to `'py`.
|
||||
type PyPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>);
|
||||
/// [`PyPatches`] plus the image content hash.
|
||||
type PyHashedPatches<'py> = (usize, usize, Bound<'py, PyArray1<u16>>, u64);
|
||||
|
||||
fn check_ps(ps: usize) -> PyResult<()> {
|
||||
if ps == 0 {
|
||||
return Err(PyValueError::new_err(
|
||||
"patch_size must be greater than zero",
|
||||
));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
let hash = common::sha256_u64(raw_bytes);
|
||||
let rgb = arr
|
||||
.as_slice()
|
||||
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
||||
.to_vec();
|
||||
let (oh, ow, out) = py.detach(move || {
|
||||
common::pool().install(|| {
|
||||
|
||||
#[pyfunction]
|
||||
fn patchify_rgb<'py>(
|
||||
py: Python<'py>,
|
||||
arr: PyReadonlyArray3<'py, u8>,
|
||||
patch_size: usize,
|
||||
) -> PyResult<Bound<'py, PyArray1<u16>>> {
|
||||
check_ps(patch_size)?;
|
||||
let shape = arr.shape();
|
||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||
if c != 3 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"expected HWC RGB array with 3 channels, got {c}"
|
||||
)));
|
||||
}
|
||||
let data = arr
|
||||
.as_slice()
|
||||
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
||||
.to_vec();
|
||||
let out = py.detach(move || patchify_alloc(&data, h, w, patch_size));
|
||||
Ok(out.into_pyarray(py))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (data, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn decode_patchify<'py>(
|
||||
py: Python<'py>,
|
||||
data: Vec<u8>,
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u16>>)> {
|
||||
check_ps(patch_size)?;
|
||||
let (h, w, out) = py
|
||||
.detach(move || {
|
||||
let (rgb, h, w) = common::decode_rescale(&data, rescale_frac, rescale_cap)?;
|
||||
Ok::<_, String>((h, w, patchify_alloc(&rgb, h, w, patch_size)))
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
Ok((h, w, out.into_pyarray(py)))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn decode_patchify_batch<'py>(
|
||||
py: Python<'py>,
|
||||
datas: Vec<Vec<u8>>,
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<Vec<PyPatches<'py>>> {
|
||||
check_ps(patch_size)?;
|
||||
let results: Vec<Patches> = py
|
||||
.detach(move || {
|
||||
par::try_map(&datas, |data| -> Result<Patches, String> {
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size)))
|
||||
})
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
results
|
||||
.into_iter()
|
||||
.map(|(h, w, v)| Ok((h, w, v.into_pyarray(py))))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (datas, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn preprocess_images<'py>(
|
||||
py: Python<'py>,
|
||||
datas: Vec<Vec<u8>>,
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<Vec<PyHashedPatches<'py>>> {
|
||||
check_ps(patch_size)?;
|
||||
let results: Vec<HashedPatches> = py
|
||||
.detach(move || {
|
||||
par::try_map(&datas, |data| -> Result<HashedPatches, String> {
|
||||
let hash = common::content_hash_u64(data);
|
||||
let (rgb, h, w) = common::decode_rescale(data, rescale_frac, rescale_cap)?;
|
||||
Ok((h, w, patchify_alloc(&rgb, h, w, patch_size), hash))
|
||||
})
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
results
|
||||
.into_iter()
|
||||
.map(|(h, w, v, hash)| Ok((h, w, v.into_pyarray(py), hash)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (arr, raw_bytes, patch_size, rescale_frac=None, rescale_cap=None))]
|
||||
fn rescale_patchify_hash<'py>(
|
||||
py: Python<'py>,
|
||||
arr: PyReadonlyArray3<'py, u8>,
|
||||
raw_bytes: &[u8],
|
||||
patch_size: usize,
|
||||
rescale_frac: Option<f64>,
|
||||
rescale_cap: Option<i64>,
|
||||
) -> PyResult<(usize, usize, Bound<'py, PyArray1<u16>>, u64)> {
|
||||
check_ps(patch_size)?;
|
||||
let shape = arr.shape();
|
||||
let (h, w, c) = (shape[0], shape[1], shape[2]);
|
||||
if c != 3 {
|
||||
return Err(PyValueError::new_err(format!(
|
||||
"expected HWC RGB array with 3 channels, got {c}"
|
||||
)));
|
||||
}
|
||||
let hash = common::content_hash_u64(raw_bytes);
|
||||
let rgb = arr
|
||||
.as_slice()
|
||||
.map_err(|_| PyValueError::new_err("array must be C-contiguous"))?
|
||||
.to_vec();
|
||||
let (oh, ow, out) = py.detach(move || {
|
||||
let (tw, th) = common::resize::scaled_dims(w, h, rescale_frac, rescale_cap);
|
||||
let (rgb, h, w) = if (tw, th) != (w, h) {
|
||||
(
|
||||
@@ -282,18 +268,21 @@ fn rescale_patchify_hash<'py>(
|
||||
(rgb, h, w)
|
||||
};
|
||||
(h, w, patchify_alloc(&rgb, h, w, patch_size))
|
||||
})
|
||||
});
|
||||
Ok((oh, ow, out.into_pyarray(py), hash))
|
||||
});
|
||||
Ok((oh, ow, out.into_pyarray(py), hash))
|
||||
}
|
||||
|
||||
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let m = PyModule::new(parent.py(), "inkling")?;
|
||||
m.add_function(wrap_pyfunction!(patchify_rgb, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(decode_patchify, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(decode_patchify_batch, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(preprocess_images, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(rescale_patchify_hash, &m)?)?;
|
||||
parent.add_submodule(&m)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let m = PyModule::new(parent.py(), "inkling")?;
|
||||
m.add_function(wrap_pyfunction!(patchify_rgb, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(decode_patchify, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(decode_patchify_batch, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(preprocess_images, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(rescale_patchify_hash, &m)?)?;
|
||||
parent.add_submodule(&m)?;
|
||||
Ok(())
|
||||
}
|
||||
#[cfg(feature = "python")]
|
||||
pub use python::register;
|
||||
|
||||
@@ -1,12 +1,26 @@
|
||||
mod common;
|
||||
mod inkling;
|
||||
//! sglang-mm: Rust-accelerated multimodal preprocessing for SGLang.
|
||||
//!
|
||||
//! Built two ways:
|
||||
//! * PyO3 extension `sglang.srt.multimodal._core` (feature `python`, default),
|
||||
//! used by Python processors (e.g. Inkling) and by parity tests.
|
||||
//! * Pure-Rust `rlib` (`default-features = false`), linked by `sglang-server`'s
|
||||
//! MM worker path — no pyo3 in that dependency graph.
|
||||
|
||||
pub mod common;
|
||||
pub mod driver;
|
||||
pub mod inkling;
|
||||
pub mod pipeline;
|
||||
pub mod qwen_vl;
|
||||
pub mod registry;
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
use pyo3::prelude::*;
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
#[pymodule]
|
||||
fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
common::register(m)?;
|
||||
inkling::register(m)?;
|
||||
qwen_vl::register(m)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
@@ -0,0 +1,132 @@
|
||||
//! The model-family seam of the server MM pipeline.
|
||||
//!
|
||||
//! Design rule: **families produce data, the driver owns control flow.** A
|
||||
//! family never sees the request loop, the thread pool, or the failure
|
||||
//! protocol — it implements [`MmFamilyProcessor`], turning decoded media into
|
||||
//! named tensors and describing its prompt geometry as a [`TokenLayout`]
|
||||
//! value. `driver::process` applies the layout mechanically, so expansion,
|
||||
//! per-item offsets, and position inputs all derive from one declarative
|
||||
//! structure and every family gets identical failure semantics for free.
|
||||
|
||||
/// Typed tensor payload. Grows a variant per dtype actually produced by a
|
||||
/// family — not speculatively.
|
||||
pub enum TensorData {
|
||||
F32(Vec<f32>),
|
||||
I64(Vec<i64>),
|
||||
}
|
||||
|
||||
pub struct Tensor {
|
||||
pub shape: Vec<usize>,
|
||||
pub data: TensorData,
|
||||
}
|
||||
|
||||
/// Named auxiliary tensors that reach the model runner as kwargs — the Rust
|
||||
/// analogue of Python's `MultimodalDataItem.model_specific_data`, e.g. qwen's
|
||||
/// `image_grid_thw`.
|
||||
pub type NamedTensors = Vec<(String, Tensor)>;
|
||||
|
||||
/// One decoded media item handed to [`MmFamilyProcessor::process_item`].
|
||||
/// Grows a variant per modality as families that need it are ported.
|
||||
pub enum DecodedMedia {
|
||||
/// HWC u8 RGB.
|
||||
Image {
|
||||
rgb: Vec<u8>,
|
||||
height: usize,
|
||||
width: usize,
|
||||
},
|
||||
}
|
||||
|
||||
/// Family-internal geometry of one processed item, consumed by
|
||||
/// [`MmFamilyProcessor::layout`] / [`MmFamilyProcessor::positions`]. Grows a
|
||||
/// variant per family style; the driver never interprets it.
|
||||
#[derive(Clone, Debug)]
|
||||
pub enum Geometry {
|
||||
/// `[t, h, w]` patch grid (`t` = 1 for still images).
|
||||
Grid([u32; 3]),
|
||||
}
|
||||
|
||||
/// One processed media item, mirroring Python's `MultimodalDataItem`: the
|
||||
/// primary feature tensor, named auxiliary tensors, and the geometry the
|
||||
/// family's own `layout`/`positions` hooks need.
|
||||
pub struct ProcessedItem {
|
||||
/// The model's feature tensor for this item (qwen: `pixel_values`).
|
||||
pub feature: Tensor,
|
||||
pub aux: NamedTensors,
|
||||
pub geometry: Geometry,
|
||||
}
|
||||
|
||||
/// The tokens one media item occupies in the expanded prompt.
|
||||
pub enum TokenPattern {
|
||||
/// N copies of one placeholder id (qwen-style).
|
||||
Repeat { id: i32, n: usize },
|
||||
/// An explicit id sequence — tile markers, row separators, wrapper
|
||||
/// tokens (minicpm/internvl-style structured expansions).
|
||||
Explicit(Vec<i32>),
|
||||
}
|
||||
|
||||
/// One span of the expanded prompt.
|
||||
pub enum Segment {
|
||||
/// Copy `src` (a range into the original ids) verbatim.
|
||||
Text(std::ops::Range<usize>),
|
||||
/// Media item `item`'s token span.
|
||||
Media { item: usize, pattern: TokenPattern },
|
||||
}
|
||||
|
||||
/// Prompt geometry as data: the family *describes* the expansion, the driver
|
||||
/// *applies* it (`common::token_layout::apply_layout`) — deriving final input ids
|
||||
/// and per-item offsets, and validating that every item is placed exactly
|
||||
/// once.
|
||||
pub struct TokenLayout {
|
||||
pub segments: Vec<Segment>,
|
||||
}
|
||||
|
||||
/// Modalities a family accepts; the server's message layer rejects anything
|
||||
/// a family does not declare.
|
||||
#[derive(Clone, Copy, Debug, Default)]
|
||||
pub struct Capabilities {
|
||||
pub video: bool,
|
||||
pub audio: bool,
|
||||
}
|
||||
|
||||
/// Position scheme of the expanded prompt.
|
||||
pub enum PositionOutput {
|
||||
/// Plain sequential positions — the scheduler needs nothing extra.
|
||||
Rope1D,
|
||||
/// M-RoPE: flattened row-major `[3, input_len]` positions + the position
|
||||
/// delta (`max + 1 - input_len`).
|
||||
MRope { positions: Vec<i64>, delta: i64 },
|
||||
}
|
||||
|
||||
/// The per-model-family hooks of the server pipeline. Adding a family =
|
||||
/// implementing this in `src/<model>/mod.rs` and adding its `family` arm to
|
||||
/// [`crate::registry::pipeline_from_spec`]. All parameters come from the
|
||||
/// runtime spec JSON (resolved from the HF config on the Python side);
|
||||
/// nothing is hardcoded per model.
|
||||
pub trait MmFamilyProcessor: Send + Sync {
|
||||
/// Modalities beyond images this family accepts. Default: images only.
|
||||
fn capabilities(&self) -> Capabilities {
|
||||
Capabilities::default()
|
||||
}
|
||||
|
||||
/// Preprocess one decoded media item: the model's HF processor
|
||||
/// equivalent (resize/tile/normalize/patchify → named tensors) plus the
|
||||
/// geometry `layout`/`positions` will need.
|
||||
fn process_item(&self, media: &DecodedMedia) -> Result<ProcessedItem, String>;
|
||||
|
||||
/// Describe how the prompt expands around the processed items (in
|
||||
/// prompt order). Sees the full original prompt and all items, so
|
||||
/// structured schemes (tile markers, separators) are expressible.
|
||||
fn layout(&self, input_ids: &[i32], items: &[Geometry]) -> Result<TokenLayout, String>;
|
||||
|
||||
/// Positions for the expanded prompt. Families without a custom scheme
|
||||
/// keep the default.
|
||||
fn positions(
|
||||
&self,
|
||||
input_len: usize,
|
||||
offsets: &[(u32, u32)],
|
||||
items: &[Geometry],
|
||||
) -> Result<PositionOutput, String> {
|
||||
let _ = (input_len, offsets, items);
|
||||
Ok(PositionOutput::Rope1D)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,645 @@
|
||||
//! Qwen VL family (Qwen2-VL / 2.5-VL / 3-VL / 3.5) server-pipeline image processor.
|
||||
//!
|
||||
//! Pure-Rust equivalent of the HF `Qwen2VLImageProcessor` pipeline the Python
|
||||
//! `QwenVLImageProcessor` drives: `smart_resize` → bicubic resize → rescale +
|
||||
//! normalize → patchify into `[grid_h*grid_w, C*tps*ps*ps]` (HF flatten order:
|
||||
//! patches by `(gh/m, gw/m, m, m)`, features by `(C, tps, ps, ps)`, temporal
|
||||
//! copies duplicated for stills) — plus the image-only M-RoPE fast path.
|
||||
//! All parameters come from the runtime spec; nothing is hardcoded per model.
|
||||
|
||||
use crate::common::{par, resize, token_layout};
|
||||
use crate::pipeline::{
|
||||
DecodedMedia, Geometry, MmFamilyProcessor, PositionOutput, ProcessedItem, Tensor, TensorData,
|
||||
TokenLayout,
|
||||
};
|
||||
|
||||
const MAX_RATIO: f64 = 200.0;
|
||||
|
||||
/// One media item's placement for M-RoPE: inclusive token range + patch grid.
|
||||
pub struct MropeItem {
|
||||
pub start: u32,
|
||||
pub end: u32,
|
||||
pub grid: [u32; 3],
|
||||
}
|
||||
|
||||
/// Resolved processor params, deserialized from the Python-side spec JSON
|
||||
/// (unknown fields like `family` are ignored here).
|
||||
#[derive(Clone, Debug, serde::Deserialize)]
|
||||
pub struct QwenVlSpec {
|
||||
pub image_token_id: i32,
|
||||
pub patch_size: usize,
|
||||
pub merge_size: usize,
|
||||
pub temporal_patch_size: usize,
|
||||
pub min_pixels: usize,
|
||||
pub max_pixels: usize,
|
||||
pub image_mean: [f32; 3],
|
||||
pub image_std: [f32; 3],
|
||||
}
|
||||
|
||||
pub struct QwenVlProcessor {
|
||||
spec: QwenVlSpec,
|
||||
/// Per-channel u8 → normalized-f32 lookup: `(v/255 - mean) / std`.
|
||||
lut: [[f32; 256]; 3],
|
||||
}
|
||||
|
||||
impl QwenVlProcessor {
|
||||
pub fn new(spec: QwenVlSpec) -> Result<Self, String> {
|
||||
if spec.patch_size == 0 || spec.merge_size == 0 || spec.temporal_patch_size == 0 {
|
||||
return Err("qwen_vl spec: sizes must be positive".into());
|
||||
}
|
||||
let lut = core::array::from_fn(|c| {
|
||||
core::array::from_fn(|v| (v as f32 / 255.0 - spec.image_mean[c]) / spec.image_std[c])
|
||||
});
|
||||
Ok(Self { spec, lut })
|
||||
}
|
||||
|
||||
pub fn from_spec_json(json: &str) -> Result<Self, String> {
|
||||
let spec: QwenVlSpec =
|
||||
serde_json::from_str(json).map_err(|e| format!("qwen_vl spec: {e}"))?;
|
||||
Self::new(spec)
|
||||
}
|
||||
|
||||
fn factor(&self) -> usize {
|
||||
self.spec.patch_size * self.spec.merge_size
|
||||
}
|
||||
|
||||
/// HF flatten: patches ordered `(gh/m, gw/m, m, m)`, features `(C, tps,
|
||||
/// ps, ps)`; parallel over merged-block rows.
|
||||
fn patchify(&self, rgb: &[u8], h: usize, w: usize) -> Vec<f32> {
|
||||
let (ps, m, tps) = (
|
||||
self.spec.patch_size,
|
||||
self.spec.merge_size,
|
||||
self.spec.temporal_patch_size,
|
||||
);
|
||||
let (gh, gw) = (h / ps, w / ps);
|
||||
let dim = 3 * tps * ps * ps;
|
||||
let block_row = gw * m * dim; // one merged-block row of patches
|
||||
let mut out = vec![0.0f32; gh * gw * dim];
|
||||
|
||||
par::for_chunks_mut(&mut out, block_row, |i, chunk| {
|
||||
let mut p = 0;
|
||||
for j in 0..gw / m {
|
||||
for mh in 0..m {
|
||||
for mw in 0..m {
|
||||
let y0 = (i * m + mh) * ps;
|
||||
let x0 = (j * m + mw) * ps;
|
||||
let patch = &mut chunk[p * dim..(p + 1) * dim];
|
||||
for c in 0..3 {
|
||||
let ch = &mut patch[c * tps * ps * ps..];
|
||||
for py in 0..ps {
|
||||
let src = ((y0 + py) * w + x0) * 3 + c;
|
||||
for px in 0..ps {
|
||||
ch[py * ps + px] = self.lut[c][rgb[src + px * 3] as usize];
|
||||
}
|
||||
}
|
||||
// Temporal copies of a still are duplicates.
|
||||
let (t0, rest) = ch.split_at_mut(ps * ps);
|
||||
for t in 0..tps - 1 {
|
||||
rest[t * ps * ps..(t + 1) * ps * ps].copy_from_slice(t0);
|
||||
}
|
||||
}
|
||||
p += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
impl QwenVlProcessor {
|
||||
fn tokens_per_image(&self, grid: &[u32; 3]) -> usize {
|
||||
(grid[0] as usize * grid[1] as usize * grid[2] as usize)
|
||||
/ (self.spec.merge_size * self.spec.merge_size)
|
||||
}
|
||||
}
|
||||
|
||||
impl MmFamilyProcessor for QwenVlProcessor {
|
||||
fn process_item(&self, media: &DecodedMedia) -> Result<ProcessedItem, String> {
|
||||
let DecodedMedia::Image { rgb, height, width } = media;
|
||||
let (h, w) = (*height, *width);
|
||||
let (th, tw) = smart_resize(
|
||||
h,
|
||||
w,
|
||||
self.factor(),
|
||||
self.spec.min_pixels,
|
||||
self.spec.max_pixels,
|
||||
)?;
|
||||
let resized;
|
||||
let data = if (th, tw) != (h, w) {
|
||||
resized = resize::resize_rgb_filter(rgb, h, w, th, tw, resize::Filter::Bicubic);
|
||||
&resized
|
||||
} else {
|
||||
rgb.as_slice()
|
||||
};
|
||||
let (gh, gw) = (th / self.spec.patch_size, tw / self.spec.patch_size);
|
||||
// `smart_resize` guarantees both: dims are positive and divisible by
|
||||
// `patch_size * merge_size`. `patchify` indexes on that (and the `dim`
|
||||
// division below needs a non-empty grid), so fail loudly rather than
|
||||
// panic if a future spec change breaks the guarantee.
|
||||
if gh == 0 || gw == 0 || gh % self.spec.merge_size != 0 || gw % self.spec.merge_size != 0 {
|
||||
return Err(format!(
|
||||
"qwen_vl: patch grid {gh}x{gw} is empty or not a multiple of \
|
||||
merge_size {}",
|
||||
self.spec.merge_size
|
||||
));
|
||||
}
|
||||
let pixel_values = self.patchify(data, th, tw);
|
||||
let dim = pixel_values.len() / (gh * gw);
|
||||
Ok(ProcessedItem {
|
||||
feature: Tensor {
|
||||
shape: vec![gh * gw, dim],
|
||||
data: TensorData::F32(pixel_values),
|
||||
},
|
||||
aux: vec![(
|
||||
"image_grid_thw".to_string(),
|
||||
Tensor {
|
||||
shape: vec![3],
|
||||
data: TensorData::I64(vec![1, gh as i64, gw as i64]),
|
||||
},
|
||||
)],
|
||||
geometry: Geometry::Grid([1, gh as u32, gw as u32]),
|
||||
})
|
||||
}
|
||||
|
||||
fn layout(&self, input_ids: &[i32], items: &[Geometry]) -> Result<TokenLayout, String> {
|
||||
let counts = items
|
||||
.iter()
|
||||
.map(|Geometry::Grid(grid)| self.tokens_per_image(grid))
|
||||
.collect::<Vec<_>>();
|
||||
token_layout::layout_by_placeholder(input_ids, self.spec.image_token_id, &counts)
|
||||
}
|
||||
|
||||
fn positions(
|
||||
&self,
|
||||
input_len: usize,
|
||||
offsets: &[(u32, u32)],
|
||||
items: &[Geometry],
|
||||
) -> Result<PositionOutput, String> {
|
||||
let mrope_items = offsets
|
||||
.iter()
|
||||
.zip(items)
|
||||
.map(|(&(start, end), Geometry::Grid(grid))| MropeItem {
|
||||
start,
|
||||
end,
|
||||
grid: *grid,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
let (positions, delta) = mrope_image_only(input_len, &mrope_items, self.spec.merge_size)?;
|
||||
Ok(PositionOutput::MRope { positions, delta })
|
||||
}
|
||||
}
|
||||
|
||||
/// Python-`round()` (round-half-to-even), which `round_by_factor` relies on.
|
||||
fn round_half_even(x: f64) -> f64 {
|
||||
if (x - x.trunc()).abs() == 0.5 {
|
||||
(x / 2.0).round() * 2.0
|
||||
} else {
|
||||
x.round()
|
||||
}
|
||||
}
|
||||
|
||||
/// The Qwen `smart_resize`: dims divisible by `factor`, total pixels within
|
||||
/// `[min_pixels, max_pixels]`, aspect ratio preserved as closely as possible.
|
||||
pub fn smart_resize(
|
||||
height: usize,
|
||||
width: usize,
|
||||
factor: usize,
|
||||
min_pixels: usize,
|
||||
max_pixels: usize,
|
||||
) -> Result<(usize, usize), String> {
|
||||
let (h, w) = (height as f64, width as f64);
|
||||
if height == 0 || width == 0 {
|
||||
return Err("empty image".into());
|
||||
}
|
||||
let ratio = h.max(w) / h.min(w);
|
||||
if ratio > MAX_RATIO {
|
||||
return Err(format!(
|
||||
"absolute aspect ratio must be smaller than {MAX_RATIO}, got {ratio}"
|
||||
));
|
||||
}
|
||||
let f = factor as f64;
|
||||
let mut h_bar = ((round_half_even(h / f) * f) as usize).max(factor);
|
||||
let mut w_bar = ((round_half_even(w / f) * f) as usize).max(factor);
|
||||
if h_bar * w_bar > max_pixels {
|
||||
let beta = (h * w / max_pixels as f64).sqrt();
|
||||
h_bar = ((h / beta / f).floor() * f) as usize;
|
||||
w_bar = ((w / beta / f).floor() * f) as usize;
|
||||
} else if h_bar * w_bar < min_pixels {
|
||||
let beta = (min_pixels as f64 / (h * w)).sqrt();
|
||||
h_bar = ((h * beta / f).ceil() * f) as usize;
|
||||
w_bar = ((w * beta / f).ceil() * f) as usize;
|
||||
}
|
||||
// The downscale branch floors without a lower clamp (as Python does), so a
|
||||
// very thin image against a small `max_pixels` can floor a side to 0.
|
||||
// Python then fails inside PIL's resize; here it would reach the resize
|
||||
// coefficient math (overflow panic in debug, garbage in release) and the
|
||||
// `dim = len / (gh * gw)` division, so reject it as a request error.
|
||||
if h_bar == 0 || w_bar == 0 {
|
||||
return Err(format!(
|
||||
"smart_resize: {height}x{width} degenerates to {h_bar}x{w_bar} at \
|
||||
max_pixels={max_pixels}; image is too thin for this pixel budget"
|
||||
));
|
||||
}
|
||||
Ok((h_bar, w_bar))
|
||||
}
|
||||
|
||||
/// Image-only M-RoPE fast path (the image branch of
|
||||
/// `MRotaryEmbedding.get_rope_index`, identical across Qwen generations):
|
||||
/// text runs sequentially on all three rows; each image spans `(t, h/m, w/m)`
|
||||
/// index grids; positions advance by `max(t, h/m, w/m)` past an image.
|
||||
/// Returns flattened row-major `[3, input_len]` positions and the delta
|
||||
/// (`max + 1 - input_len`). `items` must be in prompt order.
|
||||
pub fn mrope_image_only(
|
||||
input_len: usize,
|
||||
items: &[MropeItem],
|
||||
merge_size: usize,
|
||||
) -> Result<(Vec<i64>, i64), String> {
|
||||
let len = input_len;
|
||||
let mut pos = vec![0i64; 3 * len];
|
||||
let fill_text = |st: usize, n: usize, base: i64, pos: &mut [i64]| {
|
||||
for k in 0..n {
|
||||
let v = base + k as i64;
|
||||
pos[st + k] = v;
|
||||
pos[len + st + k] = v;
|
||||
pos[2 * len + st + k] = v;
|
||||
}
|
||||
};
|
||||
let mut st = 0usize;
|
||||
let mut next_pos = 0i64;
|
||||
for item in items {
|
||||
let (start, end) = (item.start as usize, item.end as usize);
|
||||
if start < st || end >= len {
|
||||
return Err(format!(
|
||||
"mrope: item range ({start},{end}) out of order/bounds"
|
||||
));
|
||||
}
|
||||
fill_text(st, start - st, next_pos, &mut pos);
|
||||
next_pos += (start - st) as i64;
|
||||
|
||||
let t = item.grid[0] as usize;
|
||||
let gh = item.grid[1] as usize / merge_size;
|
||||
let gw = item.grid[2] as usize / merge_size;
|
||||
if t * gh * gw != end - start + 1 {
|
||||
return Err("mrope: token span does not match grid".into());
|
||||
}
|
||||
for ti in 0..t {
|
||||
for hi in 0..gh {
|
||||
for wi in 0..gw {
|
||||
let idx = start + (ti * gh + hi) * gw + wi;
|
||||
pos[idx] = next_pos + ti as i64;
|
||||
pos[len + idx] = next_pos + hi as i64;
|
||||
pos[2 * len + idx] = next_pos + wi as i64;
|
||||
}
|
||||
}
|
||||
}
|
||||
next_pos += (t.max(gh).max(gw)) as i64;
|
||||
st = end + 1;
|
||||
}
|
||||
if st < len {
|
||||
fill_text(st, len - st, next_pos, &mut pos);
|
||||
}
|
||||
let max = pos.iter().copied().max().unwrap_or(-1);
|
||||
Ok((pos, max + 1 - len as i64))
|
||||
}
|
||||
|
||||
/// The qwen scheduler-drain shape, extracted from the generic driver
|
||||
/// [`Output`](crate::driver::Output). Shared by `sglang-server`'s MM worker
|
||||
/// and the parity binding so the mapping can't drift; replaced by a generic
|
||||
/// named-tensor handoff once a second family needs a different shape.
|
||||
pub struct QwenDrain {
|
||||
pub input_ids: Vec<i32>,
|
||||
/// All items' `pixel_values`, concatenated in prompt order.
|
||||
pub features: Vec<f32>,
|
||||
pub grids: Vec<[u32; 3]>,
|
||||
pub hashes: Vec<u64>,
|
||||
pub offsets: Vec<(u32, u32)>,
|
||||
pub mrope: Vec<i64>,
|
||||
pub mrope_delta: i64,
|
||||
}
|
||||
|
||||
pub fn pack_drain(output: crate::driver::Output) -> Result<QwenDrain, String> {
|
||||
use crate::pipeline::PositionOutput;
|
||||
|
||||
let PositionOutput::MRope { positions, delta } = output.positions else {
|
||||
return Err("qwen_vl drain: expected M-RoPE positions".into());
|
||||
};
|
||||
let mut features = Vec::new();
|
||||
let mut grids = Vec::with_capacity(output.items.len());
|
||||
let mut hashes = Vec::with_capacity(output.items.len());
|
||||
for item in output.items {
|
||||
let TensorData::F32(pixel_values) = item.feature.data else {
|
||||
return Err("qwen_vl drain: expected f32 feature".into());
|
||||
};
|
||||
features.extend(pixel_values);
|
||||
let grid = item
|
||||
.aux
|
||||
.into_iter()
|
||||
.find_map(|(name, tensor)| match (name.as_str(), tensor.data) {
|
||||
("image_grid_thw", TensorData::I64(v)) => Some(v),
|
||||
_ => None,
|
||||
})
|
||||
.ok_or("qwen_vl drain: missing image_grid_thw")?;
|
||||
grids.push([grid[0] as u32, grid[1] as u32, grid[2] as u32]);
|
||||
hashes.push(item.hash);
|
||||
}
|
||||
Ok(QwenDrain {
|
||||
input_ids: output.input_ids,
|
||||
features,
|
||||
grids,
|
||||
hashes,
|
||||
offsets: output.offsets,
|
||||
mrope: positions,
|
||||
mrope_delta: delta,
|
||||
})
|
||||
}
|
||||
|
||||
// --- Python bindings (parity tests drive the exact server pipeline) ---
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
mod python {
|
||||
use numpy::{IntoPyArray, PyArray1};
|
||||
use pyo3::exceptions::PyValueError;
|
||||
use pyo3::prelude::*;
|
||||
|
||||
use super::*;
|
||||
use crate::pipeline::TensorData;
|
||||
|
||||
/// `(pixel_values flat f32, (t, h, w))` for one preprocessed image.
|
||||
type PyProcessedImage<'py> = (Bound<'py, PyArray1<f32>>, (u32, u32, u32));
|
||||
/// Full native pipeline output at the scheduler boundary:
|
||||
/// `(input_ids, features, grids, hashes, offsets, mrope, mrope_delta)`.
|
||||
type PyNativeOutput<'py> = (
|
||||
Vec<i32>,
|
||||
Bound<'py, PyArray1<f32>>,
|
||||
Vec<(u32, u32, u32)>,
|
||||
Vec<u64>,
|
||||
Vec<(u32, u32)>,
|
||||
Bound<'py, PyArray1<i64>>,
|
||||
i64,
|
||||
);
|
||||
|
||||
/// Run the full native image path on encoded image bytes:
|
||||
/// decode → smart_resize → bicubic → normalize → patchify.
|
||||
/// Returns `(pixel_values flat f32, (t, h, w))`.
|
||||
#[pyfunction]
|
||||
fn preprocess<'py>(
|
||||
py: Python<'py>,
|
||||
data: Vec<u8>,
|
||||
spec_json: &str,
|
||||
) -> PyResult<PyProcessedImage<'py>> {
|
||||
let proc = QwenVlProcessor::from_spec_json(spec_json).map_err(PyValueError::new_err)?;
|
||||
let out = py
|
||||
.detach(move || {
|
||||
let (rgb, height, width) = crate::common::decode_rgb(&data)?;
|
||||
proc.process_item(&DecodedMedia::Image { rgb, height, width })
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
let Geometry::Grid([t, h, w]) = out.geometry;
|
||||
let TensorData::F32(pixel_values) = out.feature.data else {
|
||||
return Err(PyValueError::new_err("qwen_vl: expected f32 feature"));
|
||||
};
|
||||
Ok((pixel_values.into_pyarray(py), (t, h, w)))
|
||||
}
|
||||
|
||||
#[pyfunction]
|
||||
fn smart_resize_py(
|
||||
height: usize,
|
||||
width: usize,
|
||||
factor: usize,
|
||||
min_pixels: usize,
|
||||
max_pixels: usize,
|
||||
) -> PyResult<(usize, usize)> {
|
||||
smart_resize(height, width, factor, min_pixels, max_pixels).map_err(PyValueError::new_err)
|
||||
}
|
||||
|
||||
/// `(positions flat [3*input_len], delta)` for image-only requests;
|
||||
/// `items` = [(start, end_inclusive, t, h, w), ...] in prompt order.
|
||||
#[pyfunction]
|
||||
fn mrope_image_only_py<'py>(
|
||||
py: Python<'py>,
|
||||
input_len: usize,
|
||||
items: Vec<(u32, u32, u32, u32, u32)>,
|
||||
merge_size: usize,
|
||||
) -> PyResult<(Bound<'py, PyArray1<i64>>, i64)> {
|
||||
let items: Vec<MropeItem> = items
|
||||
.into_iter()
|
||||
.map(|(start, end, t, h, w)| MropeItem {
|
||||
start,
|
||||
end,
|
||||
grid: [t, h, w],
|
||||
})
|
||||
.collect();
|
||||
let (pos, delta) =
|
||||
mrope_image_only(input_len, &items, merge_size).map_err(PyValueError::new_err)?;
|
||||
Ok((pos.into_pyarray(py), delta))
|
||||
}
|
||||
|
||||
/// One image source: a `str` (data:/base64/file/http, resolved by
|
||||
/// `common::fetch`) or raw encoded `bytes`.
|
||||
#[derive(FromPyObject)]
|
||||
enum PyImageSource {
|
||||
Str(String),
|
||||
Bytes(Vec<u8>),
|
||||
}
|
||||
|
||||
/// Drive the same typed native Qwen request pipeline used by
|
||||
/// `sglang-server` (whose message layer owns the wire-payload parsing).
|
||||
#[pyfunction]
|
||||
#[pyo3(signature = (input_ids, images, spec_json))]
|
||||
fn process_native_mm<'py>(
|
||||
py: Python<'py>,
|
||||
input_ids: Option<Vec<i32>>,
|
||||
images: Vec<PyImageSource>,
|
||||
spec_json: String,
|
||||
) -> PyResult<PyNativeOutput<'py>> {
|
||||
let images = images
|
||||
.into_iter()
|
||||
.map(|source| match source {
|
||||
PyImageSource::Str(s) => crate::driver::ImageSource::String(s),
|
||||
PyImageSource::Bytes(b) => crate::driver::ImageSource::Bytes(b),
|
||||
})
|
||||
.collect();
|
||||
let input = crate::driver::MmInput {
|
||||
text: None,
|
||||
input_ids,
|
||||
images,
|
||||
};
|
||||
let drain = py
|
||||
.detach(move || {
|
||||
let family = crate::registry::pipeline_from_spec(&spec_json)?;
|
||||
let output = crate::driver::process(family.as_ref(), input, |_| {
|
||||
Err("native parity API requires input_ids".into())
|
||||
})?;
|
||||
pack_drain(output)
|
||||
})
|
||||
.map_err(PyValueError::new_err)?;
|
||||
Ok((
|
||||
drain.input_ids,
|
||||
drain.features.into_pyarray(py),
|
||||
drain.grids.into_iter().map(|[t, h, w]| (t, h, w)).collect(),
|
||||
drain.hashes,
|
||||
drain.offsets,
|
||||
drain.mrope.into_pyarray(py),
|
||||
drain.mrope_delta,
|
||||
))
|
||||
}
|
||||
|
||||
pub fn register(parent: &Bound<'_, PyModule>) -> PyResult<()> {
|
||||
let m = PyModule::new(parent.py(), "qwen_vl")?;
|
||||
m.add_function(wrap_pyfunction!(preprocess, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(smart_resize_py, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(mrope_image_only_py, &m)?)?;
|
||||
m.add_function(wrap_pyfunction!(process_native_mm, &m)?)?;
|
||||
parent.add_submodule(&m)?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(feature = "python")]
|
||||
pub use python::register;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn spec() -> QwenVlSpec {
|
||||
QwenVlSpec {
|
||||
image_token_id: 1,
|
||||
patch_size: 2,
|
||||
merge_size: 2,
|
||||
temporal_patch_size: 2,
|
||||
min_pixels: 4,
|
||||
max_pixels: 1 << 30,
|
||||
image_mean: [0.0; 3],
|
||||
image_std: [1.0; 3],
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn smart_resize_matches_python_reference() {
|
||||
// Values from the Python `smart_resize` (qwen_vl.py) run offline.
|
||||
assert_eq!(
|
||||
smart_resize(1365, 2048, 28, 3136, 12845056).unwrap(),
|
||||
(1372, 2044)
|
||||
);
|
||||
assert_eq!(
|
||||
smart_resize(100, 100, 28, 3136, 12845056).unwrap(),
|
||||
(112, 112)
|
||||
);
|
||||
// Downscale branch: 4000x3000 exceeds 1280*28*28 → floor_by_factor.
|
||||
assert_eq!(
|
||||
smart_resize(3000, 4000, 28, 3136, 1003520).unwrap(),
|
||||
(840, 1148)
|
||||
);
|
||||
// Upscale branch: tiny image below min_pixels → ceil_by_factor.
|
||||
assert_eq!(smart_resize(20, 20, 28, 3136, 12845056).unwrap(), (56, 56));
|
||||
// Qwen3.5 factors (patch 16 * merge 2, min 65536, max 16777216).
|
||||
assert_eq!(
|
||||
smart_resize(1365, 2048, 32, 65536, 16777216).unwrap(),
|
||||
(1376, 2048)
|
||||
);
|
||||
// Banker's rounding tie: 48/32 = 1.5 rounds to 2 (even), not 1.
|
||||
assert_eq!(smart_resize(4000, 48, 32, 4, 1 << 30).unwrap(), (4000, 64));
|
||||
// Extreme aspect ratio rejected.
|
||||
assert!(smart_resize(10000, 10, 28, 3136, 12845056).is_err());
|
||||
}
|
||||
|
||||
/// A thin image against a small `max_pixels` floors one side to 0. That
|
||||
/// used to reach the resize coefficient math and panic on a worker thread
|
||||
/// (`attempt to multiply with overflow`) instead of rejecting the request.
|
||||
#[test]
|
||||
fn degenerate_target_is_rejected_not_panicked() {
|
||||
// Aspect ratio 200 is exactly at MAX_RATIO, so it passes that guard;
|
||||
// 10 / beta then floors to 0 with factor 28.
|
||||
assert!(smart_resize(10, 2000, 28, 3136, 3136).is_err());
|
||||
|
||||
let mut spec = spec();
|
||||
spec.patch_size = 14;
|
||||
spec.min_pixels = 3136;
|
||||
spec.max_pixels = 3136;
|
||||
let proc = QwenVlProcessor::new(spec).unwrap();
|
||||
let err = proc
|
||||
.process_item(&DecodedMedia::Image {
|
||||
rgb: vec![0u8; 10 * 2000 * 3],
|
||||
height: 10,
|
||||
width: 2000,
|
||||
})
|
||||
.err()
|
||||
.expect("degenerate geometry must be an Err, never a panic");
|
||||
assert!(err.contains("smart_resize"), "unexpected error: {err}");
|
||||
}
|
||||
|
||||
/// The server's message layer gates modalities on what a family declares,
|
||||
/// so a family gaining video/audio support must not silently inherit the
|
||||
/// images-only default.
|
||||
#[test]
|
||||
fn qwen_declares_images_only() {
|
||||
let caps = QwenVlProcessor::new(spec()).unwrap().capabilities();
|
||||
assert!(!caps.video && !caps.audio);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patchify_layout_matches_hf_order() {
|
||||
// 4x8 image, ps=2, m=2, tps=2 → gh=2, gw=4, dim=3*2*2*2=24.
|
||||
// Pixel value encodes its (y, x): v = y*16 + x*2 (fits u8).
|
||||
let (h, w) = (4usize, 8usize);
|
||||
let mut rgb = vec![0u8; h * w * 3];
|
||||
for y in 0..h {
|
||||
for x in 0..w {
|
||||
for c in 0..3 {
|
||||
rgb[(y * w + x) * 3 + c] = (y * 16 + x * 2 + c) as u8;
|
||||
}
|
||||
}
|
||||
}
|
||||
let proc = QwenVlProcessor::new(spec()).unwrap();
|
||||
let pv = proc.patchify(&rgb, h, w);
|
||||
let dim = 24; // 3 * tps * ps * ps
|
||||
assert_eq!(pv.len(), 2 * 4 * dim);
|
||||
|
||||
// Patch order (gh/m=1, gw/m=2, m, m): patch 0 = block(0,0) offset (0,0),
|
||||
// patch 1 = (0,0)+(0,1) → x0=2, patch 2 = (0,0)+(1,0) → y0=2,
|
||||
// patch 4 = block(0,1) → x0=4.
|
||||
let lut = |y: usize, x: usize, c: usize| ((y * 16 + x * 2 + c) as f32) / 255.0;
|
||||
// patch 1, channel 0, t=0, (py=0, px=0) → pixel (0, 2).
|
||||
assert_eq!(pv[dim], lut(0, 2, 0));
|
||||
// patch 2, channel 0, t=0, (0,0) → pixel (2, 0).
|
||||
assert_eq!(pv[2 * dim], lut(2, 0, 0));
|
||||
// patch 4, channel 0 → pixel (0, 4).
|
||||
assert_eq!(pv[4 * dim], lut(0, 4, 0));
|
||||
// Temporal duplicate: t=1 block equals t=0 block.
|
||||
let ps2 = 4; // ps*ps
|
||||
assert_eq!(pv[dim + ps2], pv[dim]);
|
||||
// Channel 1 block of patch 0 → same pixel, c=1.
|
||||
assert_eq!(pv[2 * ps2], lut(0, 0, 1)); // c stride = tps*ps*ps = 8
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mrope_image_only_matches_reference() {
|
||||
// 3 text tokens, image of grid [1, 4, 6] (m=2 → 2x3 = 6 tokens), 2 text.
|
||||
// input: [T T T I I I I I I T T], len 11.
|
||||
let items = [MropeItem {
|
||||
start: 3,
|
||||
end: 8,
|
||||
grid: [1, 4, 6],
|
||||
}];
|
||||
let (pos, delta) = mrope_image_only(11, &items, 2).unwrap();
|
||||
let len = 11;
|
||||
// Text prefix 0..3: all rows 0,1,2.
|
||||
for k in 0..3 {
|
||||
assert_eq!(
|
||||
(pos[k], pos[len + k], pos[2 * len + k]),
|
||||
(k as i64, k as i64, k as i64)
|
||||
);
|
||||
}
|
||||
// Image tokens: t=0, h in 0..2, w in 0..3, +3 offset.
|
||||
assert_eq!((pos[3], pos[len + 3], pos[2 * len + 3]), (3, 3, 3));
|
||||
assert_eq!((pos[4], pos[len + 4], pos[2 * len + 4]), (3, 3, 4));
|
||||
assert_eq!((pos[6], pos[len + 6], pos[2 * len + 6]), (3, 4, 3));
|
||||
// Text tail resumes at 3 + max(1,2,3) = 6.
|
||||
assert_eq!((pos[9], pos[len + 9], pos[2 * len + 9]), (6, 6, 6));
|
||||
assert_eq!((pos[10], pos[len + 10], pos[2 * len + 10]), (7, 7, 7));
|
||||
// delta = max + 1 - len = 7 + 1 - 11.
|
||||
assert_eq!(delta, -3);
|
||||
}
|
||||
}
|
||||
@@ -1,7 +1,13 @@
|
||||
//! Model processor registry.
|
||||
//! Model processor registries.
|
||||
//!
|
||||
//! Each model implements `ImageProcessorSpec` and registers itself. The Python
|
||||
//! layer looks up a processor by model name at init time.
|
||||
//! Two registries live here:
|
||||
//! * [`ImageProcessorSpec`] / [`ProcessorRegistry`] — the Python-facing batch
|
||||
//! preprocess interface (e.g. Inkling), looked up by name at init time.
|
||||
//! * [`pipeline_from_spec`] — the pure-Rust request pipeline `sglang-server`'s
|
||||
//! MM workers drive. Each model family implements
|
||||
//! [`crate::pipeline::MmFamilyProcessor`] in `src/<model>/mod.rs`; the Python
|
||||
//! side selects one by serializing a spec
|
||||
//! (`{"family": ..., resolved processor params}`).
|
||||
|
||||
/// `(height, width, patches_as_u16_bits, content_hash)` for one image.
|
||||
pub type PreprocessedImage = (usize, usize, Vec<u16>, u64);
|
||||
@@ -59,3 +65,23 @@ pub fn default_registry() -> ProcessorRegistry {
|
||||
reg.register(Box::new(crate::inkling::InklingProcessor));
|
||||
reg
|
||||
}
|
||||
// --- Server (pure-Rust) request pipeline ---
|
||||
|
||||
/// Build a family processor from the Python-side spec JSON. `Err` on an
|
||||
/// unknown family or malformed spec — the caller treats that as "no Rust
|
||||
/// pipeline".
|
||||
pub fn pipeline_from_spec(
|
||||
json: &str,
|
||||
) -> Result<Box<dyn crate::pipeline::MmFamilyProcessor>, String> {
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Header {
|
||||
family: String,
|
||||
}
|
||||
let header: Header = serde_json::from_str(json).map_err(|e| format!("mm spec: {e}"))?;
|
||||
match header.family.as_str() {
|
||||
"qwen_vl" => Ok(Box::new(crate::qwen_vl::QwenVlProcessor::from_spec_json(
|
||||
json,
|
||||
)?)),
|
||||
other => Err(format!("unknown mm family: {other}")),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,7 +6,7 @@ import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench"))
|
||||
from bench_parity import PS, make_photo_like, ref_patchify
|
||||
|
||||
OUT = (
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
//! The pure-Rust `rlib` that `sglang-server` links must not own worker
|
||||
//! threads: the server supplies concurrency across requests and pins its own
|
||||
//! cores, so a library spawning pools behind its back would fight it.
|
||||
//!
|
||||
//! Guarding this from the outside (thread count of the process) rather than by
|
||||
//! inspecting the code, so it stays true no matter how the fan-out seam in
|
||||
//! `common::par` is refactored. Runs only in the default (rayon-less) build;
|
||||
//! under `--features parallel` the pools are expected.
|
||||
|
||||
#![cfg(not(feature = "parallel"))]
|
||||
|
||||
use sglang_mm_core::driver::{ImageSource, MmInput, process};
|
||||
use sglang_mm_core::registry::pipeline_from_spec;
|
||||
|
||||
const SPEC: &str = r#"{"family":"qwen_vl","image_token_id":1,"patch_size":14,
|
||||
"merge_size":2,"temporal_patch_size":2,"min_pixels":3136,
|
||||
"max_pixels":12845056,"image_mean":[0.0,0.0,0.0],"image_std":[1.0,1.0,1.0]}"#;
|
||||
|
||||
fn thread_names() -> Vec<String> {
|
||||
std::fs::read_dir("/proc/self/task")
|
||||
.expect("procfs")
|
||||
.filter_map(|entry| {
|
||||
let comm = entry.ok()?.path().join("comm");
|
||||
Some(std::fs::read_to_string(comm).ok()?.trim().to_string())
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn png(w: u32, h: u32) -> Vec<u8> {
|
||||
let img = image::RgbImage::from_fn(w, h, |x, y| image::Rgb([x as u8, y as u8, 7]));
|
||||
let mut buf = std::io::Cursor::new(Vec::new());
|
||||
img.write_to(&mut buf, image::ImageFormat::Png).unwrap();
|
||||
buf.into_inner()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn processing_a_request_spawns_no_worker_threads() {
|
||||
let before = thread_names().len();
|
||||
let family = pipeline_from_spec(SPEC).unwrap();
|
||||
|
||||
// Two images, so the per-item fan-out seam is exercised, not bypassed.
|
||||
let out = process(
|
||||
family.as_ref(),
|
||||
MmInput {
|
||||
text: None,
|
||||
input_ids: Some(vec![7, 1, 8, 1, 9]),
|
||||
images: vec![
|
||||
ImageSource::Bytes(png(112, 112)),
|
||||
ImageSource::Bytes(png(84, 140)),
|
||||
],
|
||||
},
|
||||
|_| Err("no tokenizer".into()),
|
||||
)
|
||||
.expect("request should succeed");
|
||||
assert_eq!(out.items.len(), 2);
|
||||
|
||||
let after = thread_names();
|
||||
let spawned: Vec<&String> = after.iter().filter(|t| t.starts_with("sgl-mm")).collect();
|
||||
assert!(
|
||||
spawned.is_empty(),
|
||||
"rlib build spawned crate-owned worker threads: {spawned:?}"
|
||||
);
|
||||
assert_eq!(
|
||||
after.len(),
|
||||
before,
|
||||
"rlib build changed the process thread count: {after:?}"
|
||||
);
|
||||
}
|
||||
@@ -4,7 +4,7 @@ import os
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import sglang.srt.multimodal._core.inkling
|
||||
from sglang.srt.multimodal._core import inkling as _rs_inkling
|
||||
|
||||
GOLDEN_DIR = os.environ.get(
|
||||
"INKLING_MM_GOLDEN_DIR",
|
||||
@@ -20,9 +20,7 @@ def bf16_bits_to_f32(bits: np.ndarray) -> np.ndarray:
|
||||
@pytest.mark.parametrize("path", GOLDENS, ids=[os.path.basename(p) for p in GOLDENS])
|
||||
def test_patchify_rgb_bit_exact(path):
|
||||
g = np.load(path)
|
||||
got = sglang.srt.multimodal._core.inkling.patchify_rgb(
|
||||
g["arr"], int(g["patch_size"])
|
||||
)
|
||||
got = _rs_inkling.patchify_rgb(g["arr"], int(g["patch_size"]))
|
||||
np.testing.assert_array_equal(got, g["bits"].reshape(-1))
|
||||
|
||||
|
||||
@@ -30,9 +28,7 @@ def test_patchify_rgb_bit_exact(path):
|
||||
def test_decode_patchify_png_bit_exact(path):
|
||||
g = np.load(path)
|
||||
h_ref, w_ref = g["arr"].shape[:2]
|
||||
h, w, got = sglang.srt.multimodal._core.inkling.decode_patchify(
|
||||
g["png"].tobytes(), int(g["patch_size"])
|
||||
)
|
||||
h, w, got = _rs_inkling.decode_patchify(g["png"].tobytes(), int(g["patch_size"]))
|
||||
assert (h, w) == (h_ref, w_ref)
|
||||
np.testing.assert_array_equal(got, g["bits"].reshape(-1))
|
||||
|
||||
@@ -41,9 +37,7 @@ def test_batch_matches_single():
|
||||
gs = [np.load(p) for p in GOLDENS]
|
||||
data = [g["png"].tobytes() for g in gs]
|
||||
ps = int(gs[0]["patch_size"])
|
||||
for (h, w, bits), g in zip(
|
||||
sglang.srt.multimodal._core.inkling.decode_patchify_batch(data, ps), gs
|
||||
):
|
||||
for (h, w, bits), g in zip(_rs_inkling.decode_patchify_batch(data, ps), gs):
|
||||
np.testing.assert_array_equal(bits, g["bits"].reshape(-1))
|
||||
|
||||
|
||||
|
||||
@@ -11,11 +11,15 @@ import soundfile as sf
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench"))
|
||||
from bench_parity import make_photo_like
|
||||
|
||||
from sglang.srt.managers.mm_utils import data_hash, hash_feature
|
||||
from sglang.srt.multimodal._core import common as _rs_common
|
||||
from sglang.srt.multimodal.inkling import InklingProcessor
|
||||
from sglang.srt.multimodal.inkling.image_processing_rust import (
|
||||
InklingRustImageProcessor,
|
||||
)
|
||||
from sglang.srt.multimodal.processors import inkling as prc
|
||||
|
||||
|
||||
@@ -39,7 +43,9 @@ def make_proc():
|
||||
proc.IMAGE_TOKEN_ID = 100
|
||||
proc.AUDIO_TOKEN_ID = 101
|
||||
proc.AUDIO_END_TOKEN_ID = 102
|
||||
proc.inkling_processor = InklingProcessor()
|
||||
proc.inkling_processor = InklingProcessor(
|
||||
image_processor=InklingRustImageProcessor()
|
||||
)
|
||||
return proc
|
||||
|
||||
|
||||
@@ -50,14 +56,19 @@ aud = wav_bytes()
|
||||
out = proc.assemble([1, 100, 2, 101, 3], [img], [aud])
|
||||
img_item = next(i for i in out.mm_items if i.modality.name == "IMAGE")
|
||||
aud_item = next(i for i in out.mm_items if i.modality.name == "AUDIO")
|
||||
assert img_item.hash == data_hash(img), "image hash != data_hash(raw bytes)"
|
||||
assert aud_item.hash == data_hash(aud), "audio hash != data_hash(raw bytes)"
|
||||
print(f" assemble: image hash={img_item.hash:#x} audio hash={aud_item.hash:#x} OK")
|
||||
# The rust image path stores its raw-bytes content hash (blake3) eagerly;
|
||||
# non-rust items get hashed lazily from the feature at set_pad_value time.
|
||||
assert img_item.hash == _rs_common.content_hash(img), "image hash != rust content_hash"
|
||||
assert aud_item.hash is None, "audio hash expected to be lazy (feature-based)"
|
||||
print(f" assemble: image hash={img_item.hash:#x} OK")
|
||||
|
||||
h0 = img_item.hash
|
||||
img_item.set_pad_value()
|
||||
assert img_item.hash == h0 and img_item.pad_value is not None
|
||||
print(f" set_pad_value: hash preserved, pad_value={img_item.pad_value} OK")
|
||||
aud_item.set_pad_value()
|
||||
assert aud_item.hash is not None and aud_item.pad_value is not None
|
||||
print(" set_pad_value: audio feature-hash filled OK")
|
||||
|
||||
out2 = proc.assemble([1, 100, 2], [img], [])
|
||||
assert out2.mm_items[0].hash == h0
|
||||
@@ -73,14 +84,8 @@ out3 = asyncio.run(
|
||||
assert all(i.hash == h0 for i in out3.mm_items), "data: URL roundtrip hash mismatch"
|
||||
print(" process_mm_data_async: concurrent resolve + hash OK")
|
||||
|
||||
orig = prc._resolve_media_item
|
||||
prc._resolve_media_item = lambda it: (time.sleep(0.3), orig(it))[1]
|
||||
t0 = time.perf_counter()
|
||||
asyncio.run(prc._resolve_media_items([data_url] * 8))
|
||||
elapsed = time.perf_counter() - t0
|
||||
prc._resolve_media_item = orig
|
||||
assert elapsed < 1.2, f"8x 0.3s resolves took {elapsed:.2f}s; expected ~0.3s"
|
||||
print(f" concurrency: 8 x 0.3s resolves in {elapsed:.2f}s OK")
|
||||
# (The old `_resolve_media_items` concurrency helper is gone; resolution now
|
||||
# happens inline in `process_mm_data_async`, covered by the roundtrip above.)
|
||||
|
||||
imgs_5 = [png_bytes(make_photo_like(1080, 1920, seed=s)) for s in range(5)]
|
||||
feats = [
|
||||
|
||||
@@ -6,10 +6,13 @@ import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "bench"))
|
||||
from bench_parity import make_photo_like
|
||||
|
||||
import sglang.srt.multimodal.inkling.image_processing as ip
|
||||
from sglang.srt.multimodal.inkling.image_processing import InklingImageProcessor
|
||||
from sglang.srt.multimodal.inkling.image_processing_rust import (
|
||||
InklingRustImageProcessor,
|
||||
)
|
||||
|
||||
|
||||
def encode(arr, fmt):
|
||||
@@ -21,17 +24,14 @@ def encode(arr, fmt):
|
||||
|
||||
|
||||
def run(images, use_rs: bool, rescale: bool):
|
||||
ip._rs_module = None
|
||||
os.environ["SGLANG_RS_MM_PREPROCESS"] = "1" if use_rs else "0"
|
||||
kwargs = (
|
||||
{}
|
||||
if rescale
|
||||
else {"rescale_image_frac": None, "rescale_image_max_upscaled_long_edge": None}
|
||||
)
|
||||
proc = ip.InklingImageProcessor(patch_size=40, **kwargs)
|
||||
out = proc.preprocess(images)
|
||||
assert (ip._rs_module is not False) == use_rs, "rust module gating mismatch"
|
||||
return out
|
||||
cls = InklingRustImageProcessor if use_rs else InklingImageProcessor
|
||||
proc = cls(patch_size=40, **kwargs)
|
||||
return proc.preprocess(images)
|
||||
|
||||
|
||||
def compare(tag, images, expect_exact, rescale=False):
|
||||
@@ -59,7 +59,7 @@ arr1 = make_photo_like(1080, 1920, seed=1)
|
||||
arr2 = make_photo_like(720, 1280, seed=2)
|
||||
arr3 = make_photo_like(480, 640, seed=3)
|
||||
|
||||
print("=== integration: InklingImageProcessor env-gated rust path ===")
|
||||
print("=== integration: Inkling rust vs python processor parity ===")
|
||||
compare("single PNG", [encode(arr1, "PNG")], expect_exact=True)
|
||||
compare("single JPEG", [encode(arr1, "JPEG")], expect_exact=False)
|
||||
compare(
|
||||
|
||||
@@ -6,7 +6,8 @@ import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
|
||||
import sglang.srt.multimodal._core.inkling
|
||||
from sglang.srt.multimodal._core import common as _rs_common
|
||||
from sglang.srt.multimodal._core import inkling as _rs_inkling
|
||||
|
||||
|
||||
def py_scaled_dims(
|
||||
@@ -41,9 +42,7 @@ def pil_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray:
|
||||
|
||||
|
||||
def rs_resize(arr: np.ndarray, tw: int, th: int) -> np.ndarray:
|
||||
return sglang.srt.multimodal._core.inkling.resize_rgb(arr, tw, th).reshape(
|
||||
th, tw, 3
|
||||
)
|
||||
return _rs_common.resize_rgb(arr, tw, th).reshape(th, tw, 3)
|
||||
|
||||
|
||||
CASES = [
|
||||
@@ -74,9 +73,9 @@ def test_scaled_dims_sweep():
|
||||
sizes += [(2048, 1024), (2049, 100), (1024, 2048), (1, 1), (4096, 4096)]
|
||||
for frac, cap in [(2.0, 2048), (1.5, 2048), (3.0, None), (None, None), (2.0, 1)]:
|
||||
for w, h in sizes:
|
||||
assert sglang.srt.multimodal._core.inkling.scaled_dims(
|
||||
assert _rs_common.scaled_dims(w, h, frac, cap) == py_scaled_dims(
|
||||
w, h, frac, cap
|
||||
) == py_scaled_dims(w, h, frac, cap), (
|
||||
), (
|
||||
w,
|
||||
h,
|
||||
frac,
|
||||
@@ -93,12 +92,10 @@ def test_decode_patchify_rescaled_matches_pil_pipeline():
|
||||
arr = rng.integers(0, 256, (1080, 1920, 3), dtype=np.uint8)
|
||||
buf = io.BytesIO()
|
||||
Image.fromarray(arr).save(buf, format="PNG")
|
||||
h, w, bits = sglang.srt.multimodal._core.inkling.decode_patchify(
|
||||
buf.getvalue(), 40, 2.0, 2048
|
||||
)
|
||||
h, w, bits = _rs_inkling.decode_patchify(buf.getvalue(), 40, 2.0, 2048)
|
||||
assert (w, h) == py_scaled_dims(1920, 1080, 2.0, 2048)
|
||||
ref_arr = pil_resize(arr, w, h)
|
||||
ref_bits = sglang.srt.multimodal._core.inkling.patchify_rgb(ref_arr, 40)
|
||||
ref_bits = _rs_inkling.patchify_rgb(ref_arr, 40)
|
||||
np.testing.assert_array_equal(bits, ref_bits)
|
||||
assert torch.from_numpy(bits).view(torch.bfloat16).shape[0] > 0
|
||||
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Shared fixtures for the native Rust multimodal suites.
|
||||
|
||||
Imported via ``sys.path`` from the sibling suites (unittest runs these files by
|
||||
path, so a package-relative import would break ``python <file>``); the module
|
||||
name is deliberately specific so it cannot shadow another suite's helpers on the
|
||||
process-global ``sys.path``.
|
||||
"""
|
||||
|
||||
import io
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import is_in_ci
|
||||
|
||||
register_cpu_ci(
|
||||
est_time=0, suite="base-a-test-cpu", disabled="Rust multimodal test helpers"
|
||||
)
|
||||
|
||||
|
||||
def load_core():
|
||||
"""The Rust ``_core`` extension, or ``None`` (→ skip) when not built
|
||||
locally. In CI a missing extension is a hard failure, never a silent
|
||||
skip — the CPU suite builds it from source."""
|
||||
try:
|
||||
from sglang.srt.multimodal import _core
|
||||
|
||||
return _core
|
||||
except ImportError:
|
||||
if is_in_ci():
|
||||
raise
|
||||
return None
|
||||
|
||||
|
||||
IMAGE_TOKEN_ID = 900
|
||||
VISION_START_ID = 901
|
||||
VISION_END_ID = 902
|
||||
VIDEO_TOKEN_ID = 903
|
||||
|
||||
PROCESSOR_CONFIGS = {
|
||||
"qwen2_vl": dict(
|
||||
patch_size=14,
|
||||
merge_size=2,
|
||||
temporal_patch_size=2,
|
||||
min_pixels=56 * 56,
|
||||
max_pixels=28 * 28 * 1280,
|
||||
image_mean=[0.48145466, 0.4578275, 0.40821073],
|
||||
image_std=[0.26862954, 0.26130258, 0.27577711],
|
||||
),
|
||||
"qwen2_5_vl": dict(
|
||||
patch_size=14,
|
||||
merge_size=2,
|
||||
temporal_patch_size=2,
|
||||
min_pixels=56 * 56,
|
||||
max_pixels=28 * 28 * 1280,
|
||||
image_mean=[0.5] * 3,
|
||||
image_std=[0.5] * 3,
|
||||
),
|
||||
"qwen3_5": dict(
|
||||
patch_size=16,
|
||||
merge_size=2,
|
||||
temporal_patch_size=2,
|
||||
min_pixels=65536,
|
||||
max_pixels=16777216,
|
||||
image_mean=[0.5] * 3,
|
||||
image_std=[0.5] * 3,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def make_image(width, height, seed=0):
|
||||
rng = np.random.default_rng(seed)
|
||||
y, x = np.mgrid[0:height, 0:width]
|
||||
base = np.stack(
|
||||
(x * 255 / max(width - 1, 1), y * 255 / max(height - 1, 1), (x + y) % 256),
|
||||
axis=-1,
|
||||
)
|
||||
return Image.fromarray(
|
||||
np.clip(base + rng.integers(0, 24, base.shape), 0, 255).astype(np.uint8)
|
||||
)
|
||||
|
||||
|
||||
def image_bytes(width, height, seed=0):
|
||||
buffer = io.BytesIO()
|
||||
make_image(width, height, seed).save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
def spec_json(config, image_token_id=IMAGE_TOKEN_ID):
|
||||
return json.dumps({"family": "qwen_vl", "image_token_id": image_token_id, **config})
|
||||
@@ -0,0 +1,74 @@
|
||||
"""Smoke coverage for the ``_core.inkling`` PyO3 bindings.
|
||||
|
||||
Covers the ``#[pyfunction]``s in ``rust/sglang-mm/src/inkling/mod.rs``
|
||||
(``patchify_rgb`` / ``decode_patchify`` / ``decode_patchify_batch`` /
|
||||
``preprocess_images`` / ``rescale_patchify_hash``) plus
|
||||
``_core.common.content_hash``.
|
||||
|
||||
The inkling bindings are otherwise exercised only by the GPU e2e model test;
|
||||
this pins the binding surface (signatures, dtypes, cross-binding consistency)
|
||||
in the CPU suite so a rework of the extension can't silently break them.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from _mm_rust_utils import image_bytes, load_core # noqa: E402
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
CORE = load_core()
|
||||
PATCH_SIZE = 16
|
||||
|
||||
|
||||
@unittest.skipUnless(CORE, "sglang-mm extension not built")
|
||||
class TestInklingBindings(CustomTestCase):
|
||||
def test_bindings_are_consistent(self):
|
||||
data = image_bytes(50, 34)
|
||||
h, w, patches = CORE.inkling.decode_patchify(data, PATCH_SIZE)
|
||||
self.assertEqual((h, w), (34, 50))
|
||||
patches = np.asarray(patches)
|
||||
self.assertEqual(patches.dtype, np.uint16)
|
||||
# One padded patch column, ceil rows (the inkling grid convention).
|
||||
expected_len = -(-h // PATCH_SIZE) * (w // PATCH_SIZE + 1) * PATCH_SIZE**2 * 3
|
||||
self.assertEqual(patches.size, expected_len)
|
||||
|
||||
# patchify_rgb on the decoded array must match the fused decode path.
|
||||
dh, dw, rgb = CORE.common.image_decode_rgb(data)
|
||||
arr = np.asarray(rgb).reshape(dh, dw, 3)
|
||||
np.testing.assert_array_equal(
|
||||
np.asarray(CORE.inkling.patchify_rgb(arr, PATCH_SIZE)), patches
|
||||
)
|
||||
|
||||
# Batch and hashed variants agree with the single-image call.
|
||||
[(bh, bw, batch)] = CORE.inkling.decode_patchify_batch([data], PATCH_SIZE)
|
||||
self.assertEqual((bh, bw), (h, w))
|
||||
np.testing.assert_array_equal(np.asarray(batch), patches)
|
||||
[(ph, pw, pre, phash)] = CORE.inkling.preprocess_images([data], PATCH_SIZE)
|
||||
self.assertEqual((ph, pw), (h, w))
|
||||
np.testing.assert_array_equal(np.asarray(pre), patches)
|
||||
self.assertEqual(phash, CORE.common.content_hash(data))
|
||||
|
||||
rh, rw, rpatches, rhash = CORE.inkling.rescale_patchify_hash(
|
||||
arr, data, PATCH_SIZE
|
||||
)
|
||||
self.assertEqual((rh, rw), (h, w))
|
||||
np.testing.assert_array_equal(np.asarray(rpatches), patches)
|
||||
self.assertEqual(rhash, phash)
|
||||
|
||||
def test_invalid_inputs_rejected(self):
|
||||
with self.assertRaises(ValueError):
|
||||
CORE.inkling.decode_patchify(b"junk", PATCH_SIZE)
|
||||
with self.assertRaises(ValueError):
|
||||
CORE.inkling.decode_patchify(image_bytes(16, 16), 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,109 @@
|
||||
"""Native driver error paths: out-of-scope and malformed inputs are rejected.
|
||||
|
||||
Covers ``process`` in ``rust/sglang-mm/src/driver.rs`` (via the
|
||||
``_core.qwen_vl.process_native_mm`` binding). The wire-payload parsing that
|
||||
feeds this driver (modality/shape rejection) lives in ``sglang-server``'s
|
||||
message layer and is tested with the integration PR.
|
||||
|
||||
There is no Python fallback path, so the server rejects every driver error
|
||||
back to the client as a 400; the message must say why (placeholder mismatch,
|
||||
undecodable image, missing prompt). This pins that contract for each
|
||||
rejection class.
|
||||
"""
|
||||
|
||||
import io
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from _mm_rust_utils import ( # noqa: E402
|
||||
IMAGE_TOKEN_ID,
|
||||
PROCESSOR_CONFIGS,
|
||||
VISION_END_ID,
|
||||
VISION_START_ID,
|
||||
image_bytes,
|
||||
load_core,
|
||||
spec_json,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
|
||||
QWEN_CORE = getattr(load_core(), "qwen_vl", None)
|
||||
SPEC = spec_json(PROCESSOR_CONFIGS["qwen2_5_vl"])
|
||||
IMAGE_IDS = [7, VISION_START_ID, IMAGE_TOKEN_ID, VISION_END_ID, 8]
|
||||
|
||||
|
||||
def gif_bytes():
|
||||
buffer = io.BytesIO()
|
||||
Image.fromarray(np.zeros((16, 16, 3), dtype=np.uint8)).save(buffer, format="GIF")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
QWEN_CORE and hasattr(QWEN_CORE, "process_native_mm"),
|
||||
"sglang-mm native Qwen driver not built",
|
||||
)
|
||||
class TestNativeDriverErrorPaths(CustomTestCase):
|
||||
def assert_rejected(self, input_ids, images, pattern, spec=SPEC):
|
||||
with self.assertRaisesRegex(ValueError, pattern):
|
||||
QWEN_CORE.process_native_mm(input_ids, images, spec)
|
||||
|
||||
def test_degenerate_geometry_rejected_not_panicked(self):
|
||||
"""A thin image against a tight ``max_pixels`` floors a side of the
|
||||
smart_resize target to 0. That used to panic on a worker thread inside
|
||||
the resize coefficient math; a Rust panic surfaces as ``PanicException``
|
||||
(a ``BaseException``), not ``ValueError``, so this asserts the request
|
||||
is rejected cleanly rather than crashing the pipeline."""
|
||||
config = dict(PROCESSOR_CONFIGS["qwen2_5_vl"], min_pixels=3136, max_pixels=3136)
|
||||
self.assert_rejected(
|
||||
IMAGE_IDS,
|
||||
[image_bytes(2000, 10)],
|
||||
"smart_resize",
|
||||
spec=spec_json(config),
|
||||
)
|
||||
|
||||
def test_placeholder_count_mismatches_rejected(self):
|
||||
cases = {
|
||||
"no placeholder": ([7, 8], [image_bytes(80, 80)]),
|
||||
"more images": (IMAGE_IDS, [image_bytes(80, 80), image_bytes(88, 80, 1)]),
|
||||
"more placeholders": (IMAGE_IDS + IMAGE_IDS, [image_bytes(80, 80)]),
|
||||
}
|
||||
for name, (ids, images) in cases.items():
|
||||
with self.subTest(case=name):
|
||||
self.assert_rejected(ids, images, "placeholder")
|
||||
|
||||
def test_undecodable_images_rejected(self):
|
||||
# Corrupt bytes are outside the native decoder's scope; the server
|
||||
# rejects them as a 400.
|
||||
self.assert_rejected(IMAGE_IDS, [b"junk"], "decode")
|
||||
|
||||
def test_gif_serves_through_the_pipeline(self):
|
||||
"""GIF moved from rejected to served when the pure-Rust webp/gif/bmp
|
||||
decoders were enabled; this pins the accept side of that contract flip
|
||||
(the reject side used to be asserted here and broke in CI)."""
|
||||
_, _, grids, _, offsets, _, _ = QWEN_CORE.process_native_mm(
|
||||
IMAGE_IDS, [gif_bytes()], SPEC
|
||||
)
|
||||
self.assertEqual(len(grids), 1)
|
||||
self.assertEqual(len(offsets), 1)
|
||||
|
||||
def test_missing_text_and_input_ids_rejected(self):
|
||||
for input_ids in (None, []):
|
||||
with self.subTest(input_ids=input_ids):
|
||||
self.assert_rejected(
|
||||
input_ids, [image_bytes(80, 80)], "without text or input_ids"
|
||||
)
|
||||
|
||||
def test_image_free_request_rejected(self):
|
||||
self.assert_rejected(IMAGE_IDS, [], "image sources")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,105 @@
|
||||
"""Qwen native image preprocessing parity against Transformers.
|
||||
|
||||
Covers ``QwenVlProcessor::process_item`` and ``smart_resize`` in
|
||||
``rust/sglang-mm/src/qwen_vl/mod.rs`` (via the ``_core.qwen_vl.preprocess``
|
||||
and ``smart_resize_py`` bindings), against the HF Qwen2-VL image processors
|
||||
and the Python ``smart_resize``.
|
||||
|
||||
Both HF processors are pinned, because they resample differently and only one
|
||||
of them is what a server actually runs. On transformers 5.x
|
||||
``Qwen2VLImageProcessor`` is the torchvision path (the ``Fast`` suffix was
|
||||
dropped) and is what ``AutoImageProcessor`` hands SGLang by default;
|
||||
``Qwen2VLImageProcessorPil`` is the PIL path, reachable via
|
||||
``--disable-fast-image-processor``. The Rust resize is a bit-exact clone of
|
||||
PIL's fixed-point kernel, so the PIL processor is asserted exactly and the
|
||||
torchvision one carries the cross-implementation envelope.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from _mm_rust_utils import ( # noqa: E402
|
||||
PROCESSOR_CONFIGS,
|
||||
image_bytes,
|
||||
load_core,
|
||||
make_image,
|
||||
spec_json,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=20, suite="base-a-test-cpu")
|
||||
|
||||
QWEN_CORE = getattr(load_core(), "qwen_vl", None)
|
||||
|
||||
|
||||
SIZES = ((640, 480), (1024, 683), (50, 40), (300, 301))
|
||||
|
||||
|
||||
@unittest.skipUnless(QWEN_CORE, "sglang-mm Qwen binding not built")
|
||||
class TestQwenImagePreprocess(CustomTestCase):
|
||||
def _assert_matches(self, processor, max_diff, mean_diff):
|
||||
for family, config in PROCESSOR_CONFIGS.items():
|
||||
hf = processor(**config)
|
||||
for index, size in enumerate(SIZES):
|
||||
with self.subTest(family=family, size=size):
|
||||
image = make_image(*size, seed=index)
|
||||
actual, grid = QWEN_CORE.preprocess(
|
||||
image_bytes(*size, seed=index), spec_json(config)
|
||||
)
|
||||
expected = hf(images=[image], return_tensors="pt")
|
||||
self.assertEqual(grid, tuple(expected.image_grid_thw[0].tolist()))
|
||||
diff = np.abs(
|
||||
np.asarray(actual).reshape(expected.pixel_values.shape)
|
||||
- expected.pixel_values.numpy()
|
||||
)
|
||||
# LessEqual, not Less: max_diff=0.0 is a real bound here.
|
||||
self.assertLessEqual(diff.max(), max_diff)
|
||||
self.assertLessEqual(diff.mean(), mean_diff)
|
||||
|
||||
def test_features_match_pil_processor_exactly(self):
|
||||
"""Against the PIL processor the native path is bit-exact, so this is
|
||||
asserted with zero tolerance: every stage (smart_resize geometry,
|
||||
the fixed-point bicubic kernel, rescale/normalize, HF patch order) is
|
||||
pinned, and any drift in any of them shows up here rather than being
|
||||
absorbed by a tolerance."""
|
||||
from transformers.models.qwen2_vl.image_processing_pil_qwen2_vl import (
|
||||
Qwen2VLImageProcessorPil,
|
||||
)
|
||||
|
||||
self._assert_matches(Qwen2VLImageProcessorPil, max_diff=0.0, mean_diff=0.0)
|
||||
|
||||
def test_features_match_torchvision_processor_within_envelope(self):
|
||||
"""The torchvision processor is what a default server runs, so its
|
||||
divergence is bounded separately — it is a different antialiased-bicubic
|
||||
implementation, which the bit-exact PIL assertion above says nothing
|
||||
about. Measured worst case is max 0.030 / mean 6.7e-5 (≈2 u8 levels
|
||||
after normalize with the qwen2_vl std)."""
|
||||
from transformers.models.qwen2_vl.image_processing_qwen2_vl import (
|
||||
Qwen2VLImageProcessor,
|
||||
)
|
||||
|
||||
self._assert_matches(Qwen2VLImageProcessor, max_diff=0.035, mean_diff=1e-3)
|
||||
|
||||
def test_smart_resize_matches_python(self):
|
||||
from sglang.srt.multimodal.processors.qwen_vl import smart_resize
|
||||
|
||||
cases = (
|
||||
(1365, 2048, 28, 3136, 12845056),
|
||||
(3000, 4000, 28, 3136, 1003520),
|
||||
(20, 20, 28, 3136, 12845056),
|
||||
(1365, 2048, 32, 65536, 16777216),
|
||||
(4000, 48, 32, 4, 1 << 30),
|
||||
)
|
||||
for case in cases:
|
||||
with self.subTest(case=case):
|
||||
self.assertEqual(QWEN_CORE.smart_resize_py(*case), smart_resize(*case))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,104 @@
|
||||
"""Qwen placeholder expansion, offsets, and M-RoPE parity.
|
||||
|
||||
Covers ``layout_by_placeholder`` / ``apply_layout`` in
|
||||
``rust/sglang-mm/src/common/token_layout.rs`` and ``mrope_image_only`` in
|
||||
``rust/sglang-mm/src/qwen_vl/mod.rs`` (via the
|
||||
``_core.qwen_vl.process_native_mm`` and ``mrope_image_only_py``
|
||||
bindings), against ``BaseMultimodalProcessor`` expansion/offsets and
|
||||
``MRotaryEmbedding.get_rope_index``.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from _mm_rust_utils import ( # noqa: E402
|
||||
IMAGE_TOKEN_ID,
|
||||
PROCESSOR_CONFIGS,
|
||||
VIDEO_TOKEN_ID,
|
||||
VISION_END_ID,
|
||||
VISION_START_ID,
|
||||
image_bytes,
|
||||
load_core,
|
||||
spec_json,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
|
||||
QWEN_CORE = getattr(load_core(), "qwen_vl", None)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
QWEN_CORE and hasattr(QWEN_CORE, "process_native_mm"),
|
||||
"sglang-mm native Qwen driver not built",
|
||||
)
|
||||
class TestQwenPromptGeometry(CustomTestCase):
|
||||
def test_placeholder_expansion_and_offsets(self):
|
||||
from sglang.srt.multimodal.processors.base_processor import (
|
||||
BaseMultimodalProcessor,
|
||||
)
|
||||
|
||||
config = PROCESSOR_CONFIGS["qwen2_5_vl"]
|
||||
for image_count in (1, 2):
|
||||
ids = [7]
|
||||
for _ in range(image_count):
|
||||
ids.extend((VISION_START_ID, IMAGE_TOKEN_ID, VISION_END_ID, 8))
|
||||
images = [image_bytes(96 + 8 * i, 80, i) for i in range(image_count)]
|
||||
with self.subTest(image_count=image_count):
|
||||
actual_ids, _, grids, _, offsets, _, _ = QWEN_CORE.process_native_mm(
|
||||
ids, images, spec_json(config)
|
||||
)
|
||||
counts = [t * h * w // config["merge_size"] ** 2 for t, h, w in grids]
|
||||
expected_ids = BaseMultimodalProcessor._expand_input_ids(
|
||||
ids, counts, IMAGE_TOKEN_ID
|
||||
)
|
||||
self.assertEqual(actual_ids, expected_ids)
|
||||
self.assertEqual(
|
||||
offsets,
|
||||
BaseMultimodalProcessor.get_mm_items_offset(
|
||||
torch.tensor(expected_ids), IMAGE_TOKEN_ID
|
||||
),
|
||||
)
|
||||
|
||||
def test_mrope_matches_model_reference(self):
|
||||
"""The single Rust image-only M-RoPE must match ``get_rope_index``
|
||||
for every native Qwen family (image-only makes them coincide)."""
|
||||
from sglang.srt.layers.rotary_embedding import MRotaryEmbedding
|
||||
|
||||
grids = [(1, 4, 6), (1, 6, 4)]
|
||||
ids, items = [10], []
|
||||
for grid in grids:
|
||||
ids.extend((VISION_START_ID, IMAGE_TOKEN_ID))
|
||||
start = len(ids) - 1
|
||||
ids.extend([IMAGE_TOKEN_ID] * (np.prod(grid) // 4 - 1))
|
||||
items.append((start, len(ids) - 1, *grid))
|
||||
ids.extend((VISION_END_ID, 11))
|
||||
|
||||
actual, delta = QWEN_CORE.mrope_image_only_py(len(ids), items, 2)
|
||||
actual = np.asarray(actual).reshape(3, -1)
|
||||
for model_type in ("qwen2_vl", "qwen2_5_vl", "qwen3_vl", "qwen3_5"):
|
||||
with self.subTest(model_type=model_type):
|
||||
expected, expected_delta = MRotaryEmbedding.get_rope_index(
|
||||
spatial_merge_size=2,
|
||||
image_token_id=IMAGE_TOKEN_ID,
|
||||
video_token_id=VIDEO_TOKEN_ID,
|
||||
vision_start_token_id=VISION_START_ID,
|
||||
model_type=model_type,
|
||||
tokens_per_second=2 if model_type == "qwen2_5_vl" else None,
|
||||
input_ids=torch.tensor(ids).unsqueeze(0),
|
||||
image_grid_thw=torch.tensor(grids),
|
||||
video_grid_thw=None,
|
||||
)
|
||||
np.testing.assert_array_equal(actual, expected.squeeze(1).numpy())
|
||||
self.assertEqual(delta, int(expected_delta.item()))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Model-independent parity tests for Rust multimodal source loading.
|
||||
|
||||
Covers ``fetch_bytes`` in ``rust/sglang-mm/src/common/fetch.rs`` (via the
|
||||
``_core.common.fetch_bytes`` binding), against the Python reference
|
||||
``sglang.srt.utils.common.get_image_bytes``.
|
||||
"""
|
||||
|
||||
import base64
|
||||
import http.server
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from sglang.srt.utils.common import get_image_bytes
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from _mm_rust_utils import load_core # noqa: E402
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
CORE = load_core()
|
||||
FETCH = CORE and CORE.common.fetch_bytes
|
||||
|
||||
|
||||
@unittest.skipUnless(FETCH, "sglang-mm fetch binding not built")
|
||||
class TestRustMediaSourceLoading(CustomTestCase):
|
||||
DATA = b"native-mm-source"
|
||||
|
||||
def test_inline_sources(self):
|
||||
encoded = base64.b64encode(self.DATA).decode()
|
||||
for source in (encoded, f"data:application/octet-stream;base64,{encoded}"):
|
||||
with self.subTest(source=source[:8]):
|
||||
self.assertEqual(bytes(FETCH(source)), get_image_bytes(source))
|
||||
|
||||
def test_file_sources(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "image.png"
|
||||
path.write_bytes(self.DATA)
|
||||
self.assertEqual(bytes(FETCH(str(path))), get_image_bytes(str(path)))
|
||||
# `file://` is asserted against the payload, not the Python helper,
|
||||
# on purpose: `get_image_bytes` passes the un-stripped URL straight
|
||||
# to `open()`, so the reference raises here. The native path strips
|
||||
# the scheme and succeeds — a deliberate divergence, not an omission.
|
||||
with self.assertRaises(OSError):
|
||||
get_image_bytes(path.as_uri())
|
||||
self.assertEqual(bytes(FETCH(path.as_uri())), self.DATA)
|
||||
|
||||
def test_http_source(self):
|
||||
data = self.DATA
|
||||
|
||||
class Handler(http.server.BaseHTTPRequestHandler):
|
||||
def do_GET(self):
|
||||
self.send_response(200)
|
||||
self.end_headers()
|
||||
self.wfile.write(data)
|
||||
|
||||
def log_message(self, *_):
|
||||
pass
|
||||
|
||||
server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
|
||||
thread = threading.Thread(target=server.serve_forever, daemon=True)
|
||||
thread.start()
|
||||
try:
|
||||
url = f"http://127.0.0.1:{server.server_port}/image"
|
||||
self.assertEqual(bytes(FETCH(url)), get_image_bytes(url))
|
||||
finally:
|
||||
server.shutdown()
|
||||
server.server_close()
|
||||
thread.join()
|
||||
|
||||
def test_invalid_sources_fail(self):
|
||||
for source in ("not base64!", "/definitely/missing/image.png"):
|
||||
with self.subTest(source=source):
|
||||
with self.assertRaises(ValueError):
|
||||
FETCH(source)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,102 @@
|
||||
"""Model-independent image decode parity for native Rust MM.
|
||||
|
||||
Covers ``decode_rgb`` in ``rust/sglang-mm/src/common/mod.rs`` (via the
|
||||
``_core.common.image_decode_rgb`` binding), against PIL's
|
||||
``Image.open(...).convert("RGB")``.
|
||||
"""
|
||||
|
||||
import io
|
||||
import sys
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||||
from _mm_rust_utils import load_core # noqa: E402
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
CORE = load_core()
|
||||
DECODE = CORE and CORE.common.image_decode_rgb
|
||||
|
||||
|
||||
def encode(image, fmt, **kwargs):
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format=fmt, **kwargs)
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
@unittest.skipUnless(DECODE, "sglang-mm decode binding not built")
|
||||
class TestRustImageDecode(CustomTestCase):
|
||||
def assert_matches_pil(self, data, tolerance=0):
|
||||
height, width, pixels = DECODE(data)
|
||||
expected = np.asarray(Image.open(io.BytesIO(data)).convert("RGB"))
|
||||
self.assertEqual((height, width), expected.shape[:2])
|
||||
actual = np.asarray(pixels).reshape(expected.shape)
|
||||
diff = np.abs(actual.astype(int) - expected.astype(int))
|
||||
self.assertLessEqual(diff.max(), tolerance)
|
||||
|
||||
def test_png_modes_match_pil(self):
|
||||
rgb = np.random.default_rng(1).integers(0, 256, (19, 23, 3), dtype=np.uint8)
|
||||
cases = [
|
||||
("RGB", Image.fromarray(rgb)),
|
||||
("L", Image.fromarray(rgb[..., 0])),
|
||||
("RGBA", Image.fromarray(np.dstack((rgb, rgb[..., 0])))),
|
||||
("P", Image.fromarray(rgb).quantize(colors=16)),
|
||||
]
|
||||
for mode, image in cases:
|
||||
with self.subTest(mode=mode):
|
||||
self.assert_matches_pil(encode(image, "PNG"))
|
||||
|
||||
def test_jpeg_modes_match_with_decoder_tolerance(self):
|
||||
rgb = np.random.default_rng(2).integers(0, 256, (31, 29, 3), dtype=np.uint8)
|
||||
image = Image.fromarray(rgb)
|
||||
exif = image.getexif()
|
||||
exif[274] = 6 # EXIF orientation: neither decoder applies it
|
||||
cases = [
|
||||
("RGB", encode(image, "JPEG")),
|
||||
("L", encode(Image.fromarray(rgb[..., 0]), "JPEG")),
|
||||
("CMYK", encode(image.convert("CMYK"), "JPEG")),
|
||||
("EXIF-rotated", encode(image, "JPEG", exif=exif)),
|
||||
]
|
||||
for mode, data in cases:
|
||||
with self.subTest(mode=mode):
|
||||
self.assert_matches_pil(data, tolerance=3)
|
||||
|
||||
def test_lossless_formats_match_pil_exactly(self):
|
||||
"""GIF/BMP/lossless-WebP joined the native decoder with the pure-Rust
|
||||
webp/gif/bmp enablement. Their reconstruction is exact by format spec,
|
||||
so parity with PIL is pinned at zero tolerance like PNG."""
|
||||
rgb = np.random.default_rng(4).integers(0, 256, (17, 21, 3), dtype=np.uint8)
|
||||
image = Image.fromarray(rgb)
|
||||
cases = {
|
||||
"gif": encode(image.quantize(colors=64), "GIF"),
|
||||
"bmp": encode(image, "BMP"),
|
||||
"webp-lossless": encode(image, "WEBP", lossless=True),
|
||||
}
|
||||
for name, data in cases.items():
|
||||
with self.subTest(fmt=name):
|
||||
self.assert_matches_pil(data)
|
||||
|
||||
def test_unsupported_inputs_fail(self):
|
||||
gray16 = Image.fromarray(np.zeros((9, 9), dtype=np.uint16))
|
||||
cases = {
|
||||
"corrupt": b"not an image",
|
||||
# >8-bit depths must error (the request is then rejected — there
|
||||
# is no Python fallback) — never silently diverge from PIL's
|
||||
# clipping.
|
||||
"png16": encode(gray16, "PNG"),
|
||||
}
|
||||
for name, data in cases.items():
|
||||
with self.subTest(input=name):
|
||||
with self.assertRaises(ValueError):
|
||||
DECODE(data)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user