kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/ernie/_operations/validate_common.py

@@ -0,0 +1,33 @@
+from kiapi.capabilities import ValidationError
+
+from .._settings import ErnieSettings
+from .._views.edit_request import EditRequest
+from .._views.generate_request import GenerateRequest
+
+_MULTIPLE = 16
+
+
+def validate_common(
+    req: GenerateRequest | EditRequest, settings: ErnieSettings, *, variant: str
+) -> None:
+    width = req.width if req.width is not None else settings.default_width
+    height = req.height if req.height is not None else settings.default_height
+
+    if width <= 0 or height <= 0:
+        raise ValidationError("width and height must be positive")
+    if width % _MULTIPLE or height % _MULTIPLE:
+        raise ValidationError(f"width and height must be multiples of {_MULTIPLE}")
+    if width > settings.max_width or height > settings.max_height:
+        raise ValidationError(
+            f"size {width}x{height} exceeds the cap {settings.max_width}x{settings.max_height}"
+        )
+
+    steps = (
+        req.steps if req.steps is not None else settings.default_steps.get(variant, 8)
+    )
+    if steps < 1 or steps > settings.max_steps:
+        raise ValidationError(f"steps must be in 1..{settings.max_steps}")
+    if req.quantize is not None and req.quantize not in (3, 4, 5, 6, 8):
+        raise ValidationError("quantize must be one of 3, 4, 5, 6, 8")
+    if len(req.loras) > settings.max_loras:
+        raise ValidationError(f"at most {settings.max_loras} loras may be applied")

kiapi/capabilities/ernie/_schemas/lora_ref.py

@@ -0,0 +1,27 @@
+"""A single LoRA adapter reference (file + scale) applied to an ernie run."""
+
+from pydantic import BaseModel, Field
+
+from kiapi.core.file import FileRef
+
+
+class LoraRef(BaseModel):
+    """One LoRA adapter to apply, by Files-API reference and blend scale.
+
+    Any lora forces a one-off transient model (slower, not reused). Up to 4 may
+    be passed in a request's ``loras`` list.
+    """
+
+    file: FileRef = Field(
+        description=(
+            "The LoRA adapter (`.safetensors`) to apply, as a Files-API file id, "
+            "http(s) URL, or data URL."
+        ),
+    )
+    scale: float = Field(
+        default=1.0,
+        description=(
+            "Blend strength for this adapter. 1.0 applies it at full weight; lower "
+            "values weaken its effect, higher values exaggerate it."
+        ),
+    )

kiapi/capabilities/ernie/_settings.py

@@ -0,0 +1,151 @@
+"""ERNIE-Image capability defaults + caps, read from ``KIAPI_ERNIE_`` env vars."""
+
+from typing import cast
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class ErnieSettings(BaseSettings):
+    """Settings for ERNIE-Image model variants, generation/edit defaults, and LoRA training."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_ERNIE_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    turbo_model: str = Field(
+        default="baidu/ERNIE-Image-Turbo",
+        title="Turbo model repo",
+        description="Hugging Face repo ID used as the ERNIE-Image turbo variant.",
+    )
+
+    base_model: str = Field(
+        default="baidu/ERNIE-Image",
+        title="Base model repo",
+        description="Hugging Face repo ID used as the ERNIE-Image base variant.",
+    )
+
+    default_steps: dict[str, int] = Field(
+        default_factory=lambda: {"turbo": 8, "base": 50},
+        title="Default steps by variant",
+        description=(
+            "Diffusion step count used when a generation or edit request omits steps.\n"
+            "Keys are variant names such as turbo or base; values are step counts."
+        ),
+    )
+
+    default_guidance: dict[str, float] = Field(
+        default_factory=lambda: {"turbo": 1.0, "base": 4.0},
+        title="Default guidance by variant",
+        description=(
+            "Guidance strength used when a generation or edit request omits guidance.\n"
+            "Keys are variant names; values are guidance values."
+        ),
+    )
+
+    default_quantize: dict[str, int | None] = Field(
+        default_factory=lambda: cast(dict[str, int | None], {"turbo": 8, "base": 8}),
+        title="Default quantization bits by variant",
+        description=(
+            "Default quantization bit count used when loading a model.\n"
+            "Keys are variant names; values are bit counts such as 8, or None "
+            "for no quantization."
+        ),
+    )
+
+    # --------------------------------------------------
+    # generate
+    # --------------------------------------------------
+
+    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 or edit request.",
+    )
+
+    # --------------------------------------------------
+    # edit
+    # --------------------------------------------------
+
+    edit_require_square: bool = Field(
+        default=True,
+        title="Require square edit input",
+        description=(
+            "When true, ERNIE image-to-image/edit avoids non-square sizes.\n"
+            "This conservative default works around mflux VAE latent-packing constraints."
+        ),
+    )
+
+    # --------------------------------------------------
+    # 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 an ERNIE-Image LoRA training job.",
+    )
+
+    train_steps: dict[str, int] = Field(
+        default_factory=lambda: {"turbo": 8, "base": 50},
+        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": 1, "base": 1},
+        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."
+        ),
+    )
+
+
+settings_manager = SettingsManager(ErnieSettings)

kiapi/capabilities/ernie/_views/edit_params.py

@@ -0,0 +1,36 @@
+"""The complete contract for one ERNIE-Image single-image edit (img2img) run.
+
+Built from settings + request by ``resolve_edit_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``.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+from .._schemas.lora_ref import LoraRef
+
+
+class EditParams(BaseModel):
+    kind: Literal["edit"] = "edit"
+    model: str
+
+    prompt: str
+    negative_prompt: str | None
+    image_file_id: str
+    image_path: str
+    image_strength: float
+
+    seed: int
+    width: int
+    height: int
+    steps: int
+    guidance: float
+    quantize: int | None
+    scheduler: str
+
+    format: Literal["png", "jpeg", "webp"]
+    quality: int
+
+    loras: list[LoraRef]

kiapi/capabilities/ernie/_views/edit_request.py

@@ -0,0 +1,142 @@
+"""ERNIE-Image single-image edit (img2img) request model.
+
+Same shape as the generate request plus an input ``image`` (FileRef) and
+``image_strength``. By default the output must be square (mflux 0.18.0 can fail
+on some non-square ERNIE img2img sizes); the guard is overrideable server-side.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from kiapi.core.file import FileRef
+
+from .._schemas.lora_ref import LoraRef
+
+
+class EditRequest(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    model: str | None = Field(
+        default=None,
+        description=(
+            "Model variant (see GET /v1/image/ernie/models). Omit for the default "
+            "`turbo`; `base` is higher-quality but slower. Aliases such as "
+            "`ernie-image-turbo` / `ernie-image` are 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}."
+        ),
+    )
+
+    prompt: str = Field(
+        ...,
+        min_length=1,
+        description="Text prompt describing the desired edit (required).",
+    )
+    negative_prompt: str | None = Field(
+        default=None,
+        description="Optional text describing what to avoid in the result.",
+    )
+    image: FileRef = Field(
+        description=(
+            "Input image to edit, as a Files-API file id, http(s) URL, or data URL."
+        ),
+    )
+    image_strength: float = Field(
+        default=0.4,
+        description=(
+            "How much the input image is preserved, in 0..1. Lower keeps more of "
+            "the input; higher follows the prompt more freely."
+        ),
+    )
+
+    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. By default must equal height (square)."
+        ),
+    )
+    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. By default must equal width (square)."
+        ),
+    )
+    steps: int | None = Field(
+        default=None,
+        ge=1,
+        description=(
+            "Number of denoising steps (1..100). Omit for the variant default "
+            "(turbo 8, base 50)."
+        ),
+    )
+    guidance: float | None = Field(
+        default=None,
+        description=(
+            "Classifier-free guidance scale. Omit for the variant default "
+            "(turbo 1.0, 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 (default 8). A differing value builds a one-off "
+            "transient model (slower, not reused)."
+        ),
+    )
+    scheduler: str = Field(
+        default="linear",
+        description="Noise scheduler. `linear` is the default and only tested value.",
+    )
+
+    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,
+            "image": self.image.model_dump(mode="json"),
+            "image_strength": self.image_strength,
+            "width": self.width,
+            "height": self.height,
+            "steps": self.steps,
+            "guidance": self.guidance,
+            "seed": self.seed,
+            "quantize": self.quantize,
+            "scheduler": self.scheduler,
+            "format": self.format,
+            "quality": self.quality,
+            "loras": [lr.model_dump() for lr in self.loras],
+        }

kiapi/capabilities/ernie/_views/generate_params.py

@@ -0,0 +1,32 @@
+"""The complete contract for one ERNIE-Image txt2img run.
+
+Built from settings + request by ``resolve_generate_params``; the model needs
+nothing else to produce the image and its metadata.
+"""
+
+from typing import Literal
+
+from pydantic import BaseModel
+
+from .._schemas.lora_ref import LoraRef
+
+
+class GenerateParams(BaseModel):
+    kind: Literal["txt2img"] = "txt2img"
+    model: str
+
+    prompt: str
+    negative_prompt: str | None
+
+    seed: int
+    width: int
+    height: int
+    steps: int
+    guidance: float
+    quantize: int | None
+    scheduler: str
+
+    format: Literal["png", "jpeg", "webp"]
+    quality: int
+
+    loras: list[LoraRef]

kiapi/capabilities/ernie/_views/generate_request.py

@@ -0,0 +1,127 @@
+"""ERNIE-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/ernie/models). Omit for the default "
+            "`turbo`; `base` is the higher-quality, slower variant. Aliases such as "
+            "`ernie-image-turbo` / `ernie-image` are 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}."
+        ),
+    )
+
+    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.",
+    )
+
+    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 8, base 50). More steps = slower, sometimes higher quality."
+        ),
+    )
+    guidance: float | None = Field(
+        default=None,
+        description=(
+            "Classifier-free guidance scale. Omit for the variant default "
+            "(turbo 1.0, 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 (default 8). A differing value builds a one-off "
+            "transient model (slower, not reused)."
+        ),
+    )
+    scheduler: str = Field(
+        default="linear",
+        description="Noise scheduler. `linear` is the default and only tested value.",
+    )
+
+    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,
+            "scheduler": self.scheduler,
+            "format": self.format,
+            "quality": self.quality,
+            "loras": [lr.model_dump() for lr in self.loras],
+        }

kiapi/capabilities/ernie/_views/train_params.py

@@ -0,0 +1,28 @@
+"""The complete contract for one ERNIE-Image LoRA finetune run.
+
+Built from settings + request by ``resolve_train_params``; combined with the
+resolved model spec and the on-disk dataset/output dirs, the model needs neither
+settings nor request to assemble its training config.
+"""
+
+from pydantic import BaseModel
+
+
+class TrainParams(BaseModel):
+    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/ernie/_views/train_request.py

@@ -0,0 +1,96 @@
+"""ERNIE-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/ernie/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 8, base 50)."
+        ),
+    )
+    guidance: float | None = Field(
+        default=None,
+        description="Guidance scale. Omit for the variant default (turbo 1.0, base 4.0).",
+    )
+    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 default (1).",
+    )
+    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 / time-embedding 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`."
+        ),
+    )