kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/seedvr2/_models/seedvr2.py

@@ -0,0 +1,140 @@
+"""Handler for SeedVR2 image upscale via mflux."""
+
+import gc
+import time
+from types import SimpleNamespace
+from typing import Any
+
+from PIL import Image
+
+from kiapi.capabilities import attach_mflux_progress
+from kiapi.core.file import FileStore
+from kiapi.core.model import ModelSpec
+from kiapi.core.workdir import create_work_dir
+
+from .._views.upscale_params import UpscaleParams
+
+FEATURES = {"image", "upscale"}
+
+
+def _variant(spec: ModelSpec) -> str:
+    return spec.name
+
+
+def _repo_base(spec: ModelSpec) -> str:
+    base, _sep, _variant_name = spec.repo.partition("#")
+    return base
+
+
+def _model_config(spec: ModelSpec) -> Any:
+    from mflux.models.common.config import ModelConfig  # type: ignore
+
+    return (
+        ModelConfig.seedvr2_7b() if _variant(spec) == "7b" else ModelConfig.seedvr2_3b()
+    )
+
+
+def _resolution(value: int | str) -> int | Any:
+    if isinstance(value, int):
+        return value
+    if value.endswith("x"):
+        from mflux.utils.scale_factor import ScaleFactor  # type: ignore
+
+        return ScaleFactor.parse(value)
+    return int(value)
+
+
+def build_model(spec: ModelSpec, *, quantize: int | None) -> Any:
+    from mflux.models.seedvr2 import SeedVR2  # type: ignore
+
+    model_config = _model_config(spec)
+    return SeedVR2(
+        quantize=quantize, model_path=_repo_base(spec), model_config=model_config
+    )
+
+
+def load(spec: ModelSpec) -> SimpleNamespace:
+    from .._settings import settings_manager
+
+    quantize = settings_manager.get_settings().default_quantize
+    model = build_model(spec, quantize=quantize)
+    return SimpleNamespace(model=model, spec=spec, quantize=quantize)
+
+
+def _store_image(
+    generated: Any,
+    *,
+    params: UpscaleParams,
+    files: FileStore,
+    meta: dict[str, Any],
+) -> dict[str, Any]:
+    pil = generated.image if hasattr(generated, "image") else generated
+    fmt = params.format
+    ext = "jpg" if fmt == "jpeg" else fmt
+    tmp_dir = create_work_dir("image/seedvr2")
+    out_path = tmp_dir / f"image.{ext}"
+    save_kwargs: dict[str, Any] = {}
+    if fmt in ("jpeg", "webp"):
+        save_kwargs["quality"] = params.quality
+    if fmt == "jpeg":
+        pil = pil.convert("RGB")
+    pil.save(out_path, format=fmt.upper(), **save_kwargs)
+    if not out_path.exists():
+        raise RuntimeError("upscale finished but no image file was produced")
+    rec = files.put_path(
+        out_path,
+        filename=f"seedvr2_{int(time.time())}.{ext}",
+        content_type=f"image/{fmt}",
+        meta=meta,
+        move=True,
+    )
+    return {"file_id": rec.file_id, "image_bytes": rec.size, **meta}
+
+
+def _upscale(model: Any, params: UpscaleParams, files: FileStore) -> dict[str, Any]:
+    resolution = _resolution(params.resolution)
+    with Image.open(params.image_path) as src:
+        input_size = src.size
+
+    t0 = time.time()
+    attach_mflux_progress(model)
+    image = model.generate_image(
+        seed=params.seed,
+        image_path=params.image_path,
+        resolution=resolution,
+        softness=params.softness,
+    )
+    total_s = round(time.time() - t0, 2)
+
+    pil = image.image if hasattr(image, "image") else image
+    output_size = pil.size
+    meta = {
+        "model": params.model,
+        "params": params.model_dump(exclude={"image_path"}),
+        "input_width": input_size[0],
+        "input_height": input_size[1],
+        "width": output_size[0],
+        "height": output_size[1],
+        "timings": {"total_s": total_s},
+    }
+    return _store_image(image, params=params, files=files, meta=meta)
+
+
+def run(
+    payload: SimpleNamespace, params: UpscaleParams, files: FileStore
+) -> dict[str, Any]:
+    return _upscale(payload.model, params, files)
+
+
+def run_transient(
+    spec: ModelSpec, params: UpscaleParams, files: FileStore
+) -> dict[str, Any]:
+    model = build_model(spec, quantize=params.quantize)
+    try:
+        return _upscale(model, params, files)
+    finally:
+        del model
+        gc.collect()
+        import mlx.core as mx
+
+        mx.clear_cache()

kiapi/capabilities/seedvr2/_operations/handle_upscale.py

@@ -0,0 +1,42 @@
+"""SeedVR2 service entry: resolve input image → run resident/transient model.
+
+Body of the seedvr2 job's worker-thread thunk. Stores the upscaled image in the
+global file store; the job's ``artifacts`` is the produced file_id.
+"""
+
+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
+from kiapi.core.model import model_registry
+
+from .._settings import settings_manager
+from .._views.upscale_request import UpscaleRequest
+from .is_quantize_override import is_quantize_override
+from .resolve_upscale_params import resolve_upscale_params
+
+
+def handle_upscale(
+    ctx: AppContext, req: UpscaleRequest
+) -> tuple[JobResult, list[FileID]]:
+    settings = settings_manager.get_settings()
+    spec = model_registry.resolve("seedvr2", req.model)
+    ctx.ensure_model_ready(spec)
+    image = resolve_file_ref(ctx.file_store, req.image, kind="image")
+    override = is_quantize_override(settings, req)
+    params = resolve_upscale_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)
+        result = spec.module.run_transient(spec, params, ctx.file_store)
+    else:
+        payload = ctx.memory_manager.acquire(spec)
+        result = spec.module.run(payload, params, ctx.file_store)
+
+    return result, [result["file_id"]]

kiapi/capabilities/seedvr2/_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_upscale`` runs a one-off transient model instead.
+"""
+
+from .._settings import SeedVR2Settings
+from .._views.upscale_request import UpscaleRequest
+
+
+def is_quantize_override(settings: SeedVR2Settings, req: UpscaleRequest) -> bool:
+    return req.quantize is not None and req.quantize != settings.default_quantize

kiapi/capabilities/seedvr2/_operations/resolve_upscale_params.py

@@ -0,0 +1,30 @@
+"""Merge an upscale request with settings defaults into the complete UpscaleParams."""
+
+import random
+
+from .._settings import SeedVR2Settings
+from .._views.upscale_params import UpscaleParams
+from .._views.upscale_request import UpscaleRequest
+
+
+def resolve_upscale_params(
+    settings: SeedVR2Settings,
+    req: UpscaleRequest,
+    *,
+    variant: str,
+    image_file_id: str,
+    image_path: str,
+) -> UpscaleParams:
+    return UpscaleParams(
+        model=variant,
+        image_file_id=image_file_id,
+        image_path=image_path,
+        resolution=req.resolution,
+        softness=req.softness,
+        seed=req.seed if req.seed is not None else random.randint(0, 2**31 - 1),
+        quantize=req.quantize
+        if req.quantize is not None
+        else settings.default_quantize,
+        format=req.format,
+        quality=req.quality,
+    )

kiapi/capabilities/seedvr2/_settings.py

@@ -0,0 +1,74 @@
+"""SeedVR2 capability defaults + caps, read from ``KIAPI_SEEDVR2_`` env vars."""
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class SeedVR2Settings(BaseSettings):
+    """Settings for the SeedVR2 upscaling model, scale factors, and resolution limits."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_SEEDVR2_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    model_repo: str = Field(
+        default="numz/SeedVR2_comfyUI",
+        title="SeedVR2 model repo",
+        description="Hugging Face repo ID for the SeedVR2 model used for video and image upscaling.",
+    )
+
+    # --------------------------------------------------
+    # upscale
+    # --------------------------------------------------
+
+    default_resolution: str = Field(
+        default="2x",
+        title="Default output resolution",
+        description=(
+            "Upscaling resolution used when a request omits resolution.\n"
+            "Use a scale notation such as 2x, or another resolution value "
+            "accepted by the implementation."
+        ),
+    )
+
+    default_softness: float = Field(
+        default=0.0,
+        title="Default softness",
+        description=(
+            "Correction amount used when a request omits softness.\n"
+            "0.0 is the baseline; higher values produce softer results."
+        ),
+    )
+
+    default_quantize: int | None = Field(
+        default=8,
+        title="Default quantization bits",
+        description=(
+            "Default quantization bit count used when loading the SeedVR2 model.\n"
+            "Use a bit count such as 8, or None to disable quantization."
+        ),
+    )
+
+    max_scale: float = Field(
+        default=4.0,
+        title="Maximum scale factor",
+        description="Upper limit for the upscaling scale factor accepted in a request.",
+    )
+
+    max_resolution: int = Field(
+        default=2048,
+        title="Maximum output resolution",
+        description="Upper pixel limit for the output width or height accepted in a request.",
+    )
+
+    min_resolution: int = Field(
+        default=16,
+        title="Minimum output resolution",
+        description="Lower pixel limit for the output width or height accepted in a request.",
+    )
+
+
+settings_manager = SettingsManager(SeedVR2Settings)

kiapi/capabilities/seedvr2/_views/upscale_params.py

@@ -0,0 +1,27 @@
+"""The complete contract for one SeedVR2 image upscale run.
+
+Built from settings + request by ``resolve_upscale_params``; the model needs
+nothing else to produce the image and its metadata. ``image_path`` is the locally
+resolved input file, kept out of the recorded metadata in favour of
+``image_file_id``. ``resolution`` keeps the request value (an int shortest-edge
+target or a scale like ``"2x"``); the model layer parses the scale form.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+
+class UpscaleParams(BaseModel):
+    kind: Literal["upscale"] = "upscale"
+    model: str
+
+    image_file_id: str
+    image_path: str
+    resolution: int | str
+    softness: float
+    seed: int
+    quantize: int | None
+
+    format: Literal["png", "jpeg", "webp"]
+    quality: int

kiapi/capabilities/seedvr2/_views/upscale_request.py

@@ -0,0 +1,94 @@
+"""SeedVR2 upscale request model.
+
+Image → upscaled image (diffusion super-resolution). 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 UpscaleRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = Field(
+        default=None,
+        description=(
+            "Model variant (see GET /v1/image/seedvr2/models). Omit for the "
+            "default `3b`; `7b` is the larger, higher-quality variant. Aliases "
+            "`seedvr2-3b` / `seedvr2-7b` are also accepted."
+        ),
+    )
+    mode: Literal["sync", "async"] = Field(
+        default="sync",
+        description=(
+            "`sync` waits for the image (504 on timeout); `async` returns 202 with "
+            "a job_id immediately — poll GET /v1/jobs/{job_id}."
+        ),
+    )
+
+    image: FileRef = Field(
+        ...,
+        description=(
+            "Input image to upscale (Files-API file id, http(s) URL, or data URL)."
+        ),
+    )
+    resolution: int | str = Field(
+        default="2x",
+        description=(
+            "Target size: an integer shortest-edge pixel target (16..2048), or a "
+            "scale factor string like `2x` / `1.5x` (> 0 and at most 4.0x). "
+            "Default `2x`."
+        ),
+    )
+    softness: float = Field(
+        default=0.0,
+        ge=0.0,
+        le=1.0,
+        description=(
+            "Detail smoothing in 0..1. 0 keeps maximum sharpness; higher values "
+            "soften the result. Default 0.0."
+        ),
+    )
+    seed: int | None = Field(
+        default=None,
+        description=(
+            "Random seed for reproducibility. Omit for a random seed (the resolved "
+            "seed is recorded in the result `params`)."
+        ),
+    )
+    quantize: int | None = Field(
+        default=None,
+        description=(
+            "Quantization bits, one of {3, 4, 5, 6, 8}. Omit to use the resident "
+            "model's quantization (q8). A differing value builds a one-off "
+            "transient model (slower, not reused)."
+        ),
+    )
+
+    format: Literal["png", "jpeg", "webp"] = Field(
+        default="png",
+        description="Output image encoding for the produced file.",
+    )
+    quality: int = Field(
+        default=90,
+        ge=1,
+        le=100,
+        description="Encoder quality 1..100 (used for jpeg/webp; ignored for png).",
+    )
+
+    def gen_params(self) -> dict:
+        return {
+            "image": self.image.model_dump(mode="json"),
+            "resolution": self.resolution,
+            "softness": self.softness,
+            "seed": self.seed,
+            "quantize": self.quantize,
+            "format": self.format,
+            "quality": self.quality,
+        }

kiapi/capabilities/web/README.ja.md

@@ -0,0 +1,187 @@
+# web
+
+**日本語**
+
+SearXNG と Crawl4AI を kiapi から使うための **web capability** です。
+
+- **search**: SearXNG で Web 検索し、JSON の検索結果を返す。
+- **fetch**: Crawl4AI で単一 URL をレンダリングし、Markdown または PDF を返す。
+
+web は、SearXNG / Crawl4AI を **resident subprocess model** として扱います。
+初回リクエスト時に `docker run --rm` を foreground subprocess として起動し、
+空いている localhost port を割り当て、以降は resident backend として再利用します。
+search / fetch も Job として single-flight worker に投入されます。
+
+## API
+
+| Endpoint | Name | Description |
+|---|---|---|
+| `POST /v1/web/search` | Web 検索 | SearXNG の `/search?format=json` を呼び、検索結果を返す。 |
+| `GET /v1/web/fetch?url=...` | ページ取得 | Crawl4AI で HTML ページをレンダリングし、Markdown または PDF を返す。 |
+| `GET /v1/web/openapi.json` | OpenAPI | 詳細な入出力仕様、使い方、TIPS を返す。 |
+
+## API Docs
+
+- [OpenAPI JSON](https://kiarina.github.io/kiapi/v1/web/openapi.json)
+- [Swagger UI](https://kiarina.github.io/kiapi/v1/web/docs.html)
+- [ReDoc](https://kiarina.github.io/kiapi/v1/web/redoc.html)
+
+## Dependencies
+
+| バックエンド | License | Description |
+|---|---|---|
+| [SearXNG](https://github.com/searxng/searxng) | AGPL-3.0 | メタ検索エンジン。kiapi が `docker run --rm` subprocess として起動する。 |
+| [Crawl4AI](https://github.com/unclecode/crawl4ai) | Apache-2.0 | ブラウザレンダリング + Markdown/PDF 変換。kiapi が `docker run --rm` subprocess として起動する。 |
+
+## Notes
+
+- **resident subprocess model**:
+  SearXNG と Crawl4AI は `ModelSpec(family="web")` として登録されます。`load()` は
+  foreground の `docker run --rm` subprocess を起動し、`release()` は subprocess を停止します。
+- **Job / single-flight**:
+  search / fetch はどちらも Job として single-flight worker で実行されます。
+- **search の query**:
+  `query` は SearXNG へそのまま渡します。`site:`、`!wp`、`!images`、`:ja` などの
+  SearXNG インライン演算子をそのまま使えます。
+- **search の結果数**:
+  SearXNG 自体に `max_results` はないため、kiapi が 1 ページ分の結果を受け取ったあと
+  client-side で切り詰めます。既定は `KIAPI_WEB_DEFAULT_MAX_RESULTS=10` です。
+- **fetch の対象**:
+  `fetch` は HTML ページ専用です。画像、音声、動画、PDF、その他バイナリ URL は
+  バイト列として返さず、`not_html` で拒否します。
+- **fetch の空コンテンツ保護**:
+  HTML として到達できても、レンダリング後に抽出できる本文がない場合は
+  `empty_content` を返します。
+- **fetch の成果物**:
+  fetch の raw body は Files API に保存され、レスポンスヘッダに `X-Kiapi-File-Id` と
+  `X-Kiapi-Job-Id` が付きます。
+- **主なエラー**:
+  validation は 422、`fetch` の `unsupported_accept` は 406、バックエンド到達不能や
+  HTTP エラーは 502、timeout は 504 です。
+- **非決定性**:
+  `search` はライブの検索エンジン結果に依存するため、同じクエリでも結果は変わります。
+
+## Settings
+
+| 変数 | 既定値 | 意味 |
+|---|---|---|
+| `KIAPI_WEB_SEARCH_IMAGE` | `searxng/searxng:latest` | SearXNG backend の Docker image。 |
+| `KIAPI_WEB_FETCH_IMAGE` | `unclecode/crawl4ai:latest` | Crawl4AI backend の Docker image。 |
+| `KIAPI_WEB_BACKEND_READY_TIMEOUT_S` | `60.0` | backend subprocess の起動待ち timeout。 |
+| `KIAPI_WEB_BACKEND_LOG_DIR` | `/tmp/kiapi/logs/web` | backend subprocess の stdout/stderr log 出力先。 |
+| `KIAPI_WEB_TIMEOUT_S` | `10.0` | SearXNG への HTTP timeout。 |
+| `KIAPI_WEB_DEFAULT_CATEGORIES` | `null` | `categories` 省略時の既定。`null` は SearXNG に任せる。 |
+| `KIAPI_WEB_DEFAULT_ENGINES` | `null` | `engines` 省略時の既定。`null` は SearXNG に任せる。 |
+| `KIAPI_WEB_DEFAULT_LANGUAGE` | `null` | `language` 省略時の既定。 |
+| `KIAPI_WEB_DEFAULT_SAFESEARCH` | `null` | `safesearch` 省略時の既定。 |
+| `KIAPI_WEB_DEFAULT_MAX_RESULTS` | `10` | `max_results` 省略時の client-side 切り詰め数。 |
+| `KIAPI_WEB_FETCH_TIMEOUT_S` | `30.0` | Crawl4AI への HTTP timeout。 |
+| `KIAPI_WEB_FETCH_MIN_CONTENT_CHARS` | `1` | これ未満の Markdown を `empty_content` とみなす。 |
+| `KIAPI_WEB_FETCH_FILTER` | `fit` | Crawl4AI `/md` の filter。`fit` は readability-pruned、`raw` は DOM 全体。 |
+| `KIAPI_WEB_FETCH_CACHE` | `0` | Crawl4AI `/md` の cache 指定。 |
+
+## Quickstart
+
+Docker Desktop など、`docker run` が使える状態で kiapi を起動します。SearXNG と
+Crawl4AI は初回リクエスト時に自動起動されます。
+
+### search — Web 検索
+
+```bash
+jq -n \
+'{
+  query: "apple",
+  max_results: 5
+}' |
+curl -sS -X POST http://localhost:${PORT:-8000}/v1/web/search \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq .
+```
+
+カテゴリ、エンジン、言語、新しさで絞り込む:
+
+```bash
+jq -n \
+'{
+  query: "mlx quantization",
+  categories: ["it"],
+  engines: ["google", "duckduckgo"],
+  language: "en",
+  time_range: "month",
+  max_results: 10
+}' |
+curl -sS -X POST http://localhost:${PORT:-8000}/v1/web/search \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq .
+```
+
+SearXNG のインライン演算子を使う:
+
+```bash
+jq -n \
+'{
+  query: "site:github.com kiapi"
+}' |
+curl -sS -X POST http://localhost:${PORT:-8000}/v1/web/search \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '.results[] | {title, url}'
+
+jq -n \
+'{
+  query: "!wp :ja Python"
+}' |
+curl -sS -X POST http://localhost:${PORT:-8000}/v1/web/search \
+-H 'Content-Type: application/json' \
+--data-binary @- |
+jq '.results[] | {title, url}'
+```
+
+レスポンスには `query`、`results`、`answers`、`infoboxes`、`suggestions`、
+`unresponsive_engines` が含まれます。
+
+### fetch — Markdown 取得
+
+```bash
+curl -sS -o page.md \
+"http://localhost:${PORT:-8000}/v1/web/fetch?url=https://github.com/kiarina"
+```
+
+`Accept: text/markdown` は既定なので省略できます。kiapi が検出した upstream の
+content-type は `X-Kiapi-Content-Type` ヘッダで確認できます。
+
+```bash
+curl -sS -D - -o page.md \
+"http://localhost:${PORT:-8000}/v1/web/fetch?url=https://github.com/kiarina" \
+| grep -i x-kiapi
+```
+
+### fetch — PDF 取得
+
+```bash
+curl -sS -H 'Accept: application/pdf' -o page.pdf \
+"http://localhost:${PORT:-8000}/v1/web/fetch?url=https://github.com/kiarina"
+```
+
+### fetch — 非 HTML URL の拒否
+
+```bash
+curl -sS \
+"http://localhost:${PORT:-8000}/v1/web/fetch?url=https://www-media.blazeworks.jp/content/icon/miineko.png" \
+| jq .
+```
+
+レスポンス例:
+
+```json
+{
+  "detail": {
+    "error": "not_html",
+    "message": "Resource is not an HTML page (content-type: image/png).",
+    "url": "https://www-media.blazeworks.jp/content/icon/miineko.png",
+    "content_type": "image/png"
+  }
+}
+```