kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/zimage/_settings.py

@@ -0,0 +1,166 @@
+"""Image (zimage) capability defaults + caps, read from ``KIAPI_ZIMAGE_``."""
+
+from typing import cast
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class ZimageSettings(BaseSettings):
+    """Settings for Z-Image model variants, generation defaults, and LoRA training defaults."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_ZIMAGE_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    turbo_repo: str = Field(
+        default="filipstrand/Z-Image-Turbo-mflux-4bit",
+        title="Turbo model repo",
+        description="Hugging Face repo ID used as the Z-Image turbo variant.",
+    )
+
+    base_repo: str = Field(
+        default="Tongyi-MAI/Z-Image",
+        title="Base model repo",
+        description="Hugging Face repo ID used as the Z-Image base variant.",
+    )
+
+    default_quantize: dict[str, int | None] = Field(
+        default_factory=lambda: {"turbo": None, "base": 8},
+        title="Default quantization bits by variant",
+        description=(
+            "Default quantization bit count used when loading a generation model.\n"
+            "Keys are variant names; values are bit counts such as 8, or None "
+            "for no quantization."
+        ),
+    )
+
+    # --------------------------------------------------
+    # generate
+    # --------------------------------------------------
+
+    default_steps: dict[str, int] = Field(
+        default_factory=lambda: {"turbo": 9, "base": 28},
+        title="Default steps by variant",
+        description=(
+            "Diffusion step count used when a generation request omits steps.\n"
+            "Keys are variant names such as turbo or base; values are step counts."
+        ),
+    )
+
+    default_guidance: dict[str, float | None] = Field(
+        default_factory=lambda: {"turbo": None, "base": 4.0},
+        title="Default guidance by variant",
+        description=(
+            "Guidance strength used when a generation request omits guidance.\n"
+            "Keys are variant names; values are guidance values. None uses the "
+            "model default."
+        ),
+    )
+
+    default_width: int = Field(
+        default=1024,
+        title="Default width",
+        description="Output image width in pixels used when a request omits width.",
+    )
+
+    default_height: int = Field(
+        default=1024,
+        title="Default height",
+        description="Output image height in pixels used when a request omits height.",
+    )
+
+    max_width: int = Field(
+        default=2048,
+        title="Maximum width",
+        description="Upper pixel limit for the output image width accepted in a request.",
+    )
+
+    max_height: int = Field(
+        default=2048,
+        title="Maximum height",
+        description="Upper pixel limit for the output image height accepted in a request.",
+    )
+
+    max_steps: int = Field(
+        default=100,
+        title="Maximum inference steps",
+        description="Upper limit for the diffusion step count accepted in a request.",
+    )
+
+    max_loras: int = Field(
+        default=4,
+        title="Maximum LoRA count",
+        description="Maximum number of LoRA adapters that can be applied in one generation request.",
+    )
+
+    # --------------------------------------------------
+    # train
+    # --------------------------------------------------
+
+    train_reserve_gb: float = Field(
+        default=24.0,
+        title="LoRA training reserved memory GB",
+        description="Memory in GB reserved for the runtime peak of a Z-Image LoRA training job.",
+    )
+
+    train_steps: dict[str, int] = Field(
+        default_factory=lambda: {"turbo": 9, "base": 20},
+        title="Training steps by variant",
+        description=(
+            "Step count used when a LoRA training request omits steps.\n"
+            "Keys are variant names; values are training step counts."
+        ),
+    )
+
+    train_timestep_low: dict[str, int] = Field(
+        default_factory=lambda: {"turbo": 4, "base": 0},
+        title="Training timestep lower bound by variant",
+        description=(
+            "Lower bound of the timestep range used during LoRA training.\n"
+            "Keys are variant names; values are lower-bound timestep values."
+        ),
+    )
+
+    train_quantize: dict[str, int | None] = Field(
+        default_factory=lambda: cast(dict[str, int | None], {"turbo": 8, "base": 8}),
+        title="Training quantization bits by variant",
+        description=(
+            "Quantization bit count applied to the model during LoRA training.\n"
+            "Keys are variant names; values are bit counts such as 8, or None "
+            "for no quantization."
+        ),
+    )
+
+    train_default_epochs: int = Field(
+        default=10,
+        title="Default LoRA training epochs",
+        description="Epoch count used when a LoRA training request omits epochs.",
+    )
+
+    train_default_rank: int = Field(
+        default=16,
+        title="Default LoRA training rank",
+        description=(
+            "LoRA rank used when a LoRA training request omits rank.\n"
+            "Higher values increase expressiveness and memory usage."
+        ),
+    )
+
+    train_default_lr: float = Field(
+        default=1e-4,
+        title="Default LoRA training learning rate",
+        description="Learning rate used when a LoRA training request omits lr.",
+    )
+
+    train_max_epochs: int = Field(
+        default=1000,
+        title="Maximum LoRA training epochs",
+        description="Upper limit for epochs accepted in a LoRA training request.",
+    )
+
+
+settings_manager = SettingsManager(ZimageSettings)

kiapi/capabilities/zimage/_views/generate_params.py

@@ -0,0 +1,23 @@
+"""Resolved Z-Image generation parameters passed to the model layer."""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+from .._schemas.lora_ref import LoraRef
+
+
+class GenerateParams(BaseModel):
+    model: str
+    prompt: str
+    negative_prompt: str | None
+    seed: int
+    width: int
+    height: int
+    steps: int
+    guidance: float | None
+    quantize: int | None
+    format: Literal["png", "jpeg", "webp"]
+    quality: int
+    loras: list[LoraRef]
+    mode: Literal["txt2img"] = "txt2img"

kiapi/capabilities/zimage/_views/generate_request.py

@@ -0,0 +1,127 @@
+"""Z-Image text-to-image request model.
+
+One endpoint serves both sync and async via ``mode``. Defaults for ``steps`` /
+``guidance`` / ``quantize`` are variant-dependent (turbo vs base) and resolved
+server-side when omitted. A ``quantize`` override or any ``loras`` builds a
+one-off transient model (slower, not reused).
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .._schemas.lora_ref import LoraRef
+
+
+class GenerateRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = Field(
+        default=None,
+        description=(
+            "Model variant (see GET /v1/image/zimage/models). Omit for the default "
+            "`turbo` (distilled, few-step, fast); `base` is the full, higher-quality, "
+            "slower variant."
+        ),
+    )
+    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}."
+        ),
+    )
+
+    prompt: str = Field(
+        ...,
+        min_length=1,
+        description="Text prompt describing the image to generate (required).",
+    )
+    negative_prompt: str | None = Field(
+        default=None,
+        description=(
+            "Optional text describing what to avoid in the image. Most effective on "
+            "`base` with guidance > 1; `turbo` is distilled (no guidance) so it has "
+            "little effect there."
+        ),
+    )
+
+    width: int | None = Field(
+        default=None,
+        description=(
+            "Output width in pixels. Omit for the server default (1024). Must be a "
+            "multiple of 16 and at most 2048."
+        ),
+    )
+    height: int | None = Field(
+        default=None,
+        description=(
+            "Output height in pixels. Omit for the server default (1024). Must be a "
+            "multiple of 16 and at most 2048."
+        ),
+    )
+    steps: int | None = Field(
+        default=None,
+        ge=1,
+        description=(
+            "Number of denoising steps (1..100). Omit for the variant default "
+            "(turbo 9, base 28). More steps = slower, sometimes higher quality."
+        ),
+    )
+    guidance: float | None = Field(
+        default=None,
+        description=(
+            "Classifier-free guidance scale. Omit for the variant default "
+            "(turbo: none — distilled; base: 4.0). Higher follows the prompt more "
+            "strictly."
+        ),
+    )
+    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 (turbo is pre-quantized 4-bit; base defaults to 8). "
+            "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).",
+    )
+
+    loras: list[LoraRef] = Field(
+        default_factory=list,
+        description=(
+            "Up to 4 LoRA adapters [{file, scale}] referencing Files-API ids. Any "
+            "lora forces a one-off transient model (slower, not reused)."
+        ),
+    )
+
+    def gen_params(self) -> dict:
+        return {
+            "prompt": self.prompt,
+            "negative_prompt": self.negative_prompt,
+            "width": self.width,
+            "height": self.height,
+            "steps": self.steps,
+            "guidance": self.guidance,
+            "seed": self.seed,
+            "quantize": self.quantize,
+            "format": self.format,
+            "quality": self.quality,
+            "loras": [lr.model_dump() for lr in self.loras],
+        }

kiapi/capabilities/zimage/_views/train_params.py

@@ -0,0 +1,22 @@
+"""Resolved Z-Image LoRA finetune parameters passed to the model layer."""
+
+from pydantic import BaseModel
+
+
+class TrainParams(BaseModel):
+    model: str
+    seed: int
+    steps: int
+    guidance: float
+    quantize: int | None
+    max_resolution: int | None
+    low_ram: bool
+    num_epochs: int
+    batch_size: int
+    timestep_low: int
+    timestep_high: int
+    learning_rate: float
+    save_frequency: int
+    lora_rank: int
+    lora_include_ff: bool
+    lora_layers: dict | None

kiapi/capabilities/zimage/_views/train_request.py

@@ -0,0 +1,101 @@
+"""Z-Image LoRA finetune request model. Always async (long-running).
+
+The dataset is a ZIP uploaded to the Files API, with each image accompanied by a
+same-stem ``.txt`` caption. The trained adapter is stored back in the Files API
+and returned as the job's artifact (``adapter_file_id``).
+"""
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from kiapi.core.file import FileRef
+
+
+class TrainRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = Field(
+        default=None,
+        description=(
+            "Model variant to finetune (see GET /v1/image/zimage/models). Omit for "
+            "the default `turbo`."
+        ),
+    )
+    dataset: FileRef = Field(
+        description=(
+            "Training dataset as a Files-API reference to a ZIP. Images may sit at "
+            "the top level or in a single subfolder; each image needs a same-stem "
+            "`.txt` caption (e.g. `cat.png` + `cat.txt`). `preview*` images are "
+            "exempt from the caption requirement."
+        ),
+    )
+
+    seed: int = Field(default=42, description="Random seed for the training run.")
+    steps: int | None = Field(
+        default=None,
+        description=(
+            "Inference/denoising steps used during training. Omit for the variant "
+            "default (turbo 9, base 20)."
+        ),
+    )
+    guidance: float = Field(
+        default=0.0,
+        description="Guidance scale used during training.",
+    )
+    quantize: int | None = Field(
+        default=None,
+        description=(
+            "Quantization bits during training, one of {3, 4, 5, 6, 8}. Omit for "
+            "the variant default (8)."
+        ),
+    )
+    max_resolution: int | None = Field(
+        default=512,
+        description=(
+            "Images are resized so their long side is at most this many pixels."
+        ),
+    )
+    low_ram: bool = Field(
+        default=True,
+        description="Trade speed for lower peak memory during training.",
+    )
+
+    num_epochs: int = Field(
+        default=10, ge=1, description="Number of passes over the dataset."
+    )
+    batch_size: int = Field(default=1, ge=1, description="Training batch size.")
+    timestep_low: int | None = Field(
+        default=None,
+        description=(
+            "Lower diffusion-timestep bound to sample. Omit for the variant default "
+            "(turbo 4, base 0)."
+        ),
+    )
+    timestep_high: int | None = Field(
+        default=None,
+        description="Upper diffusion-timestep bound to sample. Omit to use `steps`.",
+    )
+    learning_rate: float = Field(default=1e-4, description="AdamW learning rate.")
+    save_frequency: int | None = Field(
+        default=None,
+        description=(
+            "Checkpoint every N steps. Omit to only keep the final checkpoint."
+        ),
+    )
+
+    lora_rank: int = Field(
+        default=16, ge=1, description="Rank of the trained LoRA adapter."
+    )
+    lora_include_ff: bool = Field(
+        default=False,
+        description=(
+            "Also train feed-forward layers (larger adapter), not just attention "
+            "projections."
+        ),
+    )
+    lora_layers: dict | None = Field(
+        default=None,
+        description=(
+            "Advanced: explicit mflux LoRA target spec. Omit to use the built-in "
+            "default targets derived from `lora_rank` / `lora_include_ff`."
+        ),
+    )

kiapi/cli/__init__.py

@@ -0,0 +1,11 @@
+from ._helpers.dedupe_resources import dedupe_resources
+from ._helpers.register_all_capabilities import register_all_capabilities
+from ._helpers.run_resources import run_resources
+from ._helpers.select_specs import select_specs
+
+__all__ = [
+    "dedupe_resources",
+    "register_all_capabilities",
+    "run_resources",
+    "select_specs",
+]

kiapi/cli/_helpers/dedupe_resources.py

@@ -0,0 +1,12 @@
+from collections.abc import Iterable
+
+from kiapi.core.model import ModelSpec
+from kiapi.core.setup import SetupResource
+
+
+def dedupe_resources(specs: Iterable[ModelSpec]) -> list[SetupResource]:
+    resources: dict[str, SetupResource] = {}
+    for spec in specs:
+        for resource in spec.setup_resources:
+            resources.setdefault(resource.key, resource)
+    return list(resources.values())

kiapi/cli/_helpers/register_all_capabilities.py

@@ -0,0 +1,36 @@
+from kiapi.capabilities.acestep import register as register_acestep
+from kiapi.capabilities.audiogen import register as register_audiogen
+from kiapi.capabilities.chat import register as register_chat
+from kiapi.capabilities.depthpro import register as register_depthpro
+from kiapi.capabilities.embedding import register as register_embedding
+from kiapi.capabilities.ernie import register as register_ernie
+from kiapi.capabilities.flux2 import register as register_flux2
+from kiapi.capabilities.ideogram4 import register as register_ideogram4
+from kiapi.capabilities.ltx2 import register as register_ltx2
+from kiapi.capabilities.qwen import register as register_qwen
+from kiapi.capabilities.seedvr2 import register as register_seedvr2
+from kiapi.capabilities.web import register as register_web
+from kiapi.capabilities.zimage import register as register_zimage
+
+_registered = False
+
+
+def register_all_capabilities() -> None:
+    global _registered
+    if _registered:
+        return
+
+    register_acestep()
+    register_audiogen()
+    register_chat()
+    register_depthpro()
+    register_embedding()
+    register_ernie()
+    register_flux2()
+    register_ideogram4()
+    register_ltx2()
+    register_qwen()
+    register_seedvr2()
+    register_web()
+    register_zimage()
+    _registered = True

kiapi/cli/_helpers/run_resources.py

@@ -0,0 +1,24 @@
+import click
+
+from kiapi.core.setup import SetupManager, SetupResource
+
+
+def run_resources(
+    action: str, resources: list[SetupResource], manager: SetupManager
+) -> None:
+    if not resources:
+        click.echo("No setup resources.")
+        return
+
+    for resource in resources:
+        click.echo(f"{action}: {resource.kind} {resource.label}")
+        try:
+            if action == "activate":
+                state = manager.activate(resource)
+            else:
+                state = manager.deactivate(resource)
+        except Exception as exc:
+            raise click.ClickException(f"{resource.label}: {exc}") from exc
+
+        mark = "ok" if state.ready else "missing"
+        click.echo(f"  [{mark}] {state.detail}")

kiapi/cli/_helpers/select_specs.py

@@ -0,0 +1,67 @@
+import click
+import questionary
+
+from kiapi.core.model import ModelSpec, model_registry
+
+
+def select_specs(
+    *,
+    all_selected: bool,
+    domains: tuple[str, ...],
+    families: tuple[str, ...],
+    repos: tuple[str, ...],
+    interactive: bool,
+    action: str = "activate",
+) -> list[ModelSpec]:
+    specs = model_registry.list_specs()
+    if all_selected:
+        return specs
+
+    if domains or families or repos:
+        selected = [
+            spec
+            for spec in specs
+            if _matches(spec, domains=domains, families=families, repos=repos)
+        ]
+        if not selected:
+            raise click.ClickException("No matching models found.")
+        return selected
+
+    if not interactive:
+        raise click.ClickException("Specify --all, --domain, --family, or --repo.")
+
+    choices = [
+        questionary.Choice(title=_spec_choice_title(spec), value=spec) for spec in specs
+    ]
+    selected = questionary.checkbox(
+        f"Select models to {action}:", choices=choices
+    ).ask()
+    if not selected:
+        raise click.ClickException("No models selected.")
+    return list(selected)
+
+
+def _matches(
+    spec: ModelSpec,
+    *,
+    domains: tuple[str, ...],
+    families: tuple[str, ...],
+    repos: tuple[str, ...],
+) -> bool:
+    if domains and spec.domain in domains:
+        return True
+    if families and spec.family in families:
+        return True
+    if repos:
+        values = {spec.repo, *spec.aliases}
+        values.update(resource.label for resource in spec.setup_resources)
+        if values.intersection(repos):
+            return True
+    return False
+
+
+def _spec_choice_title(spec: ModelSpec) -> str:
+    size = sum(resource.disk_gb or 0.0 for resource in spec.setup_resources)
+    size_label = f" ~{size:g} GB" if size else ""
+    default = " default" if spec.default else ""
+    return f"{spec.domain}/{spec.family}/{spec.name}{default}{size_label}"