kiapi 0.1.0__py3-none-any.whl

kiapi/api/image/depthpro/router.py

@@ -0,0 +1,81 @@
+"""Image API (depthpro): ``POST /v1/image/depthpro/estimate``."""
+
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import Response
+
+from kiapi.api import (
+    REQUIRE_AUTH,
+    build_job_responses,
+    get_accept,
+    get_ctx,
+    get_worker,
+    register_capability_endpoints,
+    submit_and_maybe_wait,
+)
+from kiapi.capabilities.depthpro import (
+    EstimateRequest,
+    ValidationError,
+    handle_estimate,
+    validate_estimate,
+)
+from kiapi.core.app import AppContext
+from kiapi.core.model import UnknownModelError, model_registry
+from kiapi.core.worker import Worker
+
+from ._views.estimate_response import EstimateResponse
+
+router = APIRouter(dependencies=REQUIRE_AUTH)
+
+
+@router.post(
+    "/v1/image/depthpro/estimate",
+    responses=build_job_responses("image/png", result_model=EstimateResponse),
+)
+async def estimate(
+    req: EstimateRequest,
+    accept: str | None = Depends(get_accept),
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> Response:
+    """Estimate a depth map from one image (Depth Pro, image-to-depth).
+
+    Produces a grayscale depth PNG (near = bright, far = dark) and, unless
+    `include_depth_data` is false, a compressed NPZ holding the raw float depth
+    array plus min/max depth. The same endpoint serves both `sync` and `async`
+    via `mode`.
+
+    Sync content negotiation: when exactly one artifact is produced (i.e.
+    `include_depth_data=false`) and the client does not ask for JSON, the raw PNG
+    bytes are returned with `X-Kiapi-File-Id` / `X-Kiapi-Job-Id` headers — so
+    `curl -o depth.png` just works. Otherwise (JSON requested, two artifacts, or
+    async) the Job JSON is returned, whose `result` follows EstimateResponse.
+
+    A `quantize` value differing from the resident model runs a one-off transient
+    model (slower, not reused). Async returns 202 immediately; poll
+    GET /v1/jobs/{job_id} and fetch artifacts via GET /v1/files/{file_id}.
+    """
+    try:
+        spec = model_registry.resolve("depthpro", req.model)
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+
+    try:
+        validate_estimate(req)
+    except ValidationError as exc:
+        raise HTTPException(status_code=422, detail=str(exc))  # noqa: B904
+
+    def thunk():  # type: ignore
+        return handle_estimate(ctx, req)
+
+    return await submit_and_maybe_wait(
+        ctx,
+        worker,
+        type="depthpro",
+        params=req.gen_params() | {"model": spec.name, "mode": req.mode},
+        thunk=thunk,
+        mode=req.mode,
+        accept=accept,
+    )
+
+
+register_capability_endpoints(router, name="depthpro", base_path="/v1/image/depthpro")

kiapi/api/image/ernie/_views/image_response.py

@@ -0,0 +1,50 @@
+"""``/v1/image/ernie/{generate,edit}`` result payload (the Job ``result`` shape).
+
+Mirrors the dict returned by ``capabilities/ernie/_models/ernie._store_image`` so
+the produced image result is self-describing in OpenAPI. It is an output-side
+projection of the model result (not a model-coupled view), so it lives under
+``api/image/ernie/_views``. ``generate`` and ``edit`` produce the same shape, so
+both endpoints share this one view.
+
+``ImageResponse`` is the sole public element; nested pieces stay private.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class _Timings(BaseModel):
+    total_s: float = Field(
+        description="Wall-clock generation time in seconds (model run only)."
+    )
+
+
+class ImageResponse(BaseModel):
+    """Capability-specific ``result`` for a succeeded ernie generate/edit job."""
+
+    model: str = Field(
+        description="Resolved model variant that produced the image (turbo | base)."
+    )
+    prompt: str = Field(description="The prompt used for the run.")
+
+    file_id: str = Field(
+        description=(
+            "Files-API id of the produced image. Fetch metadata at "
+            "GET /v1/files/{id} or bytes at /download. This is also the artifact "
+            "returned as raw bytes by a single-artifact sync call."
+        )
+    )
+    image_bytes: int = Field(description="Size of the produced image in bytes.")
+
+    width: int = Field(description="Width in pixels of the produced image.")
+    height: int = Field(description="Height in pixels of the produced image.")
+
+    params: dict[str, Any] = Field(
+        description=(
+            "Resolved parameters actually used for the run (model, seed, steps, "
+            "guidance, quantize, scheduler, format, quality, loras, …), so the "
+            "result is reproducible."
+        )
+    )
+    timings: _Timings = Field(description="kiapi extension: server-side timing.")

kiapi/api/image/ernie/router.py

@@ -0,0 +1,181 @@
+"""Image API (ernie): ERNIE-Image generate/edit/train."""
+
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import Response
+
+from kiapi.api import (
+    REQUIRE_AUTH,
+    build_job_responses,
+    build_train_responses,
+    get_accept,
+    get_ctx,
+    get_worker,
+    register_capability_endpoints,
+    submit_and_maybe_wait,
+)
+from kiapi.capabilities.ernie import (
+    EditRequest,
+    GenerateRequest,
+    TrainRequest,
+    ValidationError,
+    handle_edit,
+    handle_generate,
+    handle_train,
+    validate_edit,
+    validate_generate,
+)
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileIDRef
+from kiapi.core.model import UnknownModelError, model_registry
+from kiapi.core.worker import Worker
+
+from ._views.image_response import ImageResponse
+
+router = APIRouter(dependencies=REQUIRE_AUTH)
+
+
+@router.post(
+    "/v1/image/ernie/generate",
+    responses=build_job_responses(
+        "image/png", "image/jpeg", "image/webp", result_model=ImageResponse
+    ),
+)
+async def generate_ernie_endpoint(
+    req: GenerateRequest,
+    accept: str | None = Depends(get_accept),
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> Response:
+    """Generate an image from a text prompt (ERNIE-Image, text-to-image).
+
+    No input image — use `/v1/image/ernie/edit` for image-to-image. The same
+    endpoint serves both `sync` and `async` via `mode`. Variant defaults differ:
+    `turbo` (default) is fast and uses few steps; `base` is higher quality and
+    slower.
+
+    Sync content negotiation: a single image is produced, so unless the client
+    asks for JSON the raw image bytes are returned with `X-Kiapi-File-Id` /
+    `X-Kiapi-Job-Id` headers — `curl -o out.png` just works. With
+    `Accept: application/json` (or async) the Job JSON is returned, whose
+    `result` follows ImageResponse.
+
+    Any `loras` or a `quantize` differing from the resident model runs a one-off
+    transient model (slower, not reused). Async returns 202 immediately; poll
+    GET /v1/jobs/{job_id} and fetch the artifact via GET /v1/files/{file_id}.
+    """
+    try:
+        spec = model_registry.resolve("ernie", req.model)
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+    try:
+        validate_generate(req, variant=spec.name)
+    except ValidationError as exc:
+        raise HTTPException(status_code=422, detail=str(exc))  # noqa: B904
+
+    def thunk():  # type: ignore
+        return handle_generate(ctx, req)
+
+    return await submit_and_maybe_wait(
+        ctx,
+        worker,
+        type="ernie",
+        params=req.gen_params() | {"model": spec.name, "mode": req.mode},
+        thunk=thunk,
+        mode=req.mode,
+        accept=accept,
+    )
+
+
+@router.post(
+    "/v1/image/ernie/edit",
+    responses=build_job_responses(
+        "image/png", "image/jpeg", "image/webp", result_model=ImageResponse
+    ),
+)
+async def edit_ernie_endpoint(
+    req: EditRequest,
+    accept: str | None = Depends(get_accept),
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> Response:
+    """Edit a single input image with a prompt (ERNIE-Image, image-to-image).
+
+    Takes one `image` (FileRef) plus a `prompt`; `image_strength` controls how
+    much of the input is preserved (lower = closer to input). Output must be
+    square by default (a guard around an mflux 0.18.0 non-square img2img issue),
+    overrideable via `KIAPI_ERNIE_EDIT_REQUIRE_SQUARE=0`. The same endpoint
+    serves both `sync` and `async` via `mode`.
+
+    Sync content negotiation, transient-model behavior (`loras` / `quantize`
+    override), and the ImageResponse `result` shape match
+    `/v1/image/ernie/generate`.
+    """
+    try:
+        spec = model_registry.resolve("ernie", req.model)
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+    try:
+        validate_edit(req, variant=spec.name)
+    except ValidationError as exc:
+        raise HTTPException(status_code=422, detail=str(exc))  # noqa: B904
+
+    def thunk():  # type: ignore
+        return handle_edit(ctx, req)
+
+    return await submit_and_maybe_wait(
+        ctx,
+        worker,
+        type="ernie-edit",
+        params=req.gen_params() | {"model": spec.name, "mode": req.mode},
+        thunk=thunk,
+        mode=req.mode,
+        accept=accept,
+    )
+
+
+@router.post("/v1/image/ernie/train", responses=build_train_responses())
+async def train_ernie_endpoint(
+    req: TrainRequest,
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> Response:
+    """Train a LoRA adapter on a captioned image set (ERNIE-Image, always async).
+
+    Long-running, so this endpoint is always async: it returns 202 with a
+    job_id — poll GET /v1/jobs/{job_id}. The `dataset` is a Files-API ZIP whose
+    images each have a same-stem `.txt` caption. On success the job's `result`
+    holds `adapter_file_id` (the trained `.safetensors`, also the artifact),
+    `adapter_bytes`, `num_images`, the resolved `config`, `lora_targets`, and
+    `timings`; the adapter can then be passed back via `loras` on generate/edit.
+    """
+    try:
+        spec = model_registry.resolve("ernie", req.model)
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+    if (
+        isinstance(req.dataset, FileIDRef)
+        and ctx.file_store.get(req.dataset.file_id) is None
+    ):
+        raise HTTPException(
+            status_code=400, detail=f"unknown dataset file_id {req.dataset.file_id!r}"
+        )
+
+    def thunk():  # type: ignore
+        return handle_train(ctx, req)
+
+    return await submit_and_maybe_wait(
+        ctx,
+        worker,
+        type="ernie-train",
+        params={
+            "model": spec.name,
+            "dataset": req.dataset.model_dump(mode="json"),
+            "num_epochs": req.num_epochs,
+            "lora_rank": req.lora_rank,
+        },
+        thunk=thunk,
+        mode="async",
+    )
+
+
+register_capability_endpoints(router, name="ernie", base_path="/v1/image/ernie")

kiapi/api/image/flux2/_views/image_response.py

@@ -0,0 +1,53 @@
+"""``/v1/image/flux2/{generate,edit}`` result payload (the Job ``result`` shape).
+
+Mirrors the dict returned by ``capabilities/flux2/_models/flux2._store_image`` so
+the produced image result is self-describing in OpenAPI. It is an output-side
+projection of the model result (not a model-coupled view), so it lives under
+``api/image/flux2/_views``. ``generate`` and ``edit`` produce the same shape, so
+both endpoints share this one view.
+
+``ImageResponse`` is the sole public element; nested pieces stay private.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class _Timings(BaseModel):
+    total_s: float = Field(
+        description="Wall-clock generation time in seconds (model run only)."
+    )
+
+
+class ImageResponse(BaseModel):
+    """Capability-specific ``result`` for a succeeded flux2 generate/edit job."""
+
+    model: str = Field(
+        description=(
+            "Resolved model variant that produced the image "
+            "(klein-9b | klein-base-4b | klein-base-9b)."
+        )
+    )
+    prompt: str = Field(description="The prompt used for the run.")
+
+    file_id: str = Field(
+        description=(
+            "Files-API id of the produced image. Fetch metadata at "
+            "GET /v1/files/{id} or bytes at /download. This is also the artifact "
+            "returned as raw bytes by a single-artifact sync call."
+        )
+    )
+    image_bytes: int = Field(description="Size of the produced image in bytes.")
+
+    width: int = Field(description="Width in pixels of the produced image.")
+    height: int = Field(description="Height in pixels of the produced image.")
+
+    params: dict[str, Any] = Field(
+        description=(
+            "Resolved parameters actually used for the run (model, seed, steps, "
+            "guidance, quantize, scheduler, format, quality, loras, …), so the "
+            "result is reproducible."
+        )
+    )
+    timings: _Timings = Field(description="kiapi extension: server-side timing.")

kiapi/api/image/flux2/router.py

@@ -0,0 +1,197 @@
+"""Image API (flux2): FLUX.2 generate/edit/train."""
+
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import Response
+
+from kiapi.api import (
+    REQUIRE_AUTH,
+    build_job_responses,
+    build_train_responses,
+    get_accept,
+    get_ctx,
+    get_worker,
+    register_capability_endpoints,
+    submit_and_maybe_wait,
+)
+from kiapi.capabilities.flux2 import (
+    EditRequest,
+    GenerateRequest,
+    TrainRequest,
+    ValidationError,
+    handle_edit,
+    handle_generate,
+    handle_train,
+    settings_manager,
+    validate_edit,
+    validate_generate,
+)
+from kiapi.core.app import AppContext
+from kiapi.core.file import FileIDRef
+from kiapi.core.model import UnknownModelError, model_registry
+from kiapi.core.worker import Worker
+
+from ._views.image_response import ImageResponse
+
+router = APIRouter(dependencies=REQUIRE_AUTH)
+
+
+@router.post(
+    "/v1/image/flux2/generate",
+    responses=build_job_responses(
+        "image/png", "image/jpeg", "image/webp", result_model=ImageResponse
+    ),
+)
+async def generate_flux2_endpoint(
+    req: GenerateRequest,
+    accept: str | None = Depends(get_accept),
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> Response:
+    """Generate an image from a text prompt (FLUX.2 Klein, text-to-image / img2img).
+
+    Plain text-to-image by default; supply `init_image` (+ optional
+    `image_strength`) to run img2img instead — use `/v1/image/flux2/edit` for
+    multi-reference editing. The same endpoint serves both `sync` and `async` via
+    `mode`. Variant defaults differ: `klein-9b` (default) is distilled and fast
+    (few steps, unquantized); `klein-base-4b` / `klein-base-9b` are slower base
+    variants (more steps, q8).
+
+    Sync content negotiation: a single image is produced, so unless the client
+    asks for JSON the raw image bytes are returned with `X-Kiapi-File-Id` /
+    `X-Kiapi-Job-Id` headers — `curl -o out.png` just works. With
+    `Accept: application/json` (or async) the Job JSON is returned, whose
+    `result` follows ImageResponse.
+
+    Any `loras` or a `quantize` differing from the resident model runs a one-off
+    transient model (slower, not reused). Async returns 202 immediately; poll
+    GET /v1/jobs/{job_id} and fetch the artifact via GET /v1/files/{file_id}.
+    """
+    try:
+        spec = model_registry.resolve("flux2", req.model)
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+    try:
+        validate_generate(req, variant=spec.name)
+    except ValidationError as exc:
+        raise HTTPException(status_code=422, detail=str(exc))  # noqa: B904
+
+    def thunk():  # type: ignore
+        return handle_generate(ctx, req)
+
+    return await submit_and_maybe_wait(
+        ctx,
+        worker,
+        type="flux2",
+        params=req.gen_params() | {"model": spec.name, "mode": req.mode},
+        thunk=thunk,
+        mode=req.mode,
+        accept=accept,
+    )
+
+
+@router.post(
+    "/v1/image/flux2/edit",
+    responses=build_job_responses(
+        "image/png", "image/jpeg", "image/webp", result_model=ImageResponse
+    ),
+)
+async def edit_flux2_endpoint(
+    req: EditRequest,
+    accept: str | None = Depends(get_accept),
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> Response:
+    """Edit / combine one or more reference images with a prompt (FLUX.2 Klein).
+
+    Takes a list of `images` (FileRefs, multi-reference editing) plus a `prompt`;
+    `image_strength` controls how much of the input is preserved (lower = closer
+    to input). For plain image-to-image from a single seed image, use `init_image`
+    on `/v1/image/flux2/generate` instead. The same endpoint serves both `sync`
+    and `async` via `mode`.
+
+    Sync content negotiation, transient-model behavior (`loras` / `quantize`
+    override), and the ImageResponse `result` shape match
+    `/v1/image/flux2/generate`.
+    """
+    try:
+        spec = model_registry.resolve("flux2", req.model)
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+    try:
+        validate_edit(req, variant=spec.name)
+    except ValidationError as exc:
+        raise HTTPException(status_code=422, detail=str(exc))  # noqa: B904
+
+    def thunk():  # type: ignore
+        return handle_edit(ctx, req)
+
+    return await submit_and_maybe_wait(
+        ctx,
+        worker,
+        type="flux2-edit",
+        params=req.gen_params() | {"model": spec.name, "mode": req.mode},
+        thunk=thunk,
+        mode=req.mode,
+        accept=accept,
+    )
+
+
+@router.post("/v1/image/flux2/train", responses=build_train_responses())
+async def train_flux2_endpoint(
+    req: TrainRequest,
+    ctx: AppContext = Depends(get_ctx),
+    worker: Worker = Depends(get_worker),
+) -> Response:
+    """Train a LoRA adapter on an image dataset (FLUX.2 Klein base, always async).
+
+    Long-running, so this endpoint is always async: it returns 202 with a
+    job_id — poll GET /v1/jobs/{job_id}. Only the base variants
+    (`klein-base-4b` default, `klein-base-9b`) are trainable. `training_mode`
+    selects the dataset layout: `text` uses captioned images (same-stem `.txt`),
+    `edit` uses `*_in.*` / `*_out.*` pairs with `*_in.txt` prompts; the `dataset`
+    is a Files-API ZIP. On success the job's `result` holds `adapter_file_id`
+    (the trained `.safetensors`, also the artifact), `adapter_bytes`,
+    `training_mode`, the resolved `config`, `lora_targets`, `checkpoint`, and
+    `timings`; the adapter can then be passed back via `loras` on generate/edit.
+    """
+    settings = settings_manager.get_settings()
+
+    try:
+        spec = model_registry.resolve(
+            "flux2", req.model or settings.train_default_model
+        )
+    except UnknownModelError as exc:
+        raise HTTPException(status_code=400, detail=str(exc))  # noqa: B904
+    if "base" not in spec.name:
+        raise HTTPException(
+            status_code=400,
+            detail="Flux2 training requires klein-base-4b or klein-base-9b",
+        )
+    if (
+        isinstance(req.dataset, FileIDRef)
+        and ctx.file_store.get(req.dataset.file_id) is None
+    ):
+        raise HTTPException(
+            status_code=400, detail=f"unknown dataset file_id {req.dataset.file_id!r}"
+        )
+
+    def thunk():  # type: ignore
+        return handle_train(ctx, req)
+
+    return await submit_and_maybe_wait(
+        ctx,
+        worker,
+        type="flux2-train",
+        params={
+            "model": spec.name,
+            "dataset": req.dataset.model_dump(mode="json"),
+            "training_mode": req.training_mode,
+            "num_epochs": req.num_epochs,
+            "lora_rank": req.lora_rank,
+        },
+        thunk=thunk,
+        mode="async",
+    )
+
+
+register_capability_endpoints(router, name="flux2", base_path="/v1/image/flux2")

kiapi/api/image/ideogram4/_views/image_response.py

@@ -0,0 +1,55 @@
+"""``/v1/image/ideogram4/generate`` result payload (the Job ``result`` shape).
+
+Mirrors the dict returned by ``capabilities/ideogram4/_models/ideogram4._store_image``
+so the produced image result is self-describing in OpenAPI. It is an output-side
+projection of the model result (not a model-coupled view), so it lives under
+``api/image/ideogram4/_views``.
+
+``ImageResponse`` is the sole public element; nested pieces stay private.
+"""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class _Timings(BaseModel):
+    total_s: float = Field(
+        description="Wall-clock generation time in seconds (model run only)."
+    )
+
+
+class ImageResponse(BaseModel):
+    """Capability-specific ``result`` for a succeeded ideogram4 generate job."""
+
+    model: str = Field(description="Resolved model variant that produced the image.")
+    prompt: str | dict[str, Any] = Field(
+        description="The prompt used for the run (plain text or JSON caption)."
+    )
+
+    file_id: str = Field(
+        description=(
+            "Files-API id of the produced image. Fetch metadata at "
+            "GET /v1/files/{id} or bytes at /download. This is also the artifact "
+            "returned as raw bytes by a single-artifact sync call."
+        )
+    )
+    image_bytes: int = Field(description="Size of the produced image in bytes.")
+
+    width: int = Field(description="Width in pixels of the produced image.")
+    height: int = Field(description="Height in pixels of the produced image.")
+
+    params: dict[str, Any] = Field(
+        description=(
+            "Resolved parameters actually used for the run (model, prompt, preset, "
+            "seed, width, height, quantize, format, quality, …), so the result is "
+            "reproducible."
+        )
+    )
+    safety_filter_note: str = Field(
+        description=(
+            "Reminder that Ideogram 4 may return an 'Image blocked by safety filter' "
+            "image, which kiapi stores as the artifact rather than raising an error."
+        )
+    )
+    timings: _Timings = Field(description="kiapi extension: server-side timing.")