From 65cb1fddd0e64c675ffcd9c49b7f6c4d28c9506e Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Thu, 22 Oct 2026 23:45:49 +0900 Subject: [PATCH 1/5] fix(rocm): route AsStrided through the HIP encoder The latest MLX pin added backend-specific AsStrided temporary retention, but the ROCm core overlay kept the upstream CUDA-or-Metal dispatch and pulled Metal/Metal.hpp into Linux builds. Select the ROCm device header and command encoder explicitly, then refresh the overlay provenance records. The ROCm overlay gate and a syntax-only compile of the affected translation unit pass. Refs #2349 --- .../overlays/rocm/CORE_RESIDUAL.diff | 26 ++++++++++++++++--- .../mlxcelverse/overlays/rocm/LOCAL_FIXES.md | 1 + .../rocm/mlx/backend/gpu/primitives.cpp | 4 +++ 3 files changed, 27 insertions(+), 4 deletions(-) diff --git a/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff b/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff index ec5175d45..e00e1a2db 100644 --- a/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff +++ b/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff @@ -24,7 +24,7 @@ # overlay-sha256: c255fb0f18fb597cd7ba1708dcc5799ea26ee8541820dedb8997a279d9257496 mlx/backend/common/buffer_cache.h # overlay-sha256: 6755397b3410b82fe451c1aa0455ba722d4106a67d72fb9ce6a4eecda61ca920 mlx/backend/common/compiled.cpp # overlay-sha256: 28759d52cd6ec6f55a6ba85ae6e594e24c042fc55bfe139ac1856e556ee42e10 mlx/backend/common/compiled.h -# overlay-sha256: 9dd0ee6035eb8fbbac8619ba81578e982c0ecdbfe2fd2952abab507874cf2a98 mlx/backend/gpu/primitives.cpp +# overlay-sha256: c538cb3119b20893647ea069d22e1368f9a84bcd574e6a4e62c1d7b90ac91cbc mlx/backend/gpu/primitives.cpp # overlay-sha256: 849838a3e3354496e7c4e571b98491e7c80f94f4b84fe9e642093f332069fc8a mlx/compile.cpp # overlay-sha256: 67e68d6e70df6271ed2f54d75957c4762653b6a32fa6c666ea3b014c5d425acd mlx/device.cpp # overlay-sha256: 26584a93a3720250aa59c6dd9f091863f8c047691ce55f9a988a01da390ab02a mlx/fast.cpp @@ -36,7 +36,7 @@ # overlay-sha256: 5223221a04503b7f9ab5215fc594e57884b1388a1ca84e6b2ad33edb451d2525 mlx/stream.cpp # note mlx/allocator.h: LOCAL_FIXES.md item 47 (owned ROCm asynchronous Data allocation friendship). # note mlx/backend/common/compiled.cpp: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's negative_strides tracking and the fork's constant-input skipping. -# note mlx/backend/gpu/primitives.cpp: LOCAL_FIXES.md item 24: DynamicSliceUpdate copies through copy_gpu (donates only when is_donatable holds) and drops the fork's graph_active() forced donation. +# note mlx/backend/gpu/primitives.cpp: LOCAL_FIXES.md items 24 and 50: DynamicSliceUpdate copies through copy_gpu (donates only when is_donatable holds) and drops the fork's graph_active() forced donation; AsStrided retains its temporary through the ROCm command encoder. # note mlx/fast_primitives.h: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's compile_options and the fork's output_input_aliases on fast::CustomKernel. # note mlx/io/safetensors.cpp: LOCAL_FIXES.md items 1 and 49: resolves the pin/fork conflict (empty-tensor guard + ROCm staged_write) and schedules load_safetensors on the GPU stream when rocm::is_available(). @@ -102,7 +102,16 @@ std::vector strides_vec{out.strides()}; --- a/mlx/backend/gpu/primitives.cpp +++ b/mlx/backend/gpu/primitives.cpp -@@ -19,13 +19,6 @@ +@@ -9,6 +9,8 @@ + #if defined(MLX_USE_CUDA) + #include + #include "mlx/backend/cuda/device.h" ++#elif defined(MLX_USE_ROCM) ++#include "mlx/backend/rocm/device.h" + #else + #include "mlx/backend/metal/device.h" + #endif +@@ -19,13 +21,6 @@ #define MLX_PROFILER_RANGE(message) nvtx3::scoped_range r(message) #else #define MLX_PROFILER_RANGE(message) @@ -116,7 +125,16 @@ #endif namespace mlx::core { -@@ -156,29 +149,14 @@ +@@ -39,6 +34,8 @@ + auto copy = contiguous_copy_gpu(inputs[0], stream()); + #if defined(MLX_USE_CUDA) + cu::get_command_encoder(stream()).add_temporary(copy); ++#elif defined(MLX_USE_ROCM) ++ rocm::get_command_encoder(stream()).add_temporary(copy); + #else + metal::get_command_encoder(stream()).add_temporary(copy); + #endif +@@ -156,29 +153,14 @@ return; } diff --git a/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md b/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md index 6f157a039..c51ea80b4 100644 --- a/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md +++ b/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md @@ -77,3 +77,4 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code 47. **Move-only allocation ownership.** Upstream changed `allocator::malloc` and `array::set_data` to transfer `allocator::Data`, removed the free wrapper and made donation require owned storage. The new core overlay `mlx/allocator.h` adds only the ROCm forward declaration and friendship corresponding to CUDA's existing asynchronous allocation entry point. `mlx/backend/rocm/allocator.h` and `mlx/backend/rocm/allocator.cpp` return genuinely owned `Data` from the public `rocm::malloc_async`; the underlying `RocmAllocator` still returns raw `Buffer` as required by the allocator interface. This preserves ordinary donation and the existing encoder's shared array retention until batch completion; destructing `Data` reaches the same ROCm cache/deferred-free logic exactly once. `mlx/backend/rocm/gemms/naive_gemm.hip` wraps slab views in borrowed `Data` with the existing shell-only deleter, so they cannot donate externally owned slab storage. `mlx/backend/rocm/moe_swiglu.cpp` and `mlx/backend/rocm/reduce/reduce_utils.hpp` remove the obsolete free/deleter arguments when transferring owned allocations. Retargeting `mlx/backend/common/compiled.{h,cpp}` brings in upstream's `Data`-returning allocation callback and avoids borrowed wrappers. The host mock regression in `scripts/mlxcelverse/allocator_ownership_test.{cpp,sh}` compiles the actual header and async handoff, checks move-only/owned/borrowed/empty allocations, pending-batch shared retention, and cleanup on exceptions before and after enqueue. It does not execute HIP or establish GPU completion correctness. 48. **New primitive and capability contracts.** `mlx/backend/rocm/fence.cpp` updates `Fence::wait` to take the target `uint32_t value` (`fence->event.wait(value)`) and `Fence::update` to return `fence->count`, matching upstream `mlx/fence.h` at `243bdf09`. `mlx/backend/rocm/scaled_dot_product_attention.cpp` accepts upstream's `has_sinks` argument and excludes sinks from WMMA admission, matching the existing execution guard; vector and scalar flash keep their sinks support, and unsupported sinks plus LSE requests take the graph fallback (or reject `force_fused`). `mlx/backend/rocm/primitives.cpp` defines `GatedDeltaUpdateVJP::use_fallback` and `eval_gpu`: the capability check always selects upstream's differentiable graph, like the existing forward fallback, while the explicit unsupported-kernel stub satisfies linking without claiming a native implementation. `mlx/backend/rocm/binary.hip` adds `FloorDivide::eval_gpu` using the existing binary dispatch. Its operation in `mlx/backend/rocm/device/binary_ops.hpp`, also shared by `DivMod`, now floors signed and floating negative quotients instead of truncating them, matching target CUDA semantics. `mlx/backend/rocm/compiled.cpp` makes the same correction for generated kernels and keeps integral division exact rather than rounding 64-bit values through float. `mlx/backend/rocm/fft.hip` documents the audited `odd_out` support: the real transform uses the exact output shape for the hipFFT plan length, cache key and inverse scale, so odd inverse-real lengths are not reconstructed as an even `2 * (input - 1)`. No FFT algorithm change is needed. Actual ROCm compile, graph-fallback execution, negative floor-division numerics and odd-length FFT execution remain to be checked on an AMD host; this retarget host has no HIP toolchain or AMD device. 49. **GPU-stream `load_safetensors`, reused pinned host staging pool, and upload counters in `Load::eval_gpu`.** In `mlx/io/safetensors.cpp`, `load_safetensors(Reader, StreamOrDevice)` checked only `cu::is_available()` when choosing the default `Load` stream, so ROCm defaulted to `Device::cpu` (`Load::eval_cpu`), bypassing `Load::eval_gpu` and introducing cross-stream CPU-to-GPU fences whenever a loader derived tensors from `Load` arrays; it now checks `cu::is_available() || rocm::is_available()`. In `mlx/backend/rocm/load.cpp`, `Load::eval_gpu` previously called `hipHostMalloc(&out_ptr, nbytes, hipHostMallocDefault)` and queued `hipLaunchHostFunc(encoder.stream(), hip_host_free_callback, out_ptr)` for every tensor. Under eager weight materialization (lablup/mlxcel#1564), `eval_all` runs all kept checkpoint tensors back-to-back, so per-tensor `hipHostMalloc`/`hipHostFree` paid 413 ms per 2 GiB in KFD page-pinning churn on gfx1151 (617 ms total for 16 x 128 MiB vs 163 ms with reused slots, 3.8x faster) and kept up to a batch's worth of duplicate pinned staging buffers live alongside the device weights. `mlx/backend/rocm/load.cpp` now uploads through four reused 64 MiB pinned host slots per HIP device (`hipHostMallocPortable` + per-slot `hipEvent_t` with `hipEventDisableTiming`), bounding pinned staging memory at 256 MiB per device and eliminating the per-tensor host callback, with the synchronous pageable `malloc` + `hipMemcpy` path kept as fallback when `MLXCEL_ROCM_PINNED_LOAD=0` (or `MLXCEL_CUDA_PINNED_LOAD=0`) or when slot allocation fails. `mlx/backend/rocm/load_staging.h` exposes `rocm::load_transfer_stats()`, wired into `weight_load_transfer_stats` in `src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp` under `MLXCEL_BRIDGE_ROCM_BACKEND` so `MLXCEL_PROFILE_TTFT=1` and `weight_materialize_tests::pinned_upload_chunks_and_wraps_the_slot_ring` inspect the same upload counters on ROCm as on CUDA. Applies to the fork (for `safetensors.cpp`); staging pool kept in mlxcelverse (lablup/mlxcel#1564). +50. **`AsStrided` retains its contiguous copy on ROCm.** Upstream's new non-row-contiguous path in `mlx/backend/gpu/primitives.cpp` keeps the temporary copy alive through the active backend's command encoder, but the pin-refresh overlay preserved upstream's CUDA-or-Metal dispatch verbatim. A ROCm build therefore included `mlx/backend/metal/device.h` and failed on Linux at `` before compiling any mlxcel Rust changes. The overlay now includes the ROCm device header and retains the copy through `rocm::get_command_encoder`; CUDA and Metal keep their upstream branches. Retarget compatibility fix for `bff10220`. diff --git a/src/lib/mlxcelverse/overlays/rocm/mlx/backend/gpu/primitives.cpp b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/gpu/primitives.cpp index 6c4d10520..5c6476121 100644 --- a/src/lib/mlxcelverse/overlays/rocm/mlx/backend/gpu/primitives.cpp +++ b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/gpu/primitives.cpp @@ -9,6 +9,8 @@ #if defined(MLX_USE_CUDA) #include #include "mlx/backend/cuda/device.h" +#elif defined(MLX_USE_ROCM) +#include "mlx/backend/rocm/device.h" #else #include "mlx/backend/metal/device.h" #endif @@ -32,6 +34,8 @@ void AsStrided::eval_gpu(const std::vector& inputs, array& out) { auto copy = contiguous_copy_gpu(inputs[0], stream()); #if defined(MLX_USE_CUDA) cu::get_command_encoder(stream()).add_temporary(copy); +#elif defined(MLX_USE_ROCM) + rocm::get_command_encoder(stream()).add_temporary(copy); #else metal::get_command_encoder(stream()).add_temporary(copy); #endif From aee9d0564be3e52e2b4964d25b8f4c7f3f7e783c Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 23 Oct 2026 00:12:46 +0900 Subject: [PATCH 2/5] feat: add dense prediction output foundation Add the `predict` CLI and shared dense predictor, interpolation, preprocessing, registry, and file-output contracts so upcoming model ports have one integrated path. Bound image and video decoding, stream safetensors tensor materialization, and safely replace only writer-owned output artifacts. Validate the ROCm compile, focused dense, CLI, and registry tests including real MP4 decoding, clippy with warnings denied, and formatting. Refs #2349 --- README.md | 1 + docs/README.md | 1 + docs/dense-prediction.md | 87 +++ src/commands/mod.rs | 2 + src/commands/predict.rs | 216 ++++++++ src/commands/predict_tests.rs | 292 ++++++++++ src/main.rs | 60 +++ src/main_tests.rs | 56 ++ src/models/detection.rs | 6 + src/models/detection_tests.rs | 16 + src/models/registry.rs | 32 ++ src/multimodal/video.rs | 63 ++- src/multimodal/video_tests.rs | 25 + src/server/webui/catalog_metadata.rs | 2 +- src/server/webui/catalog_reason_tests.rs | 34 +- src/vision/dense/interpolate.rs | 360 +++++++++++++ src/vision/dense/interpolate_tests.rs | 299 +++++++++++ src/vision/dense/mod.rs | 25 + src/vision/dense/output.rs | 497 ++++++++++++++++++ src/vision/dense/output_tests.rs | 314 +++++++++++ src/vision/dense/predictor.rs | 75 +++ src/vision/dense/preprocess.rs | 134 +++++ src/vision/dense/registry.rs | 83 +++ src/vision/dense/registry_tests.rs | 61 +++ src/vision/mod.rs | 1 + .../generate_dense_interpolate_fixtures.py | 58 ++ 26 files changed, 2796 insertions(+), 4 deletions(-) create mode 100644 docs/dense-prediction.md create mode 100644 src/commands/predict.rs create mode 100644 src/commands/predict_tests.rs create mode 100644 src/vision/dense/interpolate.rs create mode 100644 src/vision/dense/interpolate_tests.rs create mode 100644 src/vision/dense/mod.rs create mode 100644 src/vision/dense/output.rs create mode 100644 src/vision/dense/output_tests.rs create mode 100644 src/vision/dense/predictor.rs create mode 100644 src/vision/dense/preprocess.rs create mode 100644 src/vision/dense/registry.rs create mode 100644 src/vision/dense/registry_tests.rs create mode 100644 tests/fixtures/generate_dense_interpolate_fixtures.py diff --git a/README.md b/README.md index 96048c045..fd63596d1 100644 --- a/README.md +++ b/README.md @@ -356,6 +356,7 @@ Install with `pip install ./python`. See [Python client](docs/python-client.md) - [OpenAI Responses API](docs/responses-api.md) - [Embeddings and reranking](docs/embeddings.md) - [Image generation](docs/image-generation.md) +- [Dense prediction maps and keypoints](docs/dense-prediction.md) - [Audio API](docs/audio-api.md) - [Nemotron VoiceChat (full-duplex speech)](docs/nemotron-voicechat.md) - [TurboQuant KV cache](docs/turbo-kv-cache.md) diff --git a/docs/README.md b/docs/README.md index 5d336ddb7..8d1617930 100644 --- a/docs/README.md +++ b/docs/README.md @@ -41,6 +41,7 @@ Current GitHub-facing docs: 25. `webui/README.md` - the bundled browser WebUI (`--webui`): the English and Korean user guide (quickstarts, keys, remote access, models, settings and privacy, troubleshooting), the contributor workflow, and the API, state-machine, UX, and design-system contracts behind it. `webui-integration-matrix.md` is its release evidence record. 26. `model-triage.md` - dispositions for reported model-coverage gaps, each with the evidence it rests on, and what a port of the remaining candidates would add. 27. `vision-model-scope.md` - the scope record for detection, segmentation, 3D body, grounding, and image or video generation families: which are supported, which moved to the image-generation and dense-prediction epics with their port issues, which remain out of scope, and the criteria for re-scoping one. +28. `dense-prediction.md` - the `mlxcel predict` image/video CLI, tensor and preview file contract, shared interpolation/preprocessing APIs, and family-registration guide. ## mlxcelverse maintenance diff --git a/docs/dense-prediction.md b/docs/dense-prediction.md new file mode 100644 index 000000000..60c80a031 --- /dev/null +++ b/docs/dense-prediction.md @@ -0,0 +1,87 @@ +# Dense prediction + +`mlxcel predict` is the offline file interface for checkpoints that return depth, normal, mask, segmentation, point, or keypoint results. These outputs have shapes and lifetimes that do not fit the token stream from `generate` or the box rows from `detect`. + +This release establishes the shared CLI, registry, preprocessing, and file contract. The dense family table is intentionally empty until the model ports in the dense-prediction epic land, so the examples below show the target interface those ports will activate. A recognized checkpoint currently reports that its dense-prediction family is not supported by this build. + +```bash +mlxcel predict \ + --model models/video-depth-anything \ + --video input.mp4 \ + --output outputs/depth \ + --max-frames 32 \ + --target-fps 2 + +mlxcel predict \ + --model models/sapiens2 \ + --image person.png \ + --output outputs/pose \ + --task pose \ + --opts '{"confidence":0.5}' +``` + +Exactly one of `--image` and `--video` is required. `--task` may be omitted when the selected family has one task. `--opts` must be a JSON object; the command adds the selected task under the reserved `task` key before calling the predictor. Video sampling keeps at most 32 frames by default; `--max-frames` changes that cap, and the command rejects a sampled set whose decoded RGBA worst-case size exceeds 512 MiB before ffmpeg decodes it. The output directory is created when absent. Existing files with the standard names are replaced, while unrelated files are retained. + +Every successful run writes: + +| Path | Contents | +|---|---| +| `result.safetensors` | Every output tensor converted to contiguous f32, preserving the predictor's upstream dictionary key and shape. Video maps are stacked on axis 0. | +| `meta.json` | Family, selected task, source size, inference time after MLX evaluation, finite min/max for every tensor, model metadata, and flattened intrinsics when an `intrinsics` tensor is present. | +| `keypoints.json` | Present when people were detected. Each entry contains `bbox`, `(x, y, score)` points in source pixels, and optional point names. | +| `depth.png` | 16-bit grayscale. Finite map min/max become 0 and 65535; NaN and infinity become 0. | +| `normal.png` | 8-bit RGB after `(normal + 1) / 2`, clamped to `[0, 1]`. | +| `mask.png` | 8-bit grayscale, using 0 below or equal to 0.5 and 255 above 0.5. | +| `segmentation.png` | 8-bit class IDs. | + +For video, each available preview is written as `/000000.png`, `/000001.png`, and so on. Preview normalization uses the same whole-tensor range recorded in `meta.json`. + +## Model integration + +A model implements the object-safe interface in `src/vision/dense/predictor.rs`: + +```rust,ignore +pub trait DensePredictor { + fn family(&self) -> &'static str; + fn tasks(&self) -> &[&'static str]; + fn predict( + &self, + input: &DenseInput, + opts: &serde_json::Map, + ) -> anyhow::Result; +} +``` + +`DenseOutput::tensors` is a `BTreeMap>` because `MlxArray` is an opaque C++ bridge type. Keep upstream output names such as `depth`, `points`, `normal`, `mask`, `intrinsics`, `segmentation`, `normals`, `pointmaps`, and `alphas`. `DenseOutput::meta` carries family-specific JSON values that the standard writer preserves. + +Family predictors must attach stage context to internal preprocessing failures with `.context("preprocess")`. The command supplies the outer `predict ` context and separately names model loading and output writing, so the resulting error chain identifies the failing stage. + +Expose a constructor with this signature and add its row to `src/vision/dense/registry.rs`: + +```rust,ignore +fn load(path: &std::path::Path) -> anyhow::Result>; + +DenseFamily { + model_type: "video_depth_anything", + load: crate::vision::video_depth_anything::load, +} +``` + +The registry key must equal the lowercase top-level `model_type` in `config.json`. Add the same model type to `KNOWN_MODEL_TYPES` only when the family is a dense predictor; that list makes accidental `mlxcel generate` use direct the operator to `mlxcel predict` even in builds where the family has not landed. + +Shared preprocessing APIs live in `src/vision/dense`: + +```rust,ignore +resize_bilinear_nhwc(x, (height, width), align_corners, antialias) +resize_bicubic_nhwc(x, Some((height, width)), None, antialias) +resize_bicubic_nhwc(x, None, Some((height_scale, width_scale)), antialias) +resize_nearest_nhwc(x, (height, width)) +image_to_nhwc_f32(image) +normalize_imagenet(x) +resize_to_patch_multiple(x, DINOV2_PATCH_SIZE) +conv_weight_to_channels_last(weight, expected_in_channels) +``` + +The interpolation functions accept rank-4 NHWC arrays and match `torch.nn.functional.interpolate` with `align_corners=false` by default at call sites. Bilinear and bicubic interpolation math returns f32; their upstream-compatible same-size fast path returns a copy in the input dtype, and nearest preserves dtype. For bicubic `scale_factor`, the non-antialiased path retains the supplied factor while the antialiased path derives its scale from the floored output size, matching upstream. The convolution helper recognizes an already-channel-last trailing input-channel axis before applying OIHW to OHWI permutation, preventing a second conversion. A model loader should also use checkpoint-key layout detection for shapes whose channel and kernel axes are ambiguous. + +Mesh and Gaussian export belong to the SAM 3D Objects implementation. Streaming, webcam input, and a server route remain outside this file contract until a client-facing dense-output wire format is defined. diff --git a/src/commands/mod.rs b/src/commands/mod.rs index d0309ea2e..2754202ec 100644 --- a/src/commands/mod.rs +++ b/src/commands/mod.rs @@ -38,6 +38,7 @@ mod generate_vlm; mod generate_voicechat; pub(crate) mod inspect; pub(crate) mod models; +pub(crate) mod predict; pub(crate) mod rerank; pub(crate) mod run; mod serve; @@ -55,6 +56,7 @@ pub(crate) use generate::run_generate; pub(crate) use generate_image::run_generate_image; pub(crate) use inspect::run_inspect; pub(crate) use models::{run_list_local, run_remove}; +pub(crate) use predict::run_predict; pub(crate) use rerank::{RerankArgs, run_rerank}; pub(crate) use run::{RunArgs, run_run}; pub(crate) use serve::run_serve; diff --git a/src/commands/predict.rs b/src/commands/predict.rs new file mode 100644 index 000000000..f514eabff --- /dev/null +++ b/src/commands/predict.rs @@ -0,0 +1,216 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! CLI handler for `mlxcel predict` dense map and keypoint models. + +use std::io::Read; +use std::time::Instant; + +use anyhow::{Context, Result, bail, ensure}; +use serde_json::{Map, Value}; + +use mlxcel::initialize_runtime_checked; +use mlxcel::multimodal::video::{ + DEFAULT_FPS, DEFAULT_MAX_DECODED_VIDEO_BYTES, FrameSamplingPolicy, VideoLimits, VideoSource, + load_video_source_with_policy_and_budget, +}; +use mlxcel::vision::dense::{ + DenseFamily, DenseInput, DenseOutputInfo, FAMILIES, resolve_family, write_dense_output, +}; + +use crate::PredictArgs; + +const DEFAULT_DENSE_MAX_FRAMES: usize = 32; + +fn parse_opts(raw: Option<&str>) -> Result> { + let Some(raw) = raw else { + return Ok(Map::new()); + }; + match serde_json::from_str::(raw).context("parse --opts JSON")? { + Value::Object(map) => Ok(map), + _ => bail!("--opts must be a JSON object"), + } +} + +fn select_task<'a>(requested: Option<&'a str>, tasks: &[&'static str]) -> Result<&'a str> { + ensure!(!tasks.is_empty(), "dense predictor registered no tasks"); + if let Some(task) = requested { + ensure!( + tasks.contains(&task), + "unsupported task `{task}`; choose one of: {}", + tasks.join(", ") + ); + return Ok(task); + } + ensure!( + tasks.len() == 1, + "--task is required; this family supports: {}", + tasks.join(", ") + ); + Ok(tasks[0]) +} + +fn load_predict_video( + path: &std::path::Path, + target_fps: Option, + max_frames: Option, +) -> Result> { + let max_frames = max_frames.unwrap_or(DEFAULT_DENSE_MAX_FRAMES); + ensure!(max_frames > 0, "load video: --max-frames must be positive"); + let source = VideoSource::from_path(path.to_path_buf()); + load_video_source_with_policy_and_budget( + &source, + target_fps, + None, + FrameSamplingPolicy { + min_frames: 1, + max_frames, + frame_factor: 1, + reject_over_max: false, + }, + &VideoLimits::from_env(), + Some(DEFAULT_MAX_DECODED_VIDEO_BYTES), + ) + .with_context(|| format!("load video: {}", path.display())) +} + +fn load_predict_image(path: &std::path::Path) -> Result { + let limits = mlxcel::current_image_input_limits(); + let read_limit = limits + .max_payload_bytes + .checked_add(1) + .ok_or_else(|| anyhow::anyhow!("load image: configured payload limit overflowed"))?; + let mut bytes = Vec::new(); + std::fs::File::open(path) + .with_context(|| format!("load image: open {}", path.display()))? + .take(read_limit as u64) + .read_to_end(&mut bytes) + .with_context(|| format!("load image: read {}", path.display()))?; + ensure!( + bytes.len() <= limits.max_payload_bytes, + "load image: payload {} is {} bytes, exceeding the configured maximum of {}", + path.display(), + bytes.len(), + limits.max_payload_bytes + ); + let mut images = mlxcel::decode_image_payloads_with_limits(&[bytes], limits) + .with_context(|| format!("load image: decode {}", path.display()))?; + let image = images.pop().ok_or_else(|| { + anyhow::anyhow!( + "load image: decoder returned no image for {}", + path.display() + ) + })?; + validate_dense_image_allocation(image.width(), image.height(), limits.max_decode_alloc_bytes)?; + Ok(image) +} + +fn validate_dense_image_allocation(width: u32, height: u32, max_bytes: u64) -> Result<()> { + let expanded_bytes = u64::from(width) + .checked_mul(u64::from(height)) + .and_then(|pixels| pixels.checked_mul(3)) + .and_then(|channels| channels.checked_mul(std::mem::size_of::() as u64)) + .ok_or_else(|| anyhow::anyhow!("load image: dense RGB f32 allocation size overflowed"))?; + ensure!( + expanded_bytes <= max_bytes, + "load image: dense RGB f32 expansion for {width}x{height} requires {expanded_bytes} bytes, exceeding the configured maximum of {max_bytes}" + ); + Ok(()) +} + +pub(crate) fn run_predict(args: PredictArgs) -> Result<()> { + run_predict_with(args, FAMILIES) +} + +pub(crate) fn run_predict_with(args: PredictArgs, families: &[DenseFamily]) -> Result<()> { + ensure!( + args.model.is_dir(), + "load model: directory does not exist: {}", + args.model.display() + ); + let source_count = usize::from(args.image.is_some()) + usize::from(args.video.is_some()); + ensure!( + source_count == 1, + "load input: exactly one of --image or --video is required" + ); + if let Some(fps) = args.target_fps { + ensure!( + fps.is_finite() && fps > 0.0, + "load video: --target-fps must be finite and positive" + ); + } + let mut opts = parse_opts(args.opts.as_deref())?; + std::fs::create_dir_all(&args.output) + .with_context(|| format!("write output: create {}", args.output.display()))?; + + let family = + resolve_family(&args.model, families).context("load model: classify checkpoint")?; + let input = if let Some(path) = &args.image { + ensure!( + path.is_file(), + "load image: file does not exist: {}", + path.display() + ); + DenseInput::Image(load_predict_image(path)?) + } else { + let path = args + .video + .as_ref() + .ok_or_else(|| anyhow::anyhow!("load video: --video is required without --image"))?; + ensure!( + path.is_file(), + "load video: file does not exist: {}", + path.display() + ); + let frames = load_predict_video(path, args.target_fps, args.max_frames)?; + DenseInput::Video { + frames, + fps: args.target_fps.unwrap_or(DEFAULT_FPS), + } + }; + let input_size = input + .dimensions() + .ok_or_else(|| anyhow::anyhow!("load input: video decoded no frames"))?; + let is_video = input.is_video(); + + let _runtime = initialize_runtime_checked().context("load model: initialize runtime")?; + let predictor = (family.load)(&args.model).context("load model: construct predictor")?; + let task = select_task(args.task.as_deref(), predictor.tasks())?; + opts.insert("task".into(), Value::String(task.into())); + + let started = Instant::now(); + let output = predictor + .predict(&input, &opts) + .with_context(|| format!("predict {task}"))?; + for tensor in output.tensors.values() { + mlxcel_core::eval(tensor); + } + let elapsed_ms = started.elapsed().as_secs_f64() * 1000.0; + let info = DenseOutputInfo { + family: predictor.family(), + task, + input_size, + elapsed_ms, + is_video, + }; + let paths = write_dense_output(&args.output, &output, &info).context("write output")?; + for path in paths { + println!("saved {}", path.display()); + } + Ok(()) +} + +#[cfg(test)] +#[path = "predict_tests.rs"] +mod tests; diff --git a/src/commands/predict_tests.rs b/src/commands/predict_tests.rs new file mode 100644 index 000000000..a5c883f46 --- /dev/null +++ b/src/commands/predict_tests.rs @@ -0,0 +1,292 @@ +// Copyright 2025-2026 Lablup Inc. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::BTreeMap; +use std::ffi::OsString; +use std::path::{Path, PathBuf}; +use std::process::{Command, Stdio}; + +use clap::Parser; +use mlxcel::vision::dense::{DenseOutput, DensePredictor, PersonKeypoints}; + +use super::*; +use crate::{Cli, Commands}; + +struct Stub; + +impl DensePredictor for Stub { + fn family(&self) -> &'static str { + "sapiens2" + } + + fn tasks(&self) -> &[&'static str] { + &["depth", "pose"] + } + + fn predict(&self, input: &DenseInput, opts: &Map) -> Result { + let task = opts.get("task").and_then(Value::as_str).unwrap(); + let (tensors, keypoints) = match input { + DenseInput::Image(image) => { + use image::GenericImageView; + let (width, height) = image.dimensions(); + let n = width as usize * height as usize; + let depth = mlxcel_core::from_slice_f32( + &(0..n).map(|value| value as f32).collect::>(), + &[height as i32, width as i32], + ); + let mut tensors = BTreeMap::new(); + tensors.insert("depth".into(), depth); + let keypoints = if task == "pose" { + vec![PersonKeypoints { + bbox: [0.0, 0.0, width as f32, height as f32], + points: vec![[1.0, 1.0, 1.0]], + names: None, + }] + } else { + Vec::new() + }; + (tensors, keypoints) + } + DenseInput::Video { frames, .. } => { + let width = frames[0].width() as usize; + let height = frames[0].height() as usize; + let n = frames.len() * width * height; + let depth = mlxcel_core::from_slice_f32( + &(0..n).map(|value| value as f32).collect::>(), + &[frames.len() as i32, height as i32, width as i32], + ); + let mut tensors = BTreeMap::new(); + tensors.insert("depth".into(), depth); + (tensors, Vec::new()) + } + }; + Ok(DenseOutput { + tensors, + keypoints, + meta: Map::new(), + }) + } +} + +fn load_stub(_: &Path) -> Result> { + Ok(Box::new(Stub)) +} + +const STUB_FAMILIES: &[DenseFamily] = &[DenseFamily { + model_type: "sapiens2", + load: load_stub, +}]; + +fn png_chunk(name: &[u8; 4], data: &[u8]) -> Vec { + let mut out = Vec::with_capacity(12 + data.len()); + out.extend_from_slice(&(data.len() as u32).to_be_bytes()); + out.extend_from_slice(name); + out.extend_from_slice(data); + out.extend_from_slice(&crc32(name, data).to_be_bytes()); + out +} + +fn crc32(name: &[u8; 4], data: &[u8]) -> u32 { + let mut crc = u32::MAX; + for &byte in name.iter().chain(data) { + crc ^= u32::from(byte); + for _ in 0..8 { + let mask = 0u32.wrapping_sub(crc & 1); + crc = (crc >> 1) ^ (0xedb8_8320 & mask); + } + } + !crc +} + +fn synthetic_png_header(width: u32, height: u32) -> Vec { + let mut ihdr = Vec::with_capacity(13); + ihdr.extend_from_slice(&width.to_be_bytes()); + ihdr.extend_from_slice(&height.to_be_bytes()); + ihdr.extend_from_slice(&[8, 2, 0, 0, 0]); + let mut png = b"\x89PNG\r\n\x1a\n".to_vec(); + png.extend_from_slice(&png_chunk(b"IHDR", &ihdr)); + png.extend_from_slice(&png_chunk(b"IEND", &[])); + png +} + +fn model_dir() -> tempfile::TempDir { + let dir = tempfile::TempDir::new().unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type":"sapiens2"}"#, + ) + .unwrap(); + dir +} + +fn args(model: &Path, output: PathBuf) -> PredictArgs { + PredictArgs { + model: model.to_path_buf(), + image: None, + video: None, + output, + task: Some("depth".into()), + max_frames: None, + target_fps: None, + opts: None, + } +} + +fn parse_predict(argv: Vec) -> PredictArgs { + let command = std::thread::Builder::new() + .stack_size(8 * 1024 * 1024) + .spawn(move || Cli::try_parse_from(argv).map(|cli| cli.command)) + .unwrap() + .join() + .unwrap() + .unwrap(); + match command { + Commands::Predict(args) => args, + other => panic!("expected predict command, got {other:?}"), + } +} + +#[test] +fn image_runs_through_stub_family_and_writes_contract() { + let model = model_dir(); + let scratch = tempfile::TempDir::new().unwrap(); + let image_path = scratch.path().join("input.png"); + image::DynamicImage::new_rgb8(3, 2) + .save(&image_path) + .unwrap(); + let output = scratch.path().join("nested/output"); + let args = parse_predict(vec![ + "mlxcel".into(), + "predict".into(), + "--model".into(), + model.path().as_os_str().to_owned(), + "--image".into(), + image_path.as_os_str().to_owned(), + "--output".into(), + output.as_os_str().to_owned(), + "--task".into(), + "pose".into(), + "--opts".into(), + r#"{"threshold":0.5}"#.into(), + ]); + + run_predict_with(args, STUB_FAMILIES).unwrap(); + for name in [ + "result.safetensors", + "meta.json", + "keypoints.json", + "depth.png", + ] { + assert!(output.join(name).is_file(), "missing {name}"); + } +} + +#[test] +fn exactly_one_input_and_task_are_validated() { + let model = model_dir(); + let scratch = tempfile::TempDir::new().unwrap(); + let no_input = args(model.path(), scratch.path().join("none")); + assert!( + run_predict_with(no_input, STUB_FAMILIES) + .unwrap_err() + .to_string() + .contains("exactly one") + ); + + let image_path = scratch.path().join("input.png"); + image::DynamicImage::new_rgb8(2, 2) + .save(&image_path) + .unwrap(); + let mut missing_task = args(model.path(), scratch.path().join("missing-task")); + missing_task.image = Some(image_path); + missing_task.task = None; + assert!( + run_predict_with(missing_task, STUB_FAMILIES) + .unwrap_err() + .to_string() + .contains("--task is required") + ); +} + +#[test] +fn image_loader_rejects_an_oversized_header_before_allocation() { + let scratch = tempfile::TempDir::new().unwrap(); + let path = scratch.path().join("oversized.png"); + std::fs::write(&path, synthetic_png_header(100_000, 100_000)).unwrap(); + let error = load_predict_image(&path).unwrap_err(); + assert!( + format!("{error:#}").contains("Image decode rejected by configured limits"), + "{error:#}" + ); +} + +#[test] +fn image_loader_accounts_for_dense_rgb_f32_expansion() { + let limit = 512 * 1024 * 1024; + validate_dense_image_allocation(4096, 4096, limit).unwrap(); + let error = validate_dense_image_allocation(8192, 8192, limit) + .unwrap_err() + .to_string(); + assert!(error.contains("dense RGB f32 expansion"), "{error}"); + assert!(error.contains("805306368"), "{error}"); +} + +#[test] +fn video_cli_honors_fps_and_true_max_frame_cap() { + if !mlxcel::multimodal::video::ffmpeg_available() { + eprintln!("skipping: ffmpeg and ffprobe are required"); + return; + } + let model = model_dir(); + let scratch = tempfile::TempDir::new().unwrap(); + let video = scratch.path().join("input.mp4"); + let status = Command::new("ffmpeg") + .args([ + "-hide_banner", + "-loglevel", + "error", + "-f", + "lavfi", + "-i", + "color=c=black:s=4x4:r=2:d=2", + "-pix_fmt", + "yuv420p", + "-y", + ]) + .arg(&video) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .unwrap(); + assert!(status.success()); + + let output = scratch.path().join("video-output"); + let args = parse_predict(vec![ + "mlxcel".into(), + "predict".into(), + "--model".into(), + model.path().as_os_str().to_owned(), + "--video".into(), + video.as_os_str().to_owned(), + "--output".into(), + output.as_os_str().to_owned(), + "--task".into(), + "depth".into(), + "--max-frames".into(), + "3".into(), + "--target-fps".into(), + "2".into(), + ]); + run_predict_with(args, STUB_FAMILIES).unwrap(); + assert!(output.join("result.safetensors").is_file()); + assert!(output.join("depth/000000.png").is_file()); + assert!(output.join("depth/000002.png").is_file()); + assert!(!output.join("depth/000003.png").exists()); + + let short_cap = load_predict_video(&video, Some(2.0), Some(10)).unwrap(); + assert_eq!(short_cap.len(), 4, "a maximum above the clip must not fail"); + + let one = load_predict_video(&video, Some(2.0), Some(1)).unwrap(); + assert_eq!(one.len(), 1, "a one-frame cap must remain one frame"); +} diff --git a/src/main.rs b/src/main.rs index 438d86fe4..3f5956a5b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -211,6 +211,18 @@ enum Commands { #[command(verbatim_doc_comment)] Detect(DetectArgs), + /// Predict dense maps or keypoints from an image or video. + /// + /// Writes f32 tensors to `result.safetensors`, run metadata to + /// `meta.json`, optional `keypoints.json`, and PNG map previews. + /// + /// Examples: + /// + /// mlxcel predict -m models/video-depth --video clip.mp4 -o depth + /// mlxcel predict -m models/sapiens2 --image person.jpg -o pose --task pose + #[command(verbatim_doc_comment)] + Predict(PredictArgs), + /// Generate or edit an image with a Diffusers-style checkpoint. /// /// Runs text-to-image generation (`--task generate`, the default) or @@ -955,6 +967,52 @@ pub(crate) struct DetectArgs { pub(crate) format: commands::detect::OutputFormat, } +/// Arguments for the `mlxcel predict` subcommand. +#[derive(Args, Debug)] +pub(crate) struct PredictArgs { + /// Path to a dense-prediction model directory. + #[arg(short, long, value_name = "PATH")] + pub(crate) model: std::path::PathBuf, + + /// Input image. Conflicts with `--video`. + #[arg( + long, + value_name = "PATH", + conflicts_with = "video", + required_unless_present = "video" + )] + pub(crate) image: Option, + + /// Input video. Conflicts with `--image`. + #[arg( + long, + value_name = "PATH", + conflicts_with = "image", + required_unless_present = "image" + )] + pub(crate) video: Option, + + /// Output directory, created when absent. + #[arg(short, long, value_name = "DIR")] + pub(crate) output: std::path::PathBuf, + + /// Model task. Required when the family supports more than one task. + #[arg(long, value_name = "NAME")] + pub(crate) task: Option, + + /// Maximum sampled video frames (default: 32). + #[arg(long, value_name = "N")] + pub(crate) max_frames: Option, + + /// Target video sampling rate. + #[arg(long, value_name = "FPS")] + pub(crate) target_fps: Option, + + /// Model-specific options as a JSON object. + #[arg(long, value_name = "JSON")] + pub(crate) opts: Option, +} + /// Arguments for the `mlxcel generate-image` subcommand. #[derive(Args, Debug)] pub(crate) struct GenerateImageArgs { @@ -3185,6 +3243,7 @@ fn cli_model_dir(command: &Commands) -> Option<&std::path::Path> { Commands::Tune(args) => args.model.as_deref(), Commands::Inspect(args) => Some(args.model.as_path()), Commands::Detect(args) => Some(args.model.as_path()), + Commands::Predict(args) => Some(args.model.as_path()), Commands::GenerateImage(args) => Some(args.model.as_path()), Commands::List(_) | Commands::Arch(_) | Commands::Download(_) | Commands::Rm(_) => None, // Gated exactly like the variant itself and its dispatch arm below; @@ -3324,6 +3383,7 @@ fn main() -> anyhow::Result<()> { Commands::Inspect(args) => commands::run_inspect(args), Commands::Download(args) => commands::run_download(args), Commands::Detect(args) => commands::run_detect(args), + Commands::Predict(args) => commands::run_predict(args), Commands::GenerateImage(args) => commands::run_generate_image(args), Commands::Embed(args) => commands::run_embed(args), Commands::Rerank(args) => commands::run_rerank(args), diff --git a/src/main_tests.rs b/src/main_tests.rs index 89eb9876c..0498ba7a1 100644 --- a/src/main_tests.rs +++ b/src/main_tests.rs @@ -322,6 +322,24 @@ fn arch_json_output_is_stable_registry_json() { let families = value["families"].as_array().unwrap(); assert_eq!(families.len(), mlxcel::models::ALL_MODEL_TYPES.len() + 1); assert_eq!(families[0]["id"], "llama"); + assert!( + value["runtime_kinds"] + .as_array() + .unwrap() + .contains(&serde_json::json!("predict")) + ); + assert!( + value["output_kinds"] + .as_array() + .unwrap() + .contains(&serde_json::json!("dense_map")) + ); + assert!( + value["output_kinds"] + .as_array() + .unwrap() + .contains(&serde_json::json!("keypoints")) + ); let qwen3 = families .iter() @@ -377,6 +395,44 @@ fn arch_json_output_is_stable_registry_json() { assert_eq!(rt_detr["pipeline_parallel"], false); } +#[test] +fn predict_requires_exactly_one_image_or_video() { + let missing = try_parse_cli(&["mlxcel", "predict", "-m", "model", "-o", "output"]).unwrap_err(); + assert_eq!( + missing.kind(), + clap::error::ErrorKind::MissingRequiredArgument + ); + + let both = try_parse_cli(&[ + "mlxcel", + "predict", + "-m", + "model", + "--image", + "image.png", + "--video", + "video.mp4", + "-o", + "output", + ]) + .unwrap_err(); + assert_eq!(both.kind(), clap::error::ErrorKind::ArgumentConflict); + + let cli = parse_cli(&[ + "mlxcel", + "predict", + "-m", + "model", + "--image", + "image.png", + "-o", + "output", + "--opts", + r#"{"resolution_level":7}"#, + ]); + assert!(matches!(cli.command, Commands::Predict(_))); +} + /// Issue #26: the dead `docs/model_implementations.md` reference was /// removed. Refuse to let it come back. #[test] diff --git a/src/models/detection.rs b/src/models/detection.rs index f4b5e084f..8cc35d252 100644 --- a/src/models/detection.rs +++ b/src/models/detection.rs @@ -839,6 +839,12 @@ pub(crate) fn detect_model_type_with_probes( // same arm as their canonical lowercase form. let model_type = model_type_raw.to_ascii_lowercase(); + if crate::vision::dense::registry::is_dense_model_type(&model_type) { + return Err(anyhow::anyhow!( + "dense prediction checkpoint ({model_type_raw}): use `mlxcel predict`" + )); + } + if matches!(model_type.as_str(), "inkling_mm_model" | "inkling") && probes.inkling_dir_is_mtp_only(model_path, v)? { diff --git a/src/models/detection_tests.rs b/src/models/detection_tests.rs index 39a06c3d3..09845394b 100644 --- a/src/models/detection_tests.rs +++ b/src/models/detection_tests.rs @@ -2203,3 +2203,19 @@ fn flat_muse_glimmer_config_detects_as_the_text_model() { ); assert_eq!(detected, ModelType::MuseGlimmer); } +#[test] +fn dense_checkpoint_redirects_generate_to_predict() { + let dir = tempfile::TempDir::new().unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type":"sapiens2"}"#, + ) + .unwrap(); + let error = super::detection::get_model_type(dir.path()) + .unwrap_err() + .to_string(); + assert_eq!( + error, + "dense prediction checkpoint (sapiens2): use `mlxcel predict`" + ); +} diff --git a/src/models/registry.rs b/src/models/registry.rs index 46e8b7653..c736c5d15 100644 --- a/src/models/registry.rs +++ b/src/models/registry.rs @@ -26,6 +26,7 @@ pub enum Runtime { Asr, Tts, Detect, + Predict, ImageGenerate, } @@ -46,6 +47,8 @@ pub enum OutputKind { Scores, Audio, Boxes, + DenseMap, + Keypoints, Image, } @@ -89,9 +92,36 @@ pub struct FamilyCapabilities { #[derive(Debug, Serialize)] pub struct ArchitectureRegistry { pub mlxcel_version: &'static str, + /// Complete schema vocabulary, including capabilities whose model rows + /// have not landed yet. + pub runtime_kinds: &'static [Runtime], + pub output_kinds: &'static [OutputKind], pub families: Vec, } +const ALL_RUNTIMES: &[Runtime] = &[ + Runtime::Generate, + Runtime::Serve, + Runtime::Embed, + Runtime::Rerank, + Runtime::Asr, + Runtime::Tts, + Runtime::Detect, + Runtime::Predict, + Runtime::ImageGenerate, +]; + +const ALL_OUTPUT_KINDS: &[OutputKind] = &[ + OutputKind::Tokens, + OutputKind::Embeddings, + OutputKind::Scores, + OutputKind::Audio, + OutputKind::Boxes, + OutputKind::DenseMap, + OutputKind::Keypoints, + OutputKind::Image, +]; + #[derive(Debug, Serialize)] pub struct ArchitectureFamily { pub id: &'static str, @@ -675,6 +705,8 @@ pub fn build_architecture_registry(mlxcel_version: &'static str) -> Architecture ArchitectureRegistry { mlxcel_version, + runtime_kinds: ALL_RUNTIMES, + output_kinds: ALL_OUTPUT_KINDS, families, } } diff --git a/src/multimodal/video.rs b/src/multimodal/video.rs index 6a1822f8c..7c900850c 100644 --- a/src/multimodal/video.rs +++ b/src/multimodal/video.rs @@ -150,8 +150,11 @@ pub const INKLING_MAX_VIDEO_CLIPS: usize = INKLING_MAX_VIDEO_PAIRS / 2; /// No more than two unique decoded frames are required per selected pair. pub const INKLING_MAX_DECODED_FRAMES: usize = INKLING_MAX_VIDEO_PAIRS * 2; -/// Conservative request-wide cap for decoded frame storage (RGBA worst case). -pub const INKLING_MAX_DECODED_PIXEL_BYTES: usize = 512 * 1024 * 1024; +/// Conservative per-request cap for decoded frame storage (RGBA worst case). +pub const DEFAULT_MAX_DECODED_VIDEO_BYTES: usize = 512 * 1024 * 1024; + +/// Inkling uses the shared decoded-video storage cap. +pub const INKLING_MAX_DECODED_PIXEL_BYTES: usize = DEFAULT_MAX_DECODED_VIDEO_BYTES; /// One semantic part of Inkling's timestamped video prompt. #[derive(Debug, Clone, PartialEq)] @@ -1353,6 +1356,27 @@ pub fn load_video_source_with_policy( target_nframes: Option, policy: FrameSamplingPolicy, limits: &VideoLimits, +) -> Result, VideoError> { + load_video_source_with_policy_and_budget( + source, + target_fps, + target_nframes, + policy, + limits, + None, + ) +} + +/// Like [`load_video_source_with_policy`], with an optional aggregate budget +/// for decoded frames. The budget uses four bytes per source pixel and is +/// checked after sampling is planned but before ffmpeg starts decoding. +pub fn load_video_source_with_policy_and_budget( + source: &VideoSource, + target_fps: Option, + target_nframes: Option, + policy: FrameSamplingPolicy, + limits: &VideoLimits, + max_decoded_bytes: Option, ) -> Result, VideoError> { if !ffmpeg_available() { return Err(VideoError::FfmpegMissing); @@ -1373,6 +1397,15 @@ pub fn load_video_source_with_policy( }, other => other, })?; + if let Some(max_decoded_bytes) = max_decoded_bytes { + validate_decoded_video_budget( + &canonical, + meta.width, + meta.height, + nframes, + max_decoded_bytes, + )?; + } let indices = uniform_indices(meta.total_frames, nframes); let frames = extract_frames_single_pass(source, &indices, meta.fps, limits)?; @@ -1383,6 +1416,32 @@ pub fn load_video_source_with_policy( Ok(frames) } +fn validate_decoded_video_budget( + path: &Path, + width: u32, + height: u32, + frames: usize, + max_decoded_bytes: usize, +) -> Result { + let decoded_bytes = inkling_decoded_pixel_bytes(width, height, frames).ok_or_else(|| { + VideoError::Extract { + path: path.to_path_buf(), + message: format!( + "decoded frame budget overflow for {frames} frame(s) at {width}x{height}; reduce --max-frames or --target-fps" + ), + } + })?; + if decoded_bytes > max_decoded_bytes { + return Err(VideoError::Extract { + path: path.to_path_buf(), + message: format!( + "{frames} decoded frame(s) at {width}x{height} require {decoded_bytes} bytes, exceeding the {max_decoded_bytes}-byte budget; reduce --max-frames or --target-fps" + ), + }); + } + Ok(decoded_bytes) +} + /// One fd-safe source and its requested sampling rate for Inkling video. pub struct InklingVideoInput<'a> { pub source: &'a VideoSource, diff --git a/src/multimodal/video_tests.rs b/src/multimodal/video_tests.rs index 4f26c206a..941485e0e 100644 --- a/src/multimodal/video_tests.rs +++ b/src/multimodal/video_tests.rs @@ -101,6 +101,31 @@ fn inkling_decode_budget_caps_pixel_bytes_before_decode() { let over_cap = inkling_decoded_pixel_bytes(4096, 4096, 9).unwrap(); assert!(over_cap > INKLING_MAX_DECODED_PIXEL_BYTES); assert_eq!(INKLING_MAX_DECODED_FRAMES, 32); + + assert_eq!( + validate_decoded_video_budget( + Path::new("dense.mp4"), + 4096, + 4096, + 8, + DEFAULT_MAX_DECODED_VIDEO_BYTES, + ) + .unwrap(), + DEFAULT_MAX_DECODED_VIDEO_BYTES + ); + let error = validate_decoded_video_budget( + Path::new("dense.mp4"), + 4096, + 4096, + 9, + DEFAULT_MAX_DECODED_VIDEO_BYTES, + ) + .unwrap_err() + .to_string(); + assert!( + error.contains("reduce --max-frames or --target-fps"), + "{error}" + ); } #[test] diff --git a/src/server/webui/catalog_metadata.rs b/src/server/webui/catalog_metadata.rs index c7cabdd61..5d40780c2 100644 --- a/src/server/webui/catalog_metadata.rs +++ b/src/server/webui/catalog_metadata.rs @@ -355,7 +355,7 @@ fn tasks_for_family(family: &ArchitectureFamily) -> (Vec, Vec push_unique(&mut output, TaskKind::Rerank), Runtime::Asr => push_unique(&mut output, TaskKind::AudioTranscription), Runtime::Tts => push_unique(&mut output, TaskKind::AudioSpeech), - Runtime::Detect => push_unique(&mut input, TaskKind::VisionInput), + Runtime::Detect | Runtime::Predict => push_unique(&mut input, TaskKind::VisionInput), Runtime::ImageGenerate => push_unique(&mut output, TaskKind::ImageGeneration), } } diff --git a/src/server/webui/catalog_reason_tests.rs b/src/server/webui/catalog_reason_tests.rs index 7942c4f17..6ba31c4ba 100644 --- a/src/server/webui/catalog_reason_tests.rs +++ b/src/server/webui/catalog_reason_tests.rs @@ -12,9 +12,14 @@ // See the License for the specific language governing permissions and // limitations under the License. -use super::sanitize_detection_error; +use super::{sanitize_detection_error, tasks_for_family}; use std::path::Path; +use crate::models::registry::{ + ArchitectureFamily, BackendStatus, BackendSupport, OutputKind, Runtime, +}; +use crate::server::webui::catalog::TaskKind; + #[test] fn detection_reason_limit_counts_unicode_scalars_without_splitting_utf8() { for (glyph, width) in [("x", 1), ("한", 3), ("🦀", 4)] { @@ -48,3 +53,30 @@ fn detection_reason_redacts_full_unicode_path_before_applying_limit() { assert!(!output.contains("private-catalog")); assert!(!output.contains("비밀")); } + +#[test] +fn dense_predict_runtime_projects_as_a_provider_gated_vision_input() { + let family = ArchitectureFamily { + id: "stub_dense", + display_name: "Stub dense", + category: "dense", + model_types: vec!["stub_dense"], + aliases: &[], + runtimes: &[Runtime::Predict], + modalities_in: &[], + output: OutputKind::DenseMap, + backends: BackendSupport { + metal: BackendStatus::Supported, + cuda: BackendStatus::Supported, + rocm: BackendStatus::Supported, + }, + tensor_parallel: false, + pipeline_parallel: false, + drafters: &[], + kv_modes: &[], + }; + + let (input, output) = tasks_for_family(&family); + assert_eq!(input, vec![TaskKind::VisionInput]); + assert!(output.is_empty()); +} diff --git a/src/vision/dense/interpolate.rs b/src/vision/dense/interpolate.rs new file mode 100644 index 000000000..b0f2753ae --- /dev/null +++ b/src/vision/dense/interpolate.rs @@ -0,0 +1,360 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! PyTorch-compatible interpolation for channel-last MLX tensors. +//! +//! The coordinate and tap calculations are a direct port of mlx-vlm's +//! `models/interpolate.py`. Tap tables are small host values; gathering and +//! weighted reduction remain MLX graph operations on the input device. +//! +//! Used by: VideoDepthAnything, MoGe3, Sapiens2, Sam3dObjects. + +use anyhow::{Result, bail, ensure}; +use mlxcel_core::{MlxArray, UniquePtr}; + +#[derive(Debug)] +struct Weights1d { + indices: Vec, + weights: Vec, + out: usize, + taps: usize, +} + +fn checked_nhwc_shape(x: &MlxArray) -> Result<[usize; 4]> { + let shape = mlxcel_core::array_shape(x); + ensure!( + shape.len() == 4, + "dense interpolation expects NHWC rank 4, got {shape:?}" + ); + let mut out = [0; 4]; + for (index, dim) in shape.into_iter().enumerate() { + ensure!(dim > 0, "dense interpolation dimensions must be positive"); + out[index] = dim as usize; + } + Ok(out) +} + +fn checked_size(new_size: (usize, usize)) -> Result<(usize, usize)> { + ensure!( + new_size.0 > 0 && new_size.1 > 0, + "interpolation output dimensions must be positive" + ); + ensure!( + new_size.0 <= i32::MAX as usize && new_size.1 <= i32::MAX as usize, + "interpolation output dimensions exceed MLX limits" + ); + Ok(new_size) +} + +fn bilinear_weights_1d( + in_size: usize, + out_size: usize, + align_corners: bool, + antialias: bool, +) -> Weights1d { + let scale = if align_corners { + if out_size > 1 { + (in_size - 1) as f32 / (out_size - 1) as f32 + } else { + 0.0 + } + } else { + in_size as f32 / out_size as f32 + }; + + if antialias && scale > 1.0 { + let taps = (2.0 * scale).ceil() as usize; + let mut indices = Vec::with_capacity(out_size * taps); + let mut weights = Vec::with_capacity(out_size * taps); + for output in 0..out_size { + let src = if align_corners { + output as f32 * scale + } else { + (output as f32 + 0.5) * scale - 0.5 + }; + let lo = (src - scale + 1.0).floor().max(0.0); + let hi = (src + scale).floor().min((in_size - 1) as f32); + let start = weights.len(); + let mut total = 0.0f32; + for tap in 0..taps { + let pos = lo + tap as f32; + let weight = if pos <= hi { + (1.0 - (pos - src).abs() / scale).max(0.0) + } else { + 0.0 + }; + indices.push(pos.min((in_size - 1) as f32) as i32); + weights.push(weight); + total += weight; + } + for weight in &mut weights[start..] { + *weight /= total; + } + } + return Weights1d { + indices, + weights, + out: out_size, + taps, + }; + } + + let taps = 2; + let mut indices = Vec::with_capacity(out_size * taps); + let mut weights = Vec::with_capacity(out_size * taps); + for output in 0..out_size { + let src = if align_corners { + output as f32 * scale + } else { + (output as f32 + 0.5) * scale - 0.5 + } + .max(0.0); + let base = src.floor(); + let frac = src - base; + indices.push((base as i32).clamp(0, in_size as i32 - 1)); + indices.push((base as i32 + 1).clamp(0, in_size as i32 - 1)); + weights.extend_from_slice(&[1.0 - frac, frac]); + } + Weights1d { + indices, + weights, + out: out_size, + taps, + } +} + +fn cubic_weight(value: f32, a: f32) -> f32 { + let at = value.abs(); + let at2 = at * at; + let at3 = at2 * at; + if at <= 1.0 { + (a + 2.0) * at3 - (a + 3.0) * at2 + 1.0 + } else if at < 2.0 { + a * at3 - 5.0 * a * at2 + 8.0 * a * at - 4.0 * a + } else { + 0.0 + } +} + +fn bicubic_weights_1d(in_size: usize, out_size: usize, scale: f32) -> Weights1d { + let taps = 4; + let mut indices = Vec::with_capacity(out_size * taps); + let mut weights = Vec::with_capacity(out_size * taps); + for output in 0..out_size { + let src = (output as f32 + 0.5) / scale - 0.5; + let base = src.floor(); + for offset in -1..=2 { + let pos = base + offset as f32; + indices.push((pos as i32).clamp(0, in_size as i32 - 1)); + weights.push(cubic_weight(src - pos, -0.75)); + } + } + Weights1d { + indices, + weights, + out: out_size, + taps, + } +} + +fn bicubic_aa_weights_1d(in_size: usize, out_size: usize) -> Weights1d { + let scale = in_size as f32 / out_size as f32; + let support = if scale >= 1.0 { 2.0 * scale } else { 2.0 }; + let taps = support.ceil() as usize * 2 + 1; + let invscale = if scale >= 1.0 { 1.0 / scale } else { 1.0 }; + let mut indices = Vec::with_capacity(out_size * taps); + let mut weights = Vec::with_capacity(out_size * taps); + for output in 0..out_size { + let center = scale * (output as f32 + 0.5); + let xmin = (center - support + 0.5).floor().max(0.0); + let xmax = (center + support + 0.5).floor().min(in_size as f32); + let xsize = (xmax - xmin).clamp(0.0, taps as f32); + let start = weights.len(); + let mut total = 0.0f32; + for tap in 0..taps { + let pos = xmin + tap as f32; + let weight = if (tap as f32) < xsize { + cubic_weight((pos - center + 0.5) * invscale, -0.5) + } else { + 0.0 + }; + indices.push(pos.min((in_size - 1) as f32) as i32); + weights.push(weight); + total += weight; + } + if total > 0.0 { + for weight in &mut weights[start..] { + *weight /= total.max(1e-30); + } + } + } + Weights1d { + indices, + weights, + out: out_size, + taps, + } +} + +fn interpolate_axis( + x: &MlxArray, + table: &Weights1d, + axis: i32, + shape_with_taps: &[i32], + weight_shape: &[i32], +) -> UniquePtr { + let indices = mlxcel_core::from_slice_i32(&table.indices, &[table.indices.len() as i32]); + let gathered = mlxcel_core::take(x, &indices, axis); + let gathered = mlxcel_core::reshape(&gathered, shape_with_taps); + let weights = mlxcel_core::from_slice_f32(&table.weights, weight_shape); + mlxcel_core::sum_axis(&mlxcel_core::multiply(&gathered, &weights), axis + 1, false) +} + +fn separable_interpolate( + x: &MlxArray, + input: [usize; 4], + y: &Weights1d, + x_weights: &Weights1d, +) -> UniquePtr { + let [batch, _, width, channels] = input.map(|dim| dim as i32); + let as_f32 = mlxcel_core::astype(x, mlxcel_core::dtype::FLOAT32); + let height = interpolate_axis( + &as_f32, + y, + 1, + &[batch, y.out as i32, y.taps as i32, width, channels], + &[1, y.out as i32, y.taps as i32, 1, 1], + ); + interpolate_axis( + &height, + x_weights, + 2, + &[ + batch, + y.out as i32, + x_weights.out as i32, + x_weights.taps as i32, + channels, + ], + &[1, 1, x_weights.out as i32, x_weights.taps as i32, 1], + ) +} + +/// Resize `[N, H, W, C]` with PyTorch `bilinear`, `align_corners` and +/// antialias semantics. +/// +/// Interpolation is computed in f32. When the requested spatial size already +/// matches the input, the upstream-compatible identity path returns a copy in +/// the input dtype. +pub fn resize_bilinear_nhwc( + x: &MlxArray, + new_size: (usize, usize), + align_corners: bool, + antialias: bool, +) -> Result> { + let input = checked_nhwc_shape(x)?; + let (out_h, out_w) = checked_size(new_size)?; + if (input[1], input[2]) == (out_h, out_w) { + return Ok(mlxcel_core::copy(x)); + } + let y = bilinear_weights_1d(input[1], out_h, align_corners, antialias); + let x_weights = bilinear_weights_1d(input[2], out_w, align_corners, antialias); + Ok(separable_interpolate(x, input, &y, &x_weights)) +} + +/// Resize `[N, H, W, C]` with PyTorch `bicubic` semantics. +/// +/// Supply exactly one of `new_size` or `scale_factor`. A scale factor floors +/// the output dimensions. Non-antialiased sampling retains the supplied +/// factor; antialiased sampling intentionally derives its scale from the +/// floored output size, matching upstream. As with bilinear resize, the +/// identity-size path preserves the input dtype and other paths produce f32. +pub fn resize_bicubic_nhwc( + x: &MlxArray, + new_size: Option<(usize, usize)>, + scale_factor: Option<(f32, f32)>, + antialias: bool, +) -> Result> { + let input = checked_nhwc_shape(x)?; + let ((out_h, out_w), (scale_h, scale_w)) = match (new_size, scale_factor) { + (Some(size), None) => { + let size = checked_size(size)?; + ( + size, + ( + size.0 as f32 / input[1] as f32, + size.1 as f32 / input[2] as f32, + ), + ) + } + (None, Some(scale)) => { + ensure!( + scale.0.is_finite() && scale.1.is_finite() && scale.0 > 0.0 && scale.1 > 0.0, + "bicubic scale factors must be finite and positive" + ); + let size = checked_size(( + (input[1] as f32 * scale.0) as usize, + (input[2] as f32 * scale.1) as usize, + ))?; + (size, scale) + } + _ => bail!("supply exactly one of bicubic new_size or scale_factor"), + }; + if (input[1], input[2]) == (out_h, out_w) { + return Ok(mlxcel_core::copy(x)); + } + let y = if antialias { + bicubic_aa_weights_1d(input[1], out_h) + } else { + bicubic_weights_1d(input[1], out_h, scale_h) + }; + let x_weights = if antialias { + bicubic_aa_weights_1d(input[2], out_w) + } else { + bicubic_weights_1d(input[2], out_w, scale_w) + }; + Ok(separable_interpolate(x, input, &y, &x_weights)) +} + +/// Resize `[N, H, W, C]` with PyTorch's `nearest` source-index rule. +pub fn resize_nearest_nhwc(x: &MlxArray, new_size: (usize, usize)) -> Result> { + let input = checked_nhwc_shape(x)?; + let (out_h, out_w) = checked_size(new_size)?; + let mut output = mlxcel_core::copy(x); + if input[1] != out_h { + let indices: Vec = (0..out_h) + .map(|index| { + ((index as f32 * input[1] as f32 / out_h as f32).floor() as i32) + .min(input[1] as i32 - 1) + }) + .collect(); + let indices = mlxcel_core::from_slice_i32(&indices, &[out_h as i32]); + output = mlxcel_core::take(&output, &indices, 1); + } + if input[2] != out_w { + let indices: Vec = (0..out_w) + .map(|index| { + ((index as f32 * input[2] as f32 / out_w as f32).floor() as i32) + .min(input[2] as i32 - 1) + }) + .collect(); + let indices = mlxcel_core::from_slice_i32(&indices, &[out_w as i32]); + output = mlxcel_core::take(&output, &indices, 2); + } + Ok(output) +} + +#[cfg(test)] +#[path = "interpolate_tests.rs"] +mod tests; diff --git a/src/vision/dense/interpolate_tests.rs b/src/vision/dense/interpolate_tests.rs new file mode 100644 index 000000000..0d50b5769 --- /dev/null +++ b/src/vision/dense/interpolate_tests.rs @@ -0,0 +1,299 @@ +// Copyright 2025-2026 Lablup Inc. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; + +fn input(height: usize, width: usize) -> UniquePtr { + let values: Vec = (0..height * width) + .map(|index| (((index * 37) % 101) as i32 - 50) as f32 / 17.0) + .collect(); + mlxcel_core::from_slice_f32(&values, &[1, height as i32, width as i32, 1]) +} + +fn assert_probes(actual: &MlxArray, expected: &[(usize, f32)]) { + let values = mlxcel_core::utils::array_to_vec_f32(actual); + for &(index, expected) in expected { + let actual = values[index]; + assert!( + (actual - expected).abs() <= 4e-5, + "index {index}: expected {expected}, got {actual}" + ); + } +} + +fn assert_values_close(left: &MlxArray, right: &MlxArray) { + let left = mlxcel_core::utils::array_to_vec_f32(left); + let right = mlxcel_core::utils::array_to_vec_f32(right); + assert_eq!(left.len(), right.len()); + for (index, (left, right)) in left.into_iter().zip(right).enumerate() { + assert!( + (left - right).abs() <= 4e-5, + "index {index}: expected {right}, got {left}" + ); + } +} + +// Literals were generated once with tests/fixtures/generate_dense_interpolate_fixtures.py, +// which runs torch.nn.functional.interpolate on the fixed inputs below. +#[test] +fn upstream_bilinear_literals_match_in_both_directions() { + let a = input(7, 9); + assert_probes( + &resize_bilinear_nhwc(&a, (14, 5), false, false).unwrap(), + &[ + (0, -2.07058811), + (1, 0.658824205), + (4, 1.71764588), + (10, -0.747058868), + (23, 0.888234735), + (35, 0.397058904), + (65, 0.200000048), + (68, 2.4470582), + (69, 0.423528194), + ], + ); + assert_probes( + &resize_bilinear_nhwc(&a, (14, 5), false, true).unwrap(), + &[ + (0, -1.69747913), + (1, -0.172793865), + (4, 1.3445369), + (10, -0.692226946), + (23, 0.2159926), + (35, 0.324579865), + (65, -0.0210083164), + (68, 1.49632359), + (69, 0.474789947), + ], + ); + + let b = input(14, 5); + assert_probes( + &resize_bilinear_nhwc(&b, (7, 9), false, false).unwrap(), + &[ + (0, -0.470588207), + (1, -0.735294104), + (4, 0.911764741), + (9, 0.5), + (21, 0.323529392), + (31, 0.852941155), + (58, 0.794117689), + (61, -0.529412031), + (62, -0.794117689), + ], + ); + assert_probes( + &resize_bilinear_nhwc(&b, (7, 9), false, true).unwrap(), + &[ + (0, -0.260504246), + (1, -0.666666746), + (4, 0.69747901), + (9, 0.5), + (21, 0.323529303), + (31, 0.852941215), + (58, 1.00840342), + (61, -0.456582785), + (62, -0.579831958), + ], + ); +} + +#[test] +fn upstream_bicubic_size_literals_match_in_both_directions() { + let a = input(7, 9); + assert_probes( + &resize_bicubic_nhwc(&a, Some((14, 5)), None, false).unwrap(), + &[ + (0, -2.48464966), + (1, 1.29624152), + (4, 2.09783149), + (10, -0.743794441), + (23, 1.64457595), + (35, 0.132367209), + (65, -0.01364756), + (68, 3.11410022), + (69, 0.364709646), + ], + ); + assert_probes( + &resize_bicubic_nhwc(&a, Some((14, 5)), None, true).unwrap(), + &[ + (0, -1.88656294), + (1, -0.0548147932), + (4, 1.50378382), + (10, -0.717688322), + (23, 0.348062992), + (35, 0.174739793), + (65, -0.221480012), + (68, 1.66238093), + (69, 0.54034245), + ], + ); + + let b = input(14, 5); + assert_probes( + &resize_bicubic_nhwc(&b, Some((7, 9)), None, false).unwrap(), + &[ + (0, -0.296773136), + (1, -0.74775368), + (4, 1.00551474), + (9, 0.580065489), + (21, 0.675555229), + (31, 0.852941155), + (58, 0.700367689), + (61, -0.681985795), + (62, -0.967932999), + ], + ); + assert_probes( + &resize_bicubic_nhwc(&b, Some((7, 9)), None, true).unwrap(), + &[ + (0, -0.197107658), + (1, -0.748206973), + (4, 0.867831647), + (9, 0.756815016), + (21, 0.693133593), + (31, 1.13143373), + (58, 1.13635254), + (61, -0.634362221), + (62, -0.797623694), + ], + ); +} + +#[test] +fn upstream_bicubic_scale_factor_literals_match_in_both_directions() { + let a = input(7, 9); + assert_probes( + &resize_bicubic_nhwc(&a, None, Some((2.0, 0.6)), false).unwrap(), + &[ + (0, -2.59783769), + (1, 1.85225189), + (4, 0.225638911), + (10, -0.942854404), + (23, 0.0419154689), + (35, 0.404766262), + (65, 0.353826582), + (68, 2.09779048), + (69, -0.69623214), + ], + ); + assert_probes( + &resize_bicubic_nhwc(&a, None, Some((2.0, 0.6)), true).unwrap(), + &[ + (0, -1.88656294), + (1, -0.0548147932), + (4, 1.50378382), + (10, -0.717688322), + (23, 0.348062992), + (35, 0.174739793), + (65, -0.221480012), + (68, 1.66238093), + (69, 0.54034245), + ], + ); + + let b = input(14, 5); + assert_probes( + &resize_bicubic_nhwc(&b, None, Some((0.5, 1.8)), false).unwrap(), + &[ + (0, -0.296773136), + (1, -0.74775368), + (4, 1.00551474), + (9, 0.580065489), + (21, 0.675555229), + (31, 0.852941155), + (58, 0.700367689), + (61, -0.681985795), + (62, -0.967932999), + ], + ); + assert_probes( + &resize_bicubic_nhwc(&b, None, Some((0.5, 1.8)), true).unwrap(), + &[ + (0, -0.197107658), + (1, -0.748206973), + (4, 0.867831647), + (9, 0.756815016), + (21, 0.693133593), + (31, 1.13143373), + (58, 1.13635254), + (61, -0.634362221), + (62, -0.797623694), + ], + ); +} + +#[test] +fn antialiased_bicubic_uses_the_floored_output_size_scale() { + let source = input(7, 9); + let from_scale = resize_bicubic_nhwc(&source, None, Some((1.3, 0.72)), true).unwrap(); + let from_size = resize_bicubic_nhwc(&source, Some((9, 6)), None, true).unwrap(); + assert_values_close(&from_scale, &from_size); + + // The non-antialiased upstream path retains the supplied scale, so it is + // observably different from recomputing the effective 9/7 and 6/9 scales. + let from_scale = resize_bicubic_nhwc(&source, None, Some((1.3, 0.72)), false).unwrap(); + let from_size = resize_bicubic_nhwc(&source, Some((9, 6)), None, false).unwrap(); + let from_scale = mlxcel_core::utils::array_to_vec_f32(&from_scale); + let from_size = mlxcel_core::utils::array_to_vec_f32(&from_size); + assert!( + from_scale + .iter() + .zip(from_size) + .any(|(left, right)| (left - right).abs() > 1e-4) + ); +} + +#[test] +fn identity_size_paths_preserve_dtype_even_for_a_nonunit_scale() { + let source = mlxcel_core::astype(&input(7, 9), mlxcel_core::dtype::FLOAT16); + + let bilinear = resize_bilinear_nhwc(&source, (7, 9), false, true).unwrap(); + assert_eq!( + mlxcel_core::array_dtype(&bilinear), + mlxcel_core::dtype::FLOAT16 + ); + assert_values_close(&bilinear, &source); + + // Both products floor back to the original dimensions. Upstream returns + // the tensor before consulting the antialias weight path. + let bicubic = resize_bicubic_nhwc(&source, None, Some((1.01, 1.01)), true).unwrap(); + assert_eq!( + mlxcel_core::array_dtype(&bicubic), + mlxcel_core::dtype::FLOAT16 + ); + assert_values_close(&bicubic, &source); +} + +#[test] +fn upstream_nearest_literals_match_in_both_directions() { + assert_probes( + &resize_nearest_nhwc(&input(7, 9), (14, 5)).unwrap(), + &[ + (0, -2.94117641), + (1, -0.764705896), + (4, 0.411764711), + (10, -1.17647064), + (23, -0.411764711), + (35, 2.35294127), + (65, 1.70588231), + (68, 0.70588237), + (69, -0.882352948), + ], + ); + assert_probes( + &resize_nearest_nhwc(&input(14, 5), (7, 9)).unwrap(), + &[ + (0, -2.94117641), + (1, -2.94117641), + (4, 1.41176474), + (9, 1.0), + (21, 1.17647064), + (31, 1.35294116), + (58, 1.29411769), + (61, -2.47058821), + (62, -0.294117659), + ], + ); +} diff --git a/src/vision/dense/mod.rs b/src/vision/dense/mod.rs new file mode 100644 index 000000000..1f25832ce --- /dev/null +++ b/src/vision/dense/mod.rs @@ -0,0 +1,25 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Shared infrastructure for models that produce maps or keypoints. + +pub mod interpolate; +pub mod output; +pub mod predictor; +pub mod preprocess; +pub mod registry; + +pub use output::{DenseOutputInfo, write_dense_output}; +pub use predictor::{DenseInput, DenseOutput, DensePredictor, PersonKeypoints}; +pub use registry::{DenseFamily, FAMILIES, load_dense_model, resolve_family}; diff --git a/src/vision/dense/output.rs b/src/vision/dense/output.rs new file mode 100644 index 000000000..efeea399d --- /dev/null +++ b/src/vision/dense/output.rs @@ -0,0 +1,497 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Dense tensor, metadata, keypoint, and preview writers. + +use std::borrow::Cow; +use std::collections::{BTreeMap, HashMap}; +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result, bail, ensure}; +use image::{ImageBuffer, Luma, Rgb}; +use safetensors::{Dtype, View}; +use serde_json::{Map, Value, json}; + +use super::predictor::DenseOutput; + +/// Standard metadata supplied by the CLI to the output writer. +#[derive(Debug, Clone)] +pub struct DenseOutputInfo<'a> { + pub family: &'a str, + pub task: &'a str, + pub input_size: (u32, u32), + pub elapsed_ms: f64, + pub is_video: bool, +} + +struct MlxF32View<'a> { + array: &'a mlxcel_core::MlxArray, + shape: Vec, + byte_len: usize, +} + +impl View for MlxF32View<'_> { + fn dtype(&self) -> Dtype { + Dtype::F32 + } + + fn shape(&self) -> &[usize] { + &self.shape + } + + fn data(&self) -> Cow<'_, [u8]> { + Cow::Owned( + mlxcel_core::utils::array_to_vec_f32(self.array) + .into_iter() + .flat_map(f32::to_le_bytes) + .collect(), + ) + } + + fn data_len(&self) -> usize { + self.byte_len + } +} + +#[derive(Clone, Copy)] +struct MapShape { + frames: usize, + height: usize, + width: usize, + channels: usize, + is_video: bool, +} + +fn scalar_map_shape(shape: &[usize], is_video: bool) -> Result { + let parsed = match shape { + [height, width] if !is_video => MapShape { + frames: 1, + height: *height, + width: *width, + channels: 1, + is_video, + }, + [frames, height, width] if is_video => MapShape { + frames: *frames, + height: *height, + width: *width, + channels: 1, + is_video, + }, + [height, width, 1] if !is_video => MapShape { + frames: 1, + height: *height, + width: *width, + channels: 1, + is_video, + }, + [1, height, width] if !is_video => MapShape { + frames: 1, + height: *height, + width: *width, + channels: 1, + is_video, + }, + [frames, height, width, 1] if is_video => MapShape { + frames: *frames, + height: *height, + width: *width, + channels: 1, + is_video, + }, + _ => bail!("unsupported scalar-map shape {shape:?}"), + }; + ensure!( + parsed.frames > 0 && parsed.height > 0 && parsed.width > 0, + "map dimensions must be positive" + ); + Ok(parsed) +} + +fn normal_map_shape(shape: &[usize], is_video: bool) -> Result { + let parsed = match shape { + [height, width, 3] if !is_video => MapShape { + frames: 1, + height: *height, + width: *width, + channels: 3, + is_video, + }, + [frames, height, width, 3] if is_video => MapShape { + frames: *frames, + height: *height, + width: *width, + channels: 3, + is_video, + }, + [1, height, width, 3] if !is_video => MapShape { + frames: 1, + height: *height, + width: *width, + channels: 3, + is_video, + }, + _ => bail!("unsupported normal-map shape {shape:?}"), + }; + ensure!( + parsed.frames > 0 && parsed.height > 0 && parsed.width > 0, + "map dimensions must be positive" + ); + Ok(parsed) +} + +fn finite_range(values: &[f32]) -> Option<(f32, f32)> { + values + .iter() + .copied() + .filter(|value| value.is_finite()) + .fold(None, |range, value| match range { + None => Some((value, value)), + Some((min, max)) => Some((min.min(value), max.max(value))), + }) +} + +fn remove_owned_file(path: &Path) -> Result<()> { + match std::fs::symlink_metadata(path) { + Ok(metadata) if metadata.file_type().is_file() || metadata.file_type().is_symlink() => { + std::fs::remove_file(path) + .with_context(|| format!("remove stale output {}", path.display()))?; + } + Ok(_) => {} + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => return Err(error).with_context(|| format!("inspect {}", path.display())), + } + Ok(()) +} + +fn is_numbered_preview(path: &Path) -> bool { + path.file_name() + .and_then(|name| name.to_str()) + .is_some_and(|name| { + let bytes = name.as_bytes(); + bytes.len() == 10 && &bytes[6..] == b".png" && bytes[..6].iter().all(u8::is_ascii_digit) + }) +} + +fn remove_numbered_previews(path: &Path) -> Result<()> { + let metadata = match std::fs::symlink_metadata(path) { + Ok(metadata) => metadata, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(()), + Err(error) => return Err(error).with_context(|| format!("inspect {}", path.display())), + }; + if metadata.file_type().is_symlink() { + std::fs::remove_file(path) + .with_context(|| format!("remove stale preview symlink {}", path.display()))?; + return Ok(()); + } + if !metadata.file_type().is_dir() { + return Ok(()); + } + for entry in std::fs::read_dir(path) + .with_context(|| format!("read preview directory {}", path.display()))? + { + let entry = entry.with_context(|| format!("read preview entry in {}", path.display()))?; + if is_numbered_preview(&entry.path()) { + remove_owned_file(&entry.path())?; + } + } + if std::fs::read_dir(path) + .with_context(|| format!("re-read preview directory {}", path.display()))? + .next() + .is_none() + { + std::fs::remove_dir(path) + .with_context(|| format!("remove empty preview directory {}", path.display()))?; + } + Ok(()) +} + +fn remove_stale_owned_artifacts(output_dir: &Path) -> Result<()> { + for name in [ + "result.safetensors", + "meta.json", + "keypoints.json", + "depth.png", + "normal.png", + "mask.png", + "segmentation.png", + ] { + remove_owned_file(&output_dir.join(name))?; + } + for name in ["depth", "normal", "normals", "mask", "segmentation"] { + remove_numbered_previews(&output_dir.join(name))?; + } + Ok(()) +} + +fn preview_path(root: &Path, key: &str, frame: usize, shape: MapShape) -> Result { + if !shape.is_video { + let name = match key { + "normal" | "normals" => "normal.png", + other => return Ok(root.join(format!("{other}.png"))), + }; + return Ok(root.join(name)); + } + let dir = root.join(key); + std::fs::create_dir_all(&dir) + .with_context(|| format!("create preview directory {}", dir.display()))?; + Ok(dir.join(format!("{frame:06}.png"))) +} + +fn write_depth_previews( + root: &Path, + key: &str, + values: &[f32], + shape: MapShape, + range: Option<(f32, f32)>, +) -> Result> { + let mut paths = Vec::with_capacity(shape.frames); + let frame_len = shape.height * shape.width; + let (min, max) = range.unwrap_or((0.0, 0.0)); + let span = f64::from(max) - f64::from(min); + for frame in 0..shape.frames { + let pixels = &values[frame * frame_len..(frame + 1) * frame_len]; + let pixels: Vec = pixels + .iter() + .map(|value| { + if !value.is_finite() || span <= 0.0 { + 0 + } else { + (((f64::from(*value) - f64::from(min)) / span).clamp(0.0, 1.0) + * f64::from(u16::MAX)) + .round() as u16 + } + }) + .collect(); + let image = + ImageBuffer::, _>::from_raw(shape.width as u32, shape.height as u32, pixels) + .ok_or_else(|| anyhow::anyhow!("depth preview shape does not match its tensor"))?; + let path = preview_path(root, key, frame, shape)?; + image + .save(&path) + .with_context(|| format!("write depth preview {}", path.display()))?; + paths.push(path); + } + Ok(paths) +} + +fn write_normal_previews( + root: &Path, + key: &str, + values: &[f32], + shape: MapShape, +) -> Result> { + let mut paths = Vec::with_capacity(shape.frames); + let frame_len = shape.height * shape.width * shape.channels; + for frame in 0..shape.frames { + let pixels = &values[frame * frame_len..(frame + 1) * frame_len]; + let pixels: Vec = pixels + .iter() + .map(|value| { + if value.is_finite() { + (((value + 1.0) * 0.5).clamp(0.0, 1.0) * u8::MAX as f32).round() as u8 + } else { + 0 + } + }) + .collect(); + let image = + ImageBuffer::, _>::from_raw(shape.width as u32, shape.height as u32, pixels) + .ok_or_else(|| anyhow::anyhow!("normal preview shape does not match its tensor"))?; + let path = preview_path(root, key, frame, shape)?; + image + .save(&path) + .with_context(|| format!("write normal preview {}", path.display()))?; + paths.push(path); + } + Ok(paths) +} + +fn write_u8_previews( + root: &Path, + key: &str, + values: &[f32], + shape: MapShape, + binary: bool, +) -> Result> { + let mut paths = Vec::with_capacity(shape.frames); + let frame_len = shape.height * shape.width; + for frame in 0..shape.frames { + let pixels = &values[frame * frame_len..(frame + 1) * frame_len]; + let pixels: Vec = pixels + .iter() + .map(|value| { + if !value.is_finite() { + 0 + } else if binary { + u8::MAX * u8::from(*value > 0.5) + } else { + value.round().clamp(0.0, u8::MAX as f32) as u8 + } + }) + .collect(); + let image = + ImageBuffer::, _>::from_raw(shape.width as u32, shape.height as u32, pixels) + .ok_or_else(|| anyhow::anyhow!("map preview shape does not match its tensor"))?; + let path = preview_path(root, key, frame, shape)?; + image + .save(&path) + .with_context(|| format!("write map preview {}", path.display()))?; + paths.push(path); + } + Ok(paths) +} + +/// Write the stable dense-prediction file contract and return every written path. +pub fn write_dense_output( + output_dir: &Path, + output: &DenseOutput, + info: &DenseOutputInfo<'_>, +) -> Result> { + std::fs::create_dir_all(output_dir) + .with_context(|| format!("create output directory {}", output_dir.display()))?; + + // Validate every standard preview shape before replacing any prior run's + // artifacts, so a malformed predictor result leaves the old output intact. + for (name, array) in &output.tensors { + let shape: Vec = mlxcel_core::array_shape(array) + .into_iter() + .map(|dim| dim as usize) + .collect(); + match name.as_str() { + "depth" | "mask" | "segmentation" => { + scalar_map_shape(&shape, info.is_video)?; + } + "normal" | "normals" => { + normal_map_shape(&shape, info.is_video)?; + } + _ => {} + } + } + remove_stale_owned_artifacts(output_dir)?; + + // Preview and summarize one tensor at a time. The safetensors views retain + // only references to MLX arrays; the serializer materializes one tensor's + // bytes at a time instead of retaining a host copy of the full output. + let mut tensors = BTreeMap::new(); + let mut map_ranges = Map::new(); + let mut intrinsics = None; + let mut written = Vec::new(); + for (name, array) in &output.tensors { + let shape: Vec = mlxcel_core::array_shape(array) + .into_iter() + .map(|dim| dim as usize) + .collect(); + let values = mlxcel_core::utils::array_to_vec_f32(array); + ensure!( + values.len() == shape.iter().product::(), + "tensor {name} data length does not match shape {shape:?}" + ); + let byte_len = values + .len() + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| anyhow::anyhow!("tensor {name} byte length overflowed"))?; + let range = finite_range(&values); + map_ranges.insert( + name.clone(), + range.map_or(Value::Null, |(min, max)| json!({ "min": min, "max": max })), + ); + if name == "intrinsics" { + intrinsics = Some(json!({ "shape": &shape, "values": &values })); + } + + let mut preview_paths = match name.as_str() { + "depth" => write_depth_previews( + output_dir, + name, + &values, + scalar_map_shape(&shape, info.is_video)?, + range, + )?, + "normal" | "normals" => write_normal_previews( + output_dir, + name, + &values, + normal_map_shape(&shape, info.is_video)?, + )?, + "mask" => write_u8_previews( + output_dir, + name, + &values, + scalar_map_shape(&shape, info.is_video)?, + true, + )?, + "segmentation" => write_u8_previews( + output_dir, + name, + &values, + scalar_map_shape(&shape, info.is_video)?, + false, + )?, + _ => Vec::new(), + }; + written.append(&mut preview_paths); + + drop(values); + let array = array + .as_ref() + .ok_or_else(|| anyhow::anyhow!("tensor {name} has a null MLX array"))?; + tensors.insert( + name.clone(), + MlxF32View { + array, + shape, + byte_len, + }, + ); + } + + let safetensors_path = output_dir.join("result.safetensors"); + safetensors::serialize_to_file(tensors, None::>, &safetensors_path) + .with_context(|| format!("write {}", safetensors_path.display()))?; + written.insert(0, safetensors_path); + + if !output.keypoints.is_empty() { + let path = output_dir.join("keypoints.json"); + let file = + std::fs::File::create(&path).with_context(|| format!("create {}", path.display()))?; + serde_json::to_writer_pretty(file, &output.keypoints) + .with_context(|| format!("write {}", path.display()))?; + written.push(path); + } + + let mut meta = output.meta.clone(); + meta.insert("family".into(), Value::String(info.family.into())); + meta.insert("task".into(), Value::String(info.task.into())); + meta.insert( + "input_size".into(), + json!({ "width": info.input_size.0, "height": info.input_size.1 }), + ); + meta.insert("maps".into(), Value::Object(map_ranges)); + meta.insert("timing_ms".into(), json!(info.elapsed_ms)); + if let Some(intrinsics) = intrinsics { + meta.insert("intrinsics".into(), intrinsics); + } + let meta_path = output_dir.join("meta.json"); + let file = std::fs::File::create(&meta_path) + .with_context(|| format!("create {}", meta_path.display()))?; + serde_json::to_writer_pretty(file, &meta) + .with_context(|| format!("write {}", meta_path.display()))?; + written.push(meta_path); + Ok(written) +} + +#[cfg(test)] +#[path = "output_tests.rs"] +mod tests; diff --git a/src/vision/dense/output_tests.rs b/src/vision/dense/output_tests.rs new file mode 100644 index 000000000..d98c0e9e6 --- /dev/null +++ b/src/vision/dense/output_tests.rs @@ -0,0 +1,314 @@ +// Copyright 2025-2026 Lablup Inc. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::vision::dense::predictor::PersonKeypoints; + +#[test] +fn synthetic_output_round_trips_tensor_depth_and_keypoints() { + // MLX does not support f64 tensors on GPU streams. Hold the shared device + // lock while this normalization regression exercises an f64 CPU tensor. + let _device_lock = mlxcel_core::streams::lock_default_device(); + let _cpu_device = mlxcel_core::streams::DefaultDeviceGuard::cpu(); + let dir = tempfile::TempDir::new().unwrap(); + let mut output = DenseOutput::default(); + let depth = mlxcel_core::from_slice_f32(&[f32::NAN, 1.0, 2.0, 3.0], &[2, 2]); + output.tensors.insert( + "depth".into(), + mlxcel_core::astype(&depth, mlxcel_core::dtype::FLOAT64), + ); + output.tensors.insert( + "normal".into(), + mlxcel_core::from_slice_f32( + &[ + -1.0, 0.0, 1.0, 1.0, 0.0, -1.0, 0.0, 0.0, 0.0, 0.5, -0.5, 0.0, + ], + &[2, 2, 3], + ), + ); + output.keypoints.push(PersonKeypoints { + bbox: [1.0, 2.0, 3.0, 4.0], + points: vec![[5.0, 6.0, 0.9]], + names: Some(vec!["nose".into()]), + }); + let info = DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (2, 2), + elapsed_ms: 1.25, + is_video: false, + }; + + let paths = write_dense_output(dir.path(), &output, &info).unwrap(); + assert!( + paths + .iter() + .any(|path| path.ends_with("result.safetensors")) + ); + let bytes = std::fs::read(dir.path().join("result.safetensors")).unwrap(); + let tensors = safetensors::SafeTensors::deserialize(&bytes).unwrap(); + let mut names = tensors.names(); + names.sort_unstable(); + assert_eq!(names, vec!["depth", "normal"]); + assert_eq!(tensors.tensor("depth").unwrap().dtype(), Dtype::F32); + assert_eq!(tensors.tensor("depth").unwrap().shape(), &[2, 2]); + assert_eq!(tensors.tensor("normal").unwrap().shape(), &[2, 2, 3]); + + let depth = image::open(dir.path().join("depth.png")).unwrap(); + assert_eq!(depth.color(), image::ColorType::L16); + let depth = depth.to_luma16(); + assert_eq!(depth.get_pixel(0, 0).0[0], 0, "NaN preview is zero"); + assert_eq!(depth.get_pixel(1, 0).0[0], 0, "finite minimum is zero"); + assert_eq!(depth.get_pixel(1, 1).0[0], u16::MAX); + + let meta: Value = + serde_json::from_slice(&std::fs::read(dir.path().join("meta.json")).unwrap()).unwrap(); + assert_eq!(meta["maps"]["depth"]["min"], 1.0); + assert_eq!(meta["maps"]["depth"]["max"], 3.0); + let keypoints: Vec = + serde_json::from_slice(&std::fs::read(dir.path().join("keypoints.json")).unwrap()).unwrap(); + assert_eq!(keypoints, output.keypoints); +} + +#[test] +fn video_depth_writes_one_preview_per_stacked_frame() { + let dir = tempfile::TempDir::new().unwrap(); + let mut output = DenseOutput::default(); + output.tensors.insert( + "depth".into(), + mlxcel_core::from_slice_f32(&[0.0, 1.0, 2.0, 3.0], &[2, 1, 2]), + ); + write_dense_output( + dir.path(), + &output, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (2, 1), + elapsed_ms: 0.0, + is_video: true, + }, + ) + .unwrap(); + assert!(dir.path().join("depth/000000.png").is_file()); + assert!(dir.path().join("depth/000001.png").is_file()); +} + +#[test] +fn one_frame_video_still_uses_a_numbered_preview_directory() { + let dir = tempfile::TempDir::new().unwrap(); + let mut output = DenseOutput::default(); + output.tensors.insert( + "depth".into(), + mlxcel_core::from_slice_f32(&[0.0, 1.0], &[1, 1, 2]), + ); + write_dense_output( + dir.path(), + &output, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (2, 1), + elapsed_ms: 0.0, + is_video: true, + }, + ) + .unwrap(); + assert!(dir.path().join("depth/000000.png").is_file()); + assert!(!dir.path().join("depth.png").exists()); +} + +#[test] +fn video_preview_rejects_an_unstacked_image_map() { + let dir = tempfile::TempDir::new().unwrap(); + let mut output = DenseOutput::default(); + output.tensors.insert( + "depth".into(), + mlxcel_core::from_slice_f32(&[0.0, 1.0, 2.0, 3.0], &[2, 2]), + ); + let error = write_dense_output( + dir.path(), + &output, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (2, 2), + elapsed_ms: 0.0, + is_video: true, + }, + ) + .unwrap_err() + .to_string(); + assert!( + error.contains("unsupported scalar-map shape [2, 2]"), + "{error}" + ); +} + +#[test] +fn reruns_remove_only_stale_writer_owned_artifacts() { + let dir = tempfile::TempDir::new().unwrap(); + std::fs::write(dir.path().join("keep.txt"), "keep").unwrap(); + std::fs::create_dir(dir.path().join("depth")).unwrap(); + std::fs::write(dir.path().join("depth/keep.txt"), "keep").unwrap(); + + let mut image_output = DenseOutput::default(); + image_output.tensors.insert( + "depth".into(), + mlxcel_core::from_slice_f32(&[0.0, 1.0], &[1, 2]), + ); + image_output.keypoints.push(PersonKeypoints { + bbox: [0.0; 4], + points: Vec::new(), + names: None, + }); + write_dense_output( + dir.path(), + &image_output, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (2, 1), + elapsed_ms: 0.0, + is_video: false, + }, + ) + .unwrap(); + assert!(dir.path().join("depth.png").is_file()); + assert!(dir.path().join("keypoints.json").is_file()); + + let mut three_frames = DenseOutput::default(); + three_frames.tensors.insert( + "depth".into(), + mlxcel_core::from_slice_f32(&[0.0, 1.0, 2.0], &[3, 1, 1]), + ); + write_dense_output( + dir.path(), + &three_frames, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (1, 1), + elapsed_ms: 0.0, + is_video: true, + }, + ) + .unwrap(); + assert!(!dir.path().join("depth.png").exists()); + assert!(!dir.path().join("keypoints.json").exists()); + assert!(dir.path().join("depth/000002.png").is_file()); + + let mut one_frame = DenseOutput::default(); + one_frame.tensors.insert( + "depth".into(), + mlxcel_core::from_slice_f32(&[4.0], &[1, 1, 1]), + ); + write_dense_output( + dir.path(), + &one_frame, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (1, 1), + elapsed_ms: 0.0, + is_video: true, + }, + ) + .unwrap(); + assert!(dir.path().join("depth/000000.png").is_file()); + assert!(!dir.path().join("depth/000001.png").exists()); + assert!(!dir.path().join("depth/000002.png").exists()); + + write_dense_output( + dir.path(), + &image_output, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (2, 1), + elapsed_ms: 0.0, + is_video: false, + }, + ) + .unwrap(); + assert!(dir.path().join("depth.png").is_file()); + assert!(!dir.path().join("depth/000000.png").exists()); + assert_eq!( + std::fs::read_to_string(dir.path().join("keep.txt")).unwrap(), + "keep" + ); + assert_eq!( + std::fs::read_to_string(dir.path().join("depth/keep.txt")).unwrap(), + "keep" + ); +} + +#[cfg(unix)] +#[test] +fn video_preview_replaces_a_directory_symlink_without_touching_its_target() { + use std::os::unix::fs::symlink; + + let dir = tempfile::TempDir::new().unwrap(); + let external = tempfile::TempDir::new().unwrap(); + std::fs::write(external.path().join("000000.png"), "external preview").unwrap(); + std::fs::write(external.path().join("keep.txt"), "external sentinel").unwrap(); + std::fs::write( + external.path().join("result.safetensors"), + "external tensors", + ) + .unwrap(); + std::fs::write(external.path().join("meta.json"), "external metadata").unwrap(); + symlink(external.path(), dir.path().join("depth")).unwrap(); + symlink( + external.path().join("result.safetensors"), + dir.path().join("result.safetensors"), + ) + .unwrap(); + symlink( + external.path().join("meta.json"), + dir.path().join("meta.json"), + ) + .unwrap(); + + let mut output = DenseOutput::default(); + output.tensors.insert( + "depth".into(), + mlxcel_core::from_slice_f32(&[0.0, 1.0], &[1, 1, 2]), + ); + write_dense_output( + dir.path(), + &output, + &DenseOutputInfo { + family: "stub", + task: "depth", + input_size: (2, 1), + elapsed_ms: 0.0, + is_video: true, + }, + ) + .unwrap(); + + assert!(dir.path().join("depth").is_dir()); + assert!(!dir.path().join("depth").is_symlink()); + assert!(dir.path().join("depth/000000.png").is_file()); + assert!(dir.path().join("result.safetensors").is_file()); + assert!(!dir.path().join("result.safetensors").is_symlink()); + assert!(dir.path().join("meta.json").is_file()); + assert!(!dir.path().join("meta.json").is_symlink()); + assert_eq!( + std::fs::read_to_string(external.path().join("000000.png")).unwrap(), + "external preview" + ); + assert_eq!( + std::fs::read_to_string(external.path().join("keep.txt")).unwrap(), + "external sentinel" + ); + assert_eq!( + std::fs::read_to_string(external.path().join("result.safetensors")).unwrap(), + "external tensors" + ); + assert_eq!( + std::fs::read_to_string(external.path().join("meta.json")).unwrap(), + "external metadata" + ); +} diff --git a/src/vision/dense/predictor.rs b/src/vision/dense/predictor.rs new file mode 100644 index 000000000..0934fc5b7 --- /dev/null +++ b/src/vision/dense/predictor.rs @@ -0,0 +1,75 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::collections::BTreeMap; + +use anyhow::Result; +use image::DynamicImage; +use mlxcel_core::{MlxArray, UniquePtr}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +/// Image or decoded video passed to a dense-prediction family. +#[derive(Debug)] +pub enum DenseInput { + Image(DynamicImage), + Video { + frames: Vec, + /// Requested sampling rate of `frames`. + fps: f64, + }, +} + +impl DenseInput { + /// Source width and height. Video frames are guaranteed non-empty by the loader. + pub fn dimensions(&self) -> Option<(u32, u32)> { + use image::GenericImageView; + match self { + Self::Image(image) => Some(image.dimensions()), + Self::Video { frames, .. } => frames.first().map(GenericImageView::dimensions), + } + } + + pub fn is_video(&self) -> bool { + matches!(self, Self::Video { .. }) + } +} + +/// One person's keypoints in source-image pixel coordinates. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct PersonKeypoints { + pub bbox: [f32; 4], + /// `(x, y, score)` in source pixels. + pub points: Vec<[f32; 3]>, + #[serde(skip_serializing_if = "Option::is_none")] + pub names: Option>, +} + +/// Model-neutral dense-prediction result. +/// +/// MLX arrays are owned through `UniquePtr` because `MlxArray` is an opaque C++ +/// bridge type. Keys retain the upstream predictor's output-dictionary names. +#[derive(Default)] +pub struct DenseOutput { + pub tensors: BTreeMap>, + pub keypoints: Vec, + pub meta: Map, +} + +/// Common interface implemented by each dense-prediction family. +pub trait DensePredictor { + fn family(&self) -> &'static str; + fn tasks(&self) -> &[&'static str]; + fn predict(&self, input: &DenseInput, opts: &Map) -> Result; +} diff --git a/src/vision/dense/preprocess.rs b/src/vision/dense/preprocess.rs new file mode 100644 index 000000000..6b7a0eca1 --- /dev/null +++ b/src/vision/dense/preprocess.rs @@ -0,0 +1,134 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Shared dense-model image preprocessing. + +use anyhow::{Result, ensure}; +use image::DynamicImage; +use mlxcel_core::{MlxArray, UniquePtr}; + +use super::interpolate::resize_bilinear_nhwc; + +pub const IMAGENET_MEAN: [f32; 3] = [0.485, 0.456, 0.406]; +pub const IMAGENET_STD: [f32; 3] = [0.229, 0.224, 0.225]; +pub const DINOV2_PATCH_SIZE: usize = 14; + +/// Convert an image to `[1, H, W, 3]` RGB f32 in `[0, 1]`. +/// +/// `image`'s RGB32F conversion scales u8 by 255 and u16 by 65535, preserving +/// the integer dtype's full range rather than treating every source as u8. +pub fn image_to_nhwc_f32(image: &DynamicImage) -> UniquePtr { + let rgb = image.to_rgb32f(); + let (width, height) = rgb.dimensions(); + mlxcel_core::from_slice_f32(rgb.as_raw(), &[1, height as i32, width as i32, 3]) +} + +/// Apply ImageNet RGB channel normalization to an NHWC tensor. +pub fn normalize_imagenet(x: &MlxArray) -> Result> { + let shape = mlxcel_core::array_shape(x); + ensure!( + shape.len() == 4 && shape[3] == 3, + "ImageNet normalization expects NHWC RGB, got {shape:?}" + ); + let x = mlxcel_core::astype(x, mlxcel_core::dtype::FLOAT32); + let mean = mlxcel_core::from_slice_f32(&IMAGENET_MEAN, &[1, 1, 1, 3]); + let std = mlxcel_core::from_slice_f32(&IMAGENET_STD, &[1, 1, 1, 3]); + Ok(mlxcel_core::divide(&mlxcel_core::subtract(&x, &mean), &std)) +} + +/// Round a positive spatial size to the nearest positive patch multiple. +pub fn size_to_patch_multiple(size: (usize, usize), patch_size: usize) -> Result<(usize, usize)> { + ensure!(patch_size > 0, "patch size must be positive"); + ensure!( + size.0 > 0 && size.1 > 0, + "image dimensions must be positive" + ); + let align = |value: usize| { + let rounded = ((value + patch_size / 2) / patch_size) * patch_size; + rounded.max(patch_size) + }; + Ok((align(size.0), align(size.1))) +} + +/// Bilinearly resize NHWC input so both spatial dimensions are patch-aligned. +pub fn resize_to_patch_multiple(x: &MlxArray, patch_size: usize) -> Result> { + let shape = mlxcel_core::array_shape(x); + ensure!(shape.len() == 4, "patch resize expects NHWC rank 4"); + let size = size_to_patch_multiple((shape[1] as usize, shape[2] as usize), patch_size)?; + resize_bilinear_nhwc(x, size, false, true) +} + +/// Convert a PyTorch OIHW convolution kernel to MLX OHWI once. +/// +/// A trailing `expected_in_channels` axis proves the input is already OHWI and +/// returns a copy. A matching second axis identifies OIHW. Other layouts are +/// rejected instead of being silently transposed. As with other mlxcel vision +/// sanitizers, a shape where both axes equal `expected_in_channels` is treated +/// as already converted; callers should combine this guard with checkpoint-key +/// layout detection for ambiguous RGB 3x3 kernels. +pub fn conv_weight_to_channels_last( + weight: &MlxArray, + expected_in_channels: i32, +) -> Result> { + let shape = mlxcel_core::array_shape(weight); + ensure!( + shape.len() == 4, + "convolution weight must be rank 4, got {shape:?}" + ); + if shape[3] == expected_in_channels { + return Ok(mlxcel_core::copy(weight)); + } + ensure!( + shape[1] == expected_in_channels, + "convolution weight has no input-channel axis of {expected_in_channels}: {shape:?}" + ); + Ok(mlxcel_core::transpose_axes(weight, &[0, 2, 3, 1])) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn u16_images_scale_by_their_dtype_range() { + let image = DynamicImage::ImageRgb16( + image::ImageBuffer::, Vec>::from_raw( + 1, + 1, + vec![u16::MAX, 32768, 0], + ) + .unwrap(), + ); + let array = image_to_nhwc_f32(&image); + let values = mlxcel_core::utils::array_to_vec_f32(&array); + assert_eq!(values[0], 1.0); + assert!((values[1] - 32768.0 / 65535.0).abs() < 1e-6); + assert_eq!(values[2], 0.0); + } + + #[test] + fn patch_alignment_uses_nearest_multiple_with_one_patch_minimum() { + assert_eq!(size_to_patch_multiple((7, 21), 14).unwrap(), (14, 28)); + assert_eq!(size_to_patch_multiple((28, 35), 14).unwrap(), (28, 42)); + } + + #[test] + fn convolution_permutation_is_idempotent() { + let raw = mlxcel_core::from_slice_f32(&[0.0; 2 * 4 * 3 * 5], &[2, 4, 3, 5]); + let converted = conv_weight_to_channels_last(&raw, 4).unwrap(); + assert_eq!(mlxcel_core::array_shape(&converted), vec![2, 3, 5, 4]); + let twice = conv_weight_to_channels_last(&converted, 4).unwrap(); + assert_eq!(mlxcel_core::array_shape(&twice), vec![2, 3, 5, 4]); + } +} diff --git a/src/vision/dense/registry.rs b/src/vision/dense/registry.rs new file mode 100644 index 000000000..194ee564f --- /dev/null +++ b/src/vision/dense/registry.rs @@ -0,0 +1,83 @@ +// Copyright 2025-2026 Lablup Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Dense-prediction family registry. + +use std::path::Path; + +use anyhow::{Context, Result, bail}; +use serde_json::Value; + +use super::predictor::DensePredictor; + +pub type FamilyLoader = fn(&Path) -> Result>; + +/// One checkpoint `model_type` to predictor-constructor mapping. +#[derive(Debug, Clone, Copy)] +pub struct DenseFamily { + pub model_type: &'static str, + pub load: FamilyLoader, +} + +/// Families this build can execute. Model implementation issues add rows here. +pub static FAMILIES: &[DenseFamily] = &[]; + +/// Dense checkpoint types recognized even before their model ports land. +pub const KNOWN_MODEL_TYPES: &[&str] = + &["moge3", "sam3d_objects", "sapiens2", "video_depth_anything"]; + +pub fn is_dense_model_type(model_type: &str) -> bool { + KNOWN_MODEL_TYPES.contains(&model_type.to_ascii_lowercase().as_str()) +} + +fn checkpoint_model_type(path: &Path) -> Result { + let config_path = path.join("config.json"); + let raw = std::fs::read_to_string(&config_path) + .with_context(|| format!("read {}", config_path.display()))?; + let config: Value = + serde_json::from_str(&raw).with_context(|| format!("parse {}", config_path.display()))?; + config + .get("model_type") + .and_then(Value::as_str) + .map(str::to_ascii_lowercase) + .ok_or_else(|| anyhow::anyhow!("model_type not found in {}", config_path.display())) +} + +pub fn resolve_family<'a>(path: &Path, families: &'a [DenseFamily]) -> Result<&'a DenseFamily> { + let model_type = checkpoint_model_type(path)?; + if !is_dense_model_type(&model_type) { + bail!("checkpoint model_type `{model_type}` is not a recognized dense-prediction family"); + } + families + .iter() + .find(|family| family.model_type == model_type) + .ok_or_else(|| { + anyhow::anyhow!("dense-prediction family {model_type} is not supported by this build") + }) +} + +pub fn load_dense_model(path: &Path) -> Result> { + load_dense_model_from(path, FAMILIES) +} + +pub fn load_dense_model_from( + path: &Path, + families: &[DenseFamily], +) -> Result> { + (resolve_family(path, families)?.load)(path) +} + +#[cfg(test)] +#[path = "registry_tests.rs"] +mod tests; diff --git a/src/vision/dense/registry_tests.rs b/src/vision/dense/registry_tests.rs new file mode 100644 index 000000000..732527c82 --- /dev/null +++ b/src/vision/dense/registry_tests.rs @@ -0,0 +1,61 @@ +// Copyright 2025-2026 Lablup Inc. +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::vision::dense::{DenseInput, DenseOutput, DensePredictor}; +use serde_json::{Map, Value}; + +struct Stub; + +impl DensePredictor for Stub { + fn family(&self) -> &'static str { + "sapiens2" + } + + fn tasks(&self) -> &[&'static str] { + &["depth"] + } + + fn predict(&self, _: &DenseInput, _: &Map) -> Result { + Ok(DenseOutput::default()) + } +} + +fn load_stub(_: &Path) -> Result> { + Ok(Box::new(Stub)) +} + +#[test] +fn resolves_known_model_type_from_config() { + let dir = tempfile::TempDir::new().unwrap(); + std::fs::write( + dir.path().join("config.json"), + r#"{"model_type":"sapiens2"}"#, + ) + .unwrap(); + let families = [DenseFamily { + model_type: "sapiens2", + load: load_stub, + }]; + assert_eq!( + resolve_family(dir.path(), &families).unwrap().model_type, + "sapiens2" + ); + assert_eq!( + load_dense_model_from(dir.path(), &families) + .unwrap() + .family(), + "sapiens2" + ); +} + +#[test] +fn recognized_but_unported_family_is_named() { + let dir = tempfile::TempDir::new().unwrap(); + std::fs::write(dir.path().join("config.json"), r#"{"model_type":"moge3"}"#).unwrap(); + let error = resolve_family(dir.path(), &[]).unwrap_err().to_string(); + assert!( + error.contains("moge3 is not supported by this build"), + "{error}" + ); +} diff --git a/src/vision/mod.rs b/src/vision/mod.rs index 4bb8322bc..38b95a399 100644 --- a/src/vision/mod.rs +++ b/src/vision/mod.rs @@ -29,6 +29,7 @@ use anyhow::Result; pub mod config; pub mod connectors; +pub mod dense; pub mod detection; pub mod encoders; pub mod feature_cache; diff --git a/tests/fixtures/generate_dense_interpolate_fixtures.py b/tests/fixtures/generate_dense_interpolate_fixtures.py new file mode 100644 index 000000000..a7eb6e9fe --- /dev/null +++ b/tests/fixtures/generate_dense_interpolate_fixtures.py @@ -0,0 +1,58 @@ +#!/usr/bin/env python3 +"""Regenerate the literal probes in dense/interpolate_tests.rs. + +Requires PyTorch. Run from the repository root and copy the printed Rust +slices into the named test cases. Keeping fixtures as literals means Rust tests +do not need a Python environment or a second interpolation implementation. +""" + +import torch +import torch.nn.functional as F + + +def source(height: int, width: int) -> torch.Tensor: + values = [(((index * 37) % 101) - 50) / 17 for index in range(height * width)] + return torch.tensor(values, dtype=torch.float32).reshape(1, 1, height, width) + + +def probes(values: torch.Tensor) -> str: + flat = values.flatten().tolist() + indices = sorted({0, 1, 4, len(flat) // 7, len(flat) // 3, len(flat) // 2, len(flat) - 5, len(flat) - 2, len(flat) - 1}) + return "&[" + ", ".join(f"({index}, {flat[index]:.9g})" for index in indices) + "]" + + +for height, width, out_height, out_width, suffix, scale in [ + (7, 9, 14, 5, "a", (2.0, 0.6)), + (14, 5, 7, 9, "b", (0.5, 1.8)), +]: + tensor = source(height, width) + for antialias in (False, True): + result = F.interpolate( + tensor, + size=(out_height, out_width), + mode="bilinear", + align_corners=False, + antialias=antialias, + ) + print(f"bilinear_{suffix}_{antialias} = {probes(result)}") + for antialias in (False, True): + result = F.interpolate( + tensor, + size=(out_height, out_width), + mode="bicubic", + align_corners=False, + antialias=antialias, + ) + print(f"bicubic_size_{suffix}_{antialias} = {probes(result)}") + for antialias in (False, True): + result = F.interpolate( + tensor, + scale_factor=scale, + mode="bicubic", + align_corners=False, + antialias=antialias, + recompute_scale_factor=False, + ) + print(f"bicubic_scale_{suffix}_{antialias} = {probes(result)}") + result = F.interpolate(tensor, size=(out_height, out_width), mode="nearest") + print(f"nearest_{suffix} = {probes(result)}") From 4177582516bcc53641b901e0aa66a9cfbcf1fe90 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 23 Oct 2026 00:16:21 +0900 Subject: [PATCH 3/5] docs: clarify dense checkpoint redirect guidance Fix the dense family registration guide so it clearly tells operators that `KNOWN_MODEL_TYPES` redirects accidental `mlxcel generate` use to `mlxcel predict`. Validation: - `cargo fmt --all -- --check` - `git diff --check` - Python fixture syntax and referenced documentation paths checked Refs #2349 --- docs/dense-prediction.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/dense-prediction.md b/docs/dense-prediction.md index 60c80a031..6f57a63f7 100644 --- a/docs/dense-prediction.md +++ b/docs/dense-prediction.md @@ -67,7 +67,7 @@ DenseFamily { } ``` -The registry key must equal the lowercase top-level `model_type` in `config.json`. Add the same model type to `KNOWN_MODEL_TYPES` only when the family is a dense predictor; that list makes accidental `mlxcel generate` use direct the operator to `mlxcel predict` even in builds where the family has not landed. +The registry key must equal the lowercase top-level `model_type` in `config.json`. Add the same model type to `KNOWN_MODEL_TYPES` only when the family is a dense predictor; that list directs operators who accidentally use `mlxcel generate` to `mlxcel predict`, even in builds where the family has not landed. Shared preprocessing APIs live in `src/vision/dense`: From 2849f3c1f5586a1f54dea26975df6139dd8eed54 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 23 Oct 2026 00:19:49 +0900 Subject: [PATCH 4/5] docs: add technical report for PR #2545 --- ...rediction-output-foundation-20261023.en.md | 111 ++++++++++++++++++ ...rediction-output-foundation-20261023.ko.md | 111 ++++++++++++++++++ 2 files changed, 222 insertions(+) create mode 100644 TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.en.md create mode 100644 TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.ko.md diff --git a/TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.en.md b/TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.en.md new file mode 100644 index 000000000..8b02b7dfb --- /dev/null +++ b/TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.en.md @@ -0,0 +1,111 @@ +# PR #2545: Add dense prediction output foundation + +**Date**: 2026-10-23 +**Status**: Partial (PR open; full ROCm workspace gate and CI rerun pending) +**Languages**: Rust, C++, Python, Markdown +**Risk Level**: Medium + +## Executive Summary + +PR #2545 establishes the shared infrastructure for dense vision models that return maps and keypoints instead of text or bounding boxes. It adds the `mlxcel predict` CLI, an object-safe predictor and family registry, bounded image and video input, upstream-compatible interpolation and preprocessing, and a file contract built around f32 safetensors, JSON metadata, optional keypoints, and PNG previews. The family registry is intentionally empty in this foundation PR; later issues in epic #2382 add the model networks and post-processing. + +The implementation also fixes a pinned MLX ROCm overlay prerequisite: `AsStrided` now selects the HIP device header and encoder instead of reaching the Metal path on Linux. Focused gfx1151 tests, a real ffmpeg MP4 path, package-scoped clippy, formatting, and repository contract gates passed. Metal validation, the new CI run, and the full ROCm workspace regression gate were not complete when this pre-merge report was written. + +## 1. Problem Statement + +mlxcel previously had token generation and an object-detection CLI, but no common output surface for per-pixel depth, normals, masks, segmentation, point maps, or human keypoints. Each planned dense model would otherwise need to invent its own input validation, interpolation, serialization, previews, and CLI behavior. This would make formats inconsistent and duplicate resource-safety work across Video Depth Anything, MoGe-3, Sapiens2, SAM 3D Objects, and their shared DINOv2 backbone. + +The foundation also needed to compile on the project’s ROCm development path. The pinned MLX overlay retained an upstream CUDA-or-Metal dispatch in `AsStrided`, causing Linux ROCm builds to include a Metal header and preventing the new Rust paths from being compiled and tested on gfx1151. + +Without this change, model ports would have no stable contract for operators or downstream tooling, and video/image inputs or repeated output writes could consume unbounded resources or overwrite unrelated files. + +## 2. Change Summary + +| Area | Change | +|---|---| +| CLI and input | Added `mlxcel predict` with exactly one of `--image` or `--video`, model task selection, JSON object options, frame/FPS controls, bounded payload/decode/allocation checks, and stage-specific errors. | +| Predictor architecture | Added `DenseInput`, `DenseOutput`, `PersonKeypoints`, the object-safe `DensePredictor` trait, and a `model_type` constructor registry. The production family table remains empty until model PRs land. | +| Tensor processing | Ported NHWC bilinear, bicubic, and nearest interpolation behavior; added ImageNet normalization, integer scaling, patch-multiple resize, and guarded OIHW-to-OHWI convolution conversion. | +| Output contract | Writes contiguous f32 `result.safetensors`, `meta.json`, optional `keypoints.json`, and depth/normal/mask/segmentation previews. Video previews use numbered directories. | +| Safe reruns | Materializes one tensor on the host at a time and replaces only writer-owned artifacts; stale preview paths and directory symlinks are handled without following an external target. | +| Registry integration | Added `Runtime::Predict`, `OutputKind::DenseMap`, and `OutputKind::Keypoints`; known dense checkpoints used with `generate` receive a redirect to `mlxcel predict`. | +| ROCm prerequisite | Corrected the pinned MLX `AsStrided` overlay to use the HIP encoder and refreshed overlay provenance. | +| Documentation | Added the dense CLI, output formats, preprocessing API, registration guide, scope limits, and README indexes. | + +Implementation diff before this report: 29 files, +2,823/-8. Related commits are `65cb1fdd` (ROCm prerequisite), `aee9d056` (foundation), and `41775825` (documentation clarification). + +## 3. Technical Decisions + +### 3.1 A separate `predict` surface + +Dense results do not fit the token stream from `generate` or the box-row contract from `detect`. A separate command keeps those interfaces stable while giving dense models an output directory, task selection, and file-oriented lifecycle. Extending `detect` was rejected because maps, intrinsics, and keypoints require different serialization and can be video sequences. + +### 3.2 An empty registry before model ports + +This PR defines the interface and validates it through a test-only stub family, but does not pretend that any dense checkpoint is already supported. A recognized-but-unported `model_type` produces a specific unsupported-family error, while accidental `generate` use points to `predict`. Model PRs therefore add only a predictor and one registry row against an already tested contract. + +### 3.3 Safetensors as the lossless source, PNG as a preview + +All tensor outputs are stored as contiguous f32 with their upstream keys and shapes. JSON records family, task, source size, elapsed time, finite ranges, model metadata, and flattened intrinsics. PNG files are derived previews: depth uses 16-bit normalization, normals map from `[-1, 1]`, masks threshold at 0.5, and segmentation stores 8-bit class IDs. Consumers that require numeric fidelity read safetensors rather than reverse-engineering a visualization. + +### 3.4 Bounded input and conservative output replacement + +Image loading enforces payload, decoder, and prospective RGB-f32 allocation limits. Video decoding reuses the shared ffmpeg loader, defaults to 32 sampled frames, and applies a 512 MiB decoded-frame budget. Output reruns delete only known artifacts created by the writer and retain unrelated files, which makes an existing directory usable without turning it into a recursive deletion boundary. + +### 3.5 Shared interpolation follows the upstream numerical contract + +The resize helpers use NHWC tensors and fixed upstream-generated PyTorch literals for bilinear and bicubic antialiasing, size and scale-factor modes, and nearest sampling in both upsample and downsample directions. Same-size and dtype-preserving behavior is documented because later model ports depend on these details for parity. + +## 4. Architecture and Integration Contract + +```text +image / video + | +bounded loader and sampler + | +model_type registry -> DensePredictor::predict(task, opts) + | +DenseOutput { tensors, keypoints, meta } + | +result.safetensors + meta.json + keypoints.json + PNG previews +``` + +Model implementations must expose an object-safe predictor, preserve upstream tensor names, declare supported tasks, attach preprocessing context to internal errors, and register their lowercase top-level `model_type`. The CLI provides the outer load, predict, and write contexts. Mesh and Gaussian export, streaming/webcam input, and a server route remain outside this contract. + +## 5. Verification and Limits + +| Check | Result at report time | +|---|---| +| Dense module tests on gfx1151 | 17 passed, covering interpolation, serialization/previews, registry behavior, guarded layout conversion, and preprocessing. | +| `predict` command tests | Five focused tests passed; the actual ffmpeg MP4 test was rerun with ffmpeg available and passed. | +| Registry and exact contract filters | 12 registry tests passed; parser, architecture JSON, generate redirect, catalog projection, and video-budget filters passed. | +| Output safety | Synthetic round-trip, numbered video previews, stale writer-owned cleanup, and directory-symlink replacement tests passed. | +| Lint and format | Package library and `mlxcel` binary clippy passed with `-D warnings`; `cargo fmt --all -- --check` and `git diff --check` passed. | +| Repository fast contracts | The orchestrated fast contract make gates passed before report generation. | +| Full ROCm workspace regression | Running when this report was written; no pass is claimed here. | +| CI after report commit | Pending; the report commit triggers a fresh PR run. | +| Metal | Not run on the ROCm Linux host; delegated to CI. | + +No real dense model inference is claimed by this PR because the production family table is empty by design. No throughput or memory benchmark was performed; the resource work here consists of hard limits and one-host-tensor-at-a-time serialization rather than a measured performance optimization. + +## 6. Learning Points + +- Infrastructure can be tested before a model lands by registering a test-only predictor that exercises the real CLI, input decoding, task selection, and writer. +- A preview is not a numerical interchange format. Keeping lossless f32 tensors and recording the exact normalization range prevents visualization choices from becoming an implicit API. +- Resource budgets must include representation expansion. A compressed image that passes a payload limit can still exceed memory limits when decoded and converted to RGB f32. +- Safe reuse of an output directory requires an explicit ownership list. Removing only standard artifacts avoids recursive deletion of user files and makes symlink behavior testable. +- Numerical ports need fixtures produced by the reference implementation. Fixed PyTorch probes make interpolation drift visible without requiring Python or PyTorch in the Rust test environment. + +## 7. Follow-up Actions + +- Complete the full ROCm workspace regression gate and the fresh PR CI run before merge. +- Confirm Metal compilation and tests in CI. +- Add DINOv2 and the dense model families through the dependent issues in epic #2382. +- Add mesh/Gaussian export with the SAM 3D Objects implementation; define a server wire contract only when a client requirement exists. + +## References + +- PR #2545: `feat: add dense prediction output foundation` +- Issue #2349: dense-prediction output surface +- Epic #2382: dense-prediction model delivery +- `docs/dense-prediction.md`: operator and model-integration contract diff --git a/TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.ko.md b/TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.ko.md new file mode 100644 index 000000000..5af6d21f1 --- /dev/null +++ b/TECHNICAL_REPORTS/2545-dense-prediction-output-foundation-20261023.ko.md @@ -0,0 +1,111 @@ +# PR #2545: Dense prediction 출력 기반 추가 + +**작성일**: 2026-10-23 +**상태**: 부분 완료 (PR 열림; 전체 ROCm workspace 게이트와 CI 재실행 대기 중) +**언어**: Rust, C++, Python, Markdown +**위험도**: 중간 + +## 요약 + +PR #2545는 텍스트나 bounding box 대신 map과 keypoint를 반환하는 dense vision 모델의 공용 기반을 만든다. `mlxcel predict` CLI, object-safe predictor와 family registry, 제한이 있는 이미지·비디오 입력, 업스트림과 호환되는 interpolation·preprocessing, f32 safetensors·JSON metadata·선택적 keypoint·PNG preview로 구성된 파일 계약을 추가한다. 이 기반 PR의 family registry는 의도적으로 비어 있으며, epic #2382의 후속 이슈가 모델 network와 post-processing을 추가한다. + +구현에는 pinned MLX ROCm overlay의 선행 수정도 포함된다. Linux에서 Metal 경로로 들어가던 `AsStrided`가 HIP device header와 encoder를 선택하도록 고쳤다. gfx1151 집중 테스트, 실제 ffmpeg MP4 경로, package 범위 clippy, formatting, 저장소 contract 게이트는 통과했다. 이 pre-merge 보고서를 작성할 때 Metal 검증, 새 CI 실행, 전체 ROCm workspace 회귀 게이트는 완료되지 않았다. + +## 1. 문제 정의 + +기존 mlxcel에는 token generation과 object detection CLI가 있었지만, pixel 단위 depth, normal, mask, segmentation, point map, human keypoint를 위한 공용 출력 surface가 없었다. 이 기반이 없으면 Video Depth Anything, MoGe-3, Sapiens2, SAM 3D Objects와 공용 DINOv2 backbone의 각 port가 입력 검증, interpolation, serialization, preview, CLI 동작을 따로 만들어야 한다. 그 결과 format이 달라지고 resource safety 로직이 중복된다. + +또한 프로젝트의 ROCm 개발 경로에서 이 기반이 compile되어야 했다. pinned MLX overlay의 `AsStrided`에는 업스트림 CUDA-or-Metal dispatch가 남아 있어 Linux ROCm build가 Metal header를 include했고, 새 Rust 경로를 gfx1151에서 compile하고 테스트할 수 없었다. + +이 변경이 없으면 model port마다 operator와 downstream tooling을 위한 계약이 달라지고, 이미지·비디오 입력 또는 반복 출력이 무제한 resource를 쓰거나 관계없는 파일을 덮어쓸 수 있다. + +## 2. 변경 요약 + +| 영역 | 변경 | +|---|---| +| CLI와 입력 | `--image` 또는 `--video` 중 정확히 하나를 받는 `mlxcel predict`, model task 선택, JSON object option, frame/FPS 제어, payload/decode/allocation 제한, 단계별 error context를 추가했다. | +| Predictor 구조 | `DenseInput`, `DenseOutput`, `PersonKeypoints`, object-safe `DensePredictor` trait, `model_type` constructor registry를 추가했다. 실제 family table은 model PR이 들어올 때까지 비어 있다. | +| Tensor 처리 | NHWC bilinear, bicubic, nearest interpolation을 port했고, ImageNet normalization, integer scaling, patch 배수 resize, 보호된 OIHW-to-OHWI convolution 변환을 추가했다. | +| 출력 계약 | contiguous f32 `result.safetensors`, `meta.json`, 선택적 `keypoints.json`, depth/normal/mask/segmentation preview를 쓴다. 비디오 preview는 번호가 붙은 directory를 사용한다. | +| 안전한 재실행 | host에는 한 번에 tensor 하나만 materialize하고 writer가 소유한 artifact만 교체한다. 오래된 preview 경로와 directory symlink는 외부 target을 따라가지 않고 처리한다. | +| Registry 통합 | `Runtime::Predict`, `OutputKind::DenseMap`, `OutputKind::Keypoints`를 추가했다. 알려진 dense checkpoint를 `generate`로 실행하면 `mlxcel predict` 안내를 반환한다. | +| ROCm 선행 수정 | pinned MLX `AsStrided` overlay가 HIP encoder를 사용하도록 고치고 overlay provenance를 갱신했다. | +| 문서 | dense CLI, 출력 format, preprocessing API, registration guide, 범위 제한, README index를 추가했다. | + +이 보고서를 제외한 구현 diff는 29개 파일, +2,823/-8이다. 관련 commit은 `65cb1fdd`(ROCm 선행 수정), `aee9d056`(기반 구현), `41775825`(문서 표현 수정)이다. + +## 3. 기술적 선택과 그 이유 + +### 3.1 별도의 `predict` surface + +Dense 결과는 `generate`의 token stream이나 `detect`의 box row 계약에 맞지 않는다. 별도 command를 두어 기존 interface를 유지하면서 dense 모델에 output directory, task 선택, file 중심 lifecycle을 제공한다. Map, intrinsics, keypoint는 다른 serialization이 필요하고 비디오 sequence가 될 수 있으므로 `detect` 확장은 채택하지 않았다. + +### 3.2 Model port보다 먼저 빈 registry 제공 + +이 PR은 interface를 정의하고 test-only stub family로 검증하지만 dense checkpoint가 이미 지원되는 것처럼 보이게 하지 않는다. 알려졌지만 아직 port되지 않은 `model_type`은 구체적인 unsupported-family error를 내고, `generate`를 잘못 사용하면 `predict`로 안내한다. 이후 model PR은 검증된 계약에 predictor와 registry row 하나만 추가하면 된다. + +### 3.3 Safetensors를 무손실 원본으로, PNG를 preview로 사용 + +모든 tensor 출력은 upstream key와 shape를 유지한 contiguous f32로 저장한다. JSON에는 family, task, source size, elapsed time, finite range, model metadata, 평탄화한 intrinsics를 기록한다. PNG는 파생 preview다. Depth는 16-bit normalization, normal은 `[-1, 1]` mapping, mask는 0.5 threshold, segmentation은 8-bit class ID를 쓴다. 수치 정확도가 필요한 consumer는 visualization을 역변환하지 않고 safetensors를 읽는다. + +### 3.4 입력 제한과 보수적인 출력 교체 + +이미지 loading은 payload, decoder, 예상 RGB-f32 allocation 제한을 적용한다. 비디오 decoding은 공용 ffmpeg loader를 재사용하고 기본 32개 sampled frame과 512 MiB decoded-frame budget을 적용한다. 재실행할 때 writer가 만든 알려진 artifact만 지우고 관계없는 파일은 유지한다. 따라서 기존 directory를 재사용하면서도 recursive deletion 범위로 만들지 않는다. + +### 3.5 공용 interpolation은 업스트림 수치 계약을 따름 + +Resize helper는 NHWC tensor를 사용하고, 업스트림 PyTorch가 생성한 고정 literal로 bilinear·bicubic antialiasing, size·scale-factor mode, nearest sampling을 upsample과 downsample 양방향에서 비교한다. 이후 model port의 parity가 이 세부 동작에 의존하므로 same-size와 dtype 보존 동작도 문서화했다. + +## 4. 구조와 통합 계약 + +```text +image / video + | +제한이 있는 loader와 sampler + | +model_type registry -> DensePredictor::predict(task, opts) + | +DenseOutput { tensors, keypoints, meta } + | +result.safetensors + meta.json + keypoints.json + PNG previews +``` + +Model 구현은 object-safe predictor를 노출하고, upstream tensor name을 유지하고, 지원 task를 선언하고, 내부 preprocessing error에 context를 붙이고, lowercase top-level `model_type`을 등록해야 한다. CLI는 바깥쪽 load, predict, write context를 제공한다. Mesh/Gaussian export, streaming/webcam 입력, server route는 이 계약의 범위 밖이다. + +## 5. 검증과 제한 사항 + +| 확인 항목 | 보고서 작성 시점의 결과 | +|---|---| +| gfx1151 dense module 테스트 | 17개 통과. Interpolation, serialization/preview, registry 동작, 보호된 layout 변환, preprocessing을 포함한다. | +| `predict` command 테스트 | 집중 테스트 5개 통과. ffmpeg가 있는 상태에서 실제 MP4 테스트를 다시 실행해 통과했다. | +| Registry와 정확한 contract filter | Registry 테스트 12개 통과. Parser, architecture JSON, generate redirect, catalog projection, video-budget filter가 통과했다. | +| 출력 안전성 | Synthetic round-trip, 번호가 붙은 video preview, 오래된 writer-owned artifact 정리, directory-symlink 교체 테스트가 통과했다. | +| Lint와 format | Package library와 `mlxcel` binary clippy가 `-D warnings`로 통과했다. `cargo fmt --all -- --check`, `git diff --check`도 통과했다. | +| 저장소 fast contract | 보고서 생성 전에 orchestrated fast contract make 게이트가 통과했다. | +| 전체 ROCm workspace 회귀 | 보고서 작성 시 실행 중이며 여기서는 통과했다고 기록하지 않는다. | +| 보고서 commit 이후 CI | 대기 중이다. 보고서 commit이 새 PR 실행을 시작한다. | +| Metal | ROCm Linux host에서는 실행하지 않았고 CI에 맡긴다. | + +실제 production family table이 의도적으로 비어 있으므로 이 PR은 실제 dense model inference를 지원한다고 주장하지 않는다. Throughput 또는 memory benchmark도 수행하지 않았다. 여기의 resource 작업은 측정된 성능 최적화가 아니라 hard limit와 한 번에 host tensor 하나만 쓰는 serialization이다. + +## 6. 학습 포인트 + +- Model이 들어오기 전에도 test-only predictor를 등록하면 실제 CLI, input decoding, task 선택, writer를 통과하는 infrastructure를 검증할 수 있다. +- Preview는 수치 교환 format이 아니다. 무손실 f32 tensor와 정확한 normalization range를 함께 저장하면 visualization 선택이 암묵적인 API가 되는 것을 막는다. +- Resource budget은 representation 확장을 포함해야 한다. 압축된 이미지가 payload 제한을 통과해도 decode와 RGB f32 변환 뒤에는 memory 제한을 넘을 수 있다. +- Output directory를 안전하게 재사용하려면 명시적인 소유 목록이 필요하다. 표준 artifact만 제거하면 사용자 파일의 recursive deletion을 피하고 symlink 동작도 테스트할 수 있다. +- 수치 port에는 reference implementation이 만든 fixture가 필요하다. 고정 PyTorch probe는 Rust test 환경에 Python이나 PyTorch를 요구하지 않으면서 interpolation drift를 드러낸다. + +## 7. 후속 조치 + +- Merge 전에 전체 ROCm workspace 회귀 게이트와 새 PR CI 실행을 완료한다. +- CI에서 Metal compile과 테스트를 확인한다. +- Epic #2382의 dependent issue를 통해 DINOv2와 dense model family를 추가한다. +- SAM 3D Objects 구현에서 mesh/Gaussian export를 추가한다. Server wire contract는 client 요구가 생길 때 정의한다. + +## 참고 자료 + +- PR #2545: `feat: add dense prediction output foundation` +- Issue #2349: dense-prediction output surface +- Epic #2382: dense-prediction model delivery +- `docs/dense-prediction.md`: operator 및 model integration 계약 From d5da21193328aa59fc744d383d8798299c8de4a6 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 23 Oct 2026 00:47:54 +0900 Subject: [PATCH 5/5] test: preserve generated dense interpolation reference literals Keep the PyTorch reference decimals unchanged with a fixture-scoped precision lint allowance, matching the existing scheduler fixtures. CI checks test code in addition to the library and binary, so its wider scope caught these literal-only warnings. Refs #2349 --- src/vision/dense/interpolate_tests.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/vision/dense/interpolate_tests.rs b/src/vision/dense/interpolate_tests.rs index 0d50b5769..011d93f74 100644 --- a/src/vision/dense/interpolate_tests.rs +++ b/src/vision/dense/interpolate_tests.rs @@ -1,6 +1,9 @@ // Copyright 2025-2026 Lablup Inc. // SPDX-License-Identifier: Apache-2.0 +// Preserve the decimal literals emitted by the external PyTorch reference. +#![allow(clippy::excessive_precision)] + use super::*; fn input(height: usize, width: usize) -> UniquePtr {