kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/depthpro/_models/depthpro.py

@@ -0,0 +1,144 @@
+"""Handler for Depth Pro image → depth map via mflux."""
+
+import gc
+import time
+from types import SimpleNamespace
+from typing import Any
+
+import numpy as np
+from PIL import Image
+
+from kiapi.core.file import FileStore
+from kiapi.core.model import ModelSpec
+from kiapi.core.workdir import create_work_dir
+
+from .._views.estimate_params import EstimateParams
+
+FEATURES = {"image", "depth"}
+
+
+def build_model(*, quantize: int | None) -> Any:
+    from mflux.models.depth_pro.model.depth_pro import DepthPro  # type: ignore
+
+    return DepthPro(quantize=quantize)
+
+
+def load(spec: ModelSpec) -> SimpleNamespace:
+    from .._settings import settings_manager
+
+    quantize = settings_manager.get_settings().default_quantize
+    model = build_model(quantize=quantize)
+    return SimpleNamespace(model=model, spec=spec, quantize=quantize)
+
+
+def _store_depth_image(
+    result: Any,
+    *,
+    files: FileStore,
+    meta: dict[str, Any],
+) -> dict[str, Any]:
+    pil = result.depth_image
+    tmp_dir = create_work_dir("image/depthpro")
+    out_path = tmp_dir / "depth.png"
+    pil.save(out_path, format="PNG")
+    if not out_path.exists():
+        raise RuntimeError("depth estimation finished but no depth image was produced")
+    rec = files.put_path(
+        out_path,
+        filename=f"depthpro_{int(time.time())}.png",
+        content_type="image/png",
+        meta=meta | {"kind": "depth_image"},
+        move=True,
+    )
+    return {"depth_image_file_id": rec.file_id, "depth_image_bytes": rec.size}
+
+
+def _store_depth_data(
+    result: Any,
+    *,
+    files: FileStore,
+    meta: dict[str, Any],
+) -> dict[str, Any]:
+    tmp_dir = create_work_dir("image/depthpro")
+    out_path = tmp_dir / "depth.npz"
+    np.savez_compressed(
+        out_path,
+        depth=np.asarray(result.depth_array),
+        min_depth=float(result.min_depth),
+        max_depth=float(result.max_depth),
+    )
+    if not out_path.exists():
+        raise RuntimeError(
+            "depth estimation finished but no depth data file was produced"
+        )
+    rec = files.put_path(
+        out_path,
+        filename=f"depthpro_{int(time.time())}.npz",
+        content_type="application/octet-stream",
+        meta=meta | {"kind": "depth_data", "format": "npz"},
+        move=True,
+    )
+    return {"depth_data_file_id": rec.file_id, "depth_data_bytes": rec.size}
+
+
+def _estimate(
+    model: Any,
+    spec: ModelSpec,
+    params: EstimateParams,
+    files: FileStore,
+) -> dict[str, Any]:
+    with Image.open(params.image_path) as src:
+        input_size = src.size
+        input_mode = src.mode
+
+    t0 = time.time()
+    result = model.create_depth_map(params.image_path)
+    total_s = round(time.time() - t0, 2)
+
+    depth_array = np.asarray(result.depth_array)
+    meta = {
+        "model": spec.name,
+        "params": params.model_dump(exclude={"image_path"}),
+        "input_width": input_size[0],
+        "input_height": input_size[1],
+        "input_mode": input_mode,
+        "width": result.depth_image.width,
+        "height": result.depth_image.height,
+        "mode": result.depth_image.mode,
+        "array_shape": list(depth_array.shape),
+        "min_depth": round(float(result.min_depth), 6),
+        "max_depth": round(float(result.max_depth), 6),
+        "timings": {"total_s": total_s},
+    }
+
+    stored = _store_depth_image(result, files=files, meta=meta)
+    if params.include_depth_data:
+        stored |= _store_depth_data(result, files=files, meta=meta)
+    else:
+        stored["depth_data_file_id"] = None
+        stored["depth_data_bytes"] = None
+    return {**stored, **meta}
+
+
+def run(
+    payload: SimpleNamespace,
+    params: EstimateParams,
+    files: FileStore,
+) -> dict[str, Any]:
+    return _estimate(payload.model, payload.spec, params, files)
+
+
+def run_transient(
+    spec: ModelSpec,
+    params: EstimateParams,
+    files: FileStore,
+) -> dict[str, Any]:
+    model = build_model(quantize=params.quantize)
+    try:
+        return _estimate(model, spec, params, files)
+    finally:
+        del model
+        gc.collect()
+        import mlx.core as mx
+
+        mx.clear_cache()

kiapi/capabilities/depthpro/_operations/handle_estimate.py

@@ -0,0 +1,50 @@
+"""Depth Pro service entry: resolve input image → run resident/transient model.
+
+Body of the depthpro job's worker-thread thunk. Stores the depth PNG (and an
+optional NPZ of the raw depth array) in the global file store; the job's
+``artifacts`` is the produced file_id(s).
+"""
+
+from kiapi.capabilities import resolve_file_ref
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileID
+from kiapi.core.job import JobResult, creep_progress
+from kiapi.core.model import model_registry
+
+from .._helpers.validate_estimate import validate_estimate
+from .._settings import settings_manager
+from .._views.estimate_request import EstimateRequest
+from .is_quantize_override import is_quantize_override
+from .resolve_estimate_params import resolve_estimate_params
+
+
+def handle_estimate(
+    ctx: AppContext, req: EstimateRequest
+) -> tuple[JobResult, list[FileID]]:
+    settings = settings_manager.get_settings()
+    spec = model_registry.resolve("depthpro", req.model)
+    ctx.ensure_model_ready(spec)
+    image = resolve_file_ref(ctx.file_store, req.image, kind="image")
+    validate_estimate(req, image_path=image.path)
+    override = is_quantize_override(settings, req)
+    params = resolve_estimate_params(
+        settings,
+        req,
+        variant=spec.name,
+        image_file_id=image.file_id,
+        image_path=image.path,
+    )
+
+    if override:
+        ctx.memory_manager.reserve(spec.weight_gb + spec.peak_headroom_gb)
+        with creep_progress(eta_s=settings.progress_eta_s):
+            result = spec.module.run_transient(spec, params, ctx.file_store)
+    else:
+        payload = ctx.memory_manager.acquire(spec)
+        with creep_progress(eta_s=settings.progress_eta_s):
+            result = spec.module.run(payload, params, ctx.file_store)
+
+    artifacts = [result["depth_image_file_id"]]
+    if result.get("depth_data_file_id") is not None:
+        artifacts.append(result["depth_data_file_id"])
+    return result, artifacts

kiapi/capabilities/depthpro/_operations/is_quantize_override.py

@@ -0,0 +1,12 @@
+"""Decide whether a request's quantize differs from the default.
+
+A True result means the resident model (quantized at its default) can't serve the
+request, so ``handle_estimate`` runs a one-off transient model instead.
+"""
+
+from .._settings import DepthProSettings
+from .._views.estimate_request import EstimateRequest
+
+
+def is_quantize_override(settings: DepthProSettings, req: EstimateRequest) -> bool:
+    return req.quantize is not None and req.quantize != settings.default_quantize

kiapi/capabilities/depthpro/_operations/resolve_estimate_params.py

@@ -0,0 +1,24 @@
+"""Merge an estimate request with settings defaults into the complete EstimateParams."""
+
+from .._settings import DepthProSettings
+from .._views.estimate_params import EstimateParams
+from .._views.estimate_request import EstimateRequest
+
+
+def resolve_estimate_params(
+    settings: DepthProSettings,
+    req: EstimateRequest,
+    *,
+    variant: str,
+    image_file_id: str,
+    image_path: str,
+) -> EstimateParams:
+    return EstimateParams(
+        model=variant,
+        image_file_id=image_file_id,
+        image_path=image_path,
+        quantize=req.quantize
+        if req.quantize is not None
+        else settings.default_quantize,
+        include_depth_data=req.include_depth_data,
+    )

kiapi/capabilities/depthpro/_settings.py

@@ -0,0 +1,46 @@
+"""Depth Pro capability defaults + caps, read from ``KIAPI_DEPTHPRO_`` env vars."""
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class DepthProSettings(BaseSettings):
+    """Settings for Depth Pro quantization, input image limits, and progress display."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_DEPTHPRO_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    default_quantize: int | None = Field(
+        default=8,
+        title="Default quantization bits",
+        description=(
+            "Default quantization bit count used when loading the Depth Pro model.\n"
+            "Use a bit count such as 8, or None to disable quantization."
+        ),
+    )
+
+    max_input_pixels: int = Field(
+        default=4096 * 4096,
+        title="Maximum input pixels",
+        description=(
+            "Maximum total pixel count for an input image passed to Depth Pro.\n"
+            "This protects the single worker queue from being blocked for a long "
+            "time by very large images."
+        ),
+    )
+
+    progress_eta_s: float = Field(
+        default=8.0,
+        title="Progress ETA seconds",
+        description=(
+            "Expected processing seconds used to show synthetic progress while Depth Pro runs.\n"
+            "Set to 0 to keep progress at the coarse value until completion."
+        ),
+    )
+
+
+settings_manager = SettingsManager(DepthProSettings)

kiapi/capabilities/depthpro/_views/estimate_params.py

@@ -0,0 +1,21 @@
+"""The complete contract for one Depth Pro image → depth map run.
+
+Built from settings + request by ``resolve_estimate_params``; the model needs
+nothing else to produce the depth map and its metadata. ``image_path`` is the
+locally resolved input file, kept out of the recorded metadata in favour of
+``image_file_id``.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class EstimateParams(BaseModel):
+    kind: Literal["estimate"] = "estimate"
+    model: str
+
+    image_file_id: str
+    image_path: str
+    quantize: int | None
+    include_depth_data: bool

kiapi/capabilities/depthpro/_views/estimate_request.py

@@ -0,0 +1,63 @@
+"""Depth Pro request model.
+
+Image → depth map (Depth Pro). One endpoint serves both sync and async via
+``mode``. The input image is referenced by ``image`` (FileRef). A ``quantize``
+override runs a one-off transient model; otherwise the resident model is used.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from kiapi.core.file import FileRef
+
+
+class EstimateRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = Field(
+        default=None,
+        description=(
+            "Model variant (see GET /v1/image/depthpro/models). Omit for the "
+            "default `base`. Aliases `depthpro`, `depth-pro`, `DepthPro` are accepted."
+        ),
+    )
+    mode: Literal["sync", "async"] = Field(
+        default="sync",
+        description=(
+            "`sync` waits for the depth map (504 on timeout); `async` returns 202 "
+            "with a job_id immediately — poll GET /v1/jobs/{job_id}."
+        ),
+    )
+
+    image: FileRef = Field(
+        description=(
+            "Input image to estimate depth for, as a Files-API file id, http(s) "
+            "URL, or data URL. Must not exceed the configured input pixel cap."
+        ),
+    )
+    quantize: int | None = Field(
+        default=None,
+        description=(
+            "Quantization bits, one of {3, 4, 5, 6, 8}. Omit to use the server "
+            "default (8). A value differing from the resident model's quantization "
+            "runs a one-off transient model (no resident reuse), so it is slower."
+        ),
+    )
+
+    include_depth_data: bool = Field(
+        default=True,
+        description=(
+            "When true, also store a compressed NPZ with the raw float depth array "
+            "and min/max depth (`depth_data_file_id`) alongside the grayscale PNG. "
+            "When false, only the PNG is produced, so a sync call can return the raw "
+            "PNG bytes directly (single artifact, unless Accept: application/json)."
+        ),
+    )
+
+    def gen_params(self) -> dict:
+        return {
+            "image": self.image.model_dump(mode="json"),
+            "quantize": self.quantize,
+            "include_depth_data": self.include_depth_data,
+        }

kiapi/capabilities/embedding/README.ja.md

@@ -0,0 +1,137 @@
+# embedding
+
+[English](README.md) | **日本語**
+
+[mlx-embeddings](https://github.com/Blaizzy/mlx-embeddings) で、テキスト / 画像の埋め込み API を提供します。
+
+- **text embedding**:
+  - Qwen3-Embedding-8B
+- **multimodal embedding**:
+  - Qwen3-VL-Embedding-2B
+
+## API
+
+| Endpoint | Name | Description |
+|---|---|---|
+| `POST /v1/embedding` | 埋め込み | 1 アイテムをベクトル化。レスポンスは `{ model, embedding, dimension, usage, timings }`。 |
+| `GET /v1/embedding/models` | モデル一覧 | 利用可能なモデルの一覧を返す。 |
+| `GET /v1/embedding/openapi.json` | OpenAPI | 詳細な入出力仕様、使い方、TIPS を返す。 |
+
+## API Docs
+
+- [OpenAPI JSON](https://kiarina.github.io/kiapi/v1/embedding/openapi.json)
+- [Swagger UI](https://kiarina.github.io/kiapi/v1/embedding/docs.html)
+- [ReDoc](https://kiarina.github.io/kiapi/v1/embedding/redoc.html)
+
+## Dependencies
+
+| Package | License | Description |
+|---|---|---|
+| [mlx-embeddings](https://github.com/Blaizzy/mlx-embeddings) | MIT | Qwen3 / Qwen3-VL 埋め込みモデルを MLX 上で駆動。 |
+
+## Models
+
+| Model | License | Terms | Size | Description |
+|---|---|---|---:|---|
+| [mlx-community/Qwen3-Embedding-8B-mxfp8](https://huggingface.co/mlx-community/Qwen3-Embedding-8B-mxfp8) | Apache-2.0 | 不要 | 7.82 GB | `qwen3-embedding-8b`（デフォルト）。text 専用、4096 次元。alias: `text`, `qwen3-embedding`, `qwen3_embedding`。 |
+| [mlx-community/Qwen3-VL-Embedding-2B-mxfp8](https://huggingface.co/mlx-community/Qwen3-VL-Embedding-2B-mxfp8) | Apache-2.0 | 不要 | 2.59 GB | `qwen3-vl-embedding-2b`。text + image、2048 次元。alias: `vl`, `qwen3-vl-embedding`, `qwen3_vl_embedding`。 |
+
+- 選択したモデルが対応しないモダリティを送ると **HTTP 400**。
+  例: `qwen3-embedding-8b` に `image` を送るとエラー。
+- VL モデルでは `text` と `image` を同時に送ると、1 つの結合アイテムとして
+  埋め込みます。
+
+## Notes
+
+- **レスポンス形式**:
+  `embedding` は float 配列、`dimension` はその長さです。
+  `usage` は `prompt_tokens` / `total_tokens` を返します。
+  現在、モデル実行経路によっては token 数は `0` のままです。
+- **出力ベクトル**:
+  `mlx-embeddings` の出力を `_utils/to_vector.py` で Python の `list[float]` に変換します。
+  ヘルプでは L2 正規化・last-token pooling のベクトルとして扱います。
+- **入力形式**:
+  `text` は文字列をそのまま渡します。
+  `image` は base64、data URL、http(s) URL を受け付け、
+  VL モデル実行前に一時ファイルへ materialize します。
+- **max_length**:
+  tokenizer へ渡す最大長は `KIAPI_EMBEDDING_MAX_LENGTH` で設定します。
+  既定値は `512` です。
+- **2 つの埋め込み経路**:
+  text モデルは `mlx_embeddings.generate(...)`、
+  VL モデルは `model.process([item], processor=...)` を使います。
+  同じ依存ライブラリ内でも経路が違うため、`mlx-embeddings` を更新したときは両方を検証してください。
+- **VL processor の補正**:
+  MLX 変換版 Qwen3-VL processor は `chat_template` や image/video token-id 属性を欠くことがあります。
+  `_models/qwen3_vl.py` の `_patch_processor` が best-effort で補います。
+  依存更新時はこの補正がまだ必要か確認してください。
+
+## Quickstart
+
+### text - デフォルトモデル
+
+```bash
+jq -n \
+'{
+  text: "今日はいい天気ですね"
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, head: .embedding[0:5], timings}'
+```
+
+### text - VL モデル
+
+```bash
+jq -n \
+'{
+  model: "vl",
+  text: "a photo of a cat"
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, timings}'
+```
+
+### image - VL モデル
+
+```bash
+IMG=$(base64 -i your_image.png)
+
+jq -n \
+--arg image "$IMG" \
+'{
+  model: "vl",
+  image: $image
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, timings}'
+```
+
+### text + image - VL モデル
+
+```bash
+IMG=$(base64 -i your_image.png)
+
+jq -n \
+--arg image "$IMG" \
+'{
+  model: "vl",
+  text: "a photo of a cat",
+  image: $image
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, timings}'
+```
+
+### モデル一覧
+
+```bash
+curl -sS http://localhost:${PORT:-8000}/v1/embedding/models | jq .
+```

kiapi/capabilities/embedding/README.md

@@ -0,0 +1,128 @@
+# embedding
+
+**English** | [日本語](README.ja.md)
+
+[mlx-embeddings](https://github.com/Blaizzy/mlx-embeddings) provides a text/image embedding API.
+
+- **text embedding**:
+  - Qwen3-Embedding-8B
+- **multimodal embedding**:
+  - Qwen3-VL-Embedding-2B
+
+## API
+
+| Endpoint | Name | Description |
+|---|---|---|
+| `POST /v1/embedding` | Embed | 1 Vectorize item. The response is `{ model, embedding, dimension, usage, timings }`. |
+| `GET /v1/embedding/models` | Model list | Returns a list of available models. |
+| `GET /v1/embedding/openapi.json` | OpenAPI | Returns detailed input/output specifications, usage, and TIPS. |
+
+## API Docs
+
+- [OpenAPI JSON](https://kiarina.github.io/kiapi/v1/embedding/openapi.json)
+- [Swagger UI](https://kiarina.github.io/kiapi/v1/embedding/docs.html)
+- [ReDoc](https://kiarina.github.io/kiapi/v1/embedding/redoc.html)
+
+## Dependencies
+
+| Package | License | Description |
+|---|---|---|
+| [mlx-embeddings](https://github.com/Blaizzy/mlx-embeddings) | MIT | Drive Qwen3 / Qwen3-VL embedding models on MLX. |
+
+## Models
+
+| Model | License | Terms | Size | Description |
+|---|---|---|---:|---|
+| [mlx-community/Qwen3-Embedding-8B-mxfp8](https://huggingface.co/mlx-community/Qwen3-Embedding-8B-mxfp8) | Apache-2.0 | Not required | 7.82 GB | `qwen3-embedding-8b` (default). text only, 4096 dimensions. alias: `text`, `qwen3-embedding`, `qwen3_embedding`. |
+| [mlx-community/Qwen3-VL-Embedding-2B-mxfp8](https://huggingface.co/mlx-community/Qwen3-VL-Embedding-2B-mxfp8) | Apache-2.0 | Not required | 2.59 GB | `qwen3-vl-embedding-2b`. text + image, 2048 dimensions. alias: `vl`, `qwen3-vl-embedding`, `qwen3_vl_embedding`. |
+
+- **HTTP 400** when sending a modality that is not supported by the selected model.
+  Example: Error when sending `image` to `qwen3-embedding-8b`.
+- In the VL model, if you send `text` and `image` at the same time, they will be sent as one combined item.
+  Embed.
+
+## Notes
+
+- **Response format**:
+  `embedding` is a float array and `dimension` is its length.
+  `usage` returns `prompt_tokens` / `total_tokens`.
+  Currently, the number of tokens remains `0` depending on the model execution path.
+- **Output Vector**:
+  Convert the output of `mlx-embeddings` to a Python `list[float]` with `_utils/to_vector.py`.
+  In the help, it is treated as a vector for L2 normalization and last-token pooling.
+- **Input format**:
+  `text` passes the string as is.
+  `image` accepts base64, data URLs, http(s) URLs,
+  Materialize to a temporary file before running the VL model.
+- **max_length**:
+  The maximum length passed to the tokenizer is set with `KIAPI_EMBEDDING_MAX_LENGTH`.
+  Default value is `512`.
+- **Two embedded paths**:
+  The text model is `mlx_embeddings.generate(...)`,
+  VL models use `model.process([item], processor=...)`.
+  The paths are different even within the same dependent library, so please verify both when updating `mlx-embeddings`.
+- **VL processor correction**:
+  MLX converted Qwen3-VL processors may lack `chat_template` and image/video token-id attributes.
+  `_patch_processor` in `_models/qwen3_vl.py` compensates with best-effort.
+  When updating dependencies, please check whether this correction is still required.
+
+## Quickstart
+
+### text - default model
+```bash
+jq -n \
+'{
+  text: "今日はいい天気ですね"
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, head: .embedding[0:5], timings}'
+```
+### text - VL model
+```bash
+jq -n \
+'{
+  model: "vl",
+  text: "a photo of a cat"
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, timings}'
+```
+### image - VL model
+```bash
+IMG=$(base64 -i your_image.png)
+
+jq -n \
+--arg image "$IMG" \
+'{
+  model: "vl",
+  image: $image
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, timings}'
+```
+### text + image - VL model
+```bash
+IMG=$(base64 -i your_image.png)
+
+jq -n \
+--arg image "$IMG" \
+'{
+  model: "vl",
+  text: "a photo of a cat",
+  image: $image
+}' |
+curl -sS http://localhost:${PORT:-8000}/v1/embedding \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '{model, dimension, timings}'
+```
+### Model list
+```bash
+curl -sS http://localhost:${PORT:-8000}/v1/embedding/models | jq .
+```

kiapi/capabilities/embedding/__init__.py

@@ -0,0 +1,24 @@
+"""Embedding capability — multi-model text/image embeddings via mlx-embeddings.
+
+Integrates mlx-embedding-server: served from the shared single-flight worker
+under the global memory budget. The `model` field selects which model answers;
+one item per request with one field per modality. Sync only.
+
+These are small models that agents call frequently, so they get a positive
+``priority`` (higher = kept longer): they survive eviction churn from big chat /
+generation models, which is exactly the intended use of priority.
+``weight_gb`` / ``peak_headroom_gb`` are seeded from the original server's
+estimates and reconciled with on-device measurement at load.
+"""
+
+from ._helpers.register import register
+from ._operations.handle_embed import handle_embed
+from ._settings import settings_manager
+from ._views.embed_request import EmbedRequest
+
+__all__ = [
+    "EmbedRequest",  # ._views
+    "handle_embed",  # ._operations
+    "register",  # ._helpers
+    "settings_manager",  # ._settings
+]