kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/flux2/_operations/handle_train.py

@@ -0,0 +1,105 @@
+"""FLUX.2 LoRA finetune service entry (always-async worker-thread thunk)."""
+
+import shutil
+import zipfile
+from pathlib import Path
+
+from kiapi.capabilities import CapabilityError, get_file_path
+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 kiapi.core.workdir import create_work_dir
+
+from .._settings import settings_manager
+from .._views.train_request import TrainRequest
+from .resolve_train_params import resolve_train_params
+
+_IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp"}
+
+
+def handle_train(ctx: AppContext, req: TrainRequest) -> tuple[JobResult, list[FileID]]:
+    settings = settings_manager.get_settings()
+    spec = model_registry.resolve("flux2", req.model or settings.train_default_model)
+    ctx.ensure_model_ready(spec)
+    if "base" not in spec.name:
+        raise CapabilityError(
+            "Flux2 training requires a base variant: klein-base-4b or klein-base-9b"
+        )
+
+    dataset_path = get_file_path(ctx.file_store, req.dataset, kind="dataset")
+
+    workdir = create_work_dir("train/flux2")
+    try:
+        data_dir = _extract_dataset(dataset_path, workdir / "data")
+        n_images = _validate_dataset(data_dir, training_mode=req.training_mode)
+        params = resolve_train_params(settings, req, variant=spec.name)
+        ctx.memory_manager.reserve(settings.train_reserve_gb)
+        result = spec.module.train(spec, params, ctx.file_store, str(data_dir))
+        result["num_images"] = n_images
+        return result, [result["adapter_file_id"]]
+    finally:
+        shutil.rmtree(workdir, ignore_errors=True)
+
+
+def _extract_dataset(zip_path: str, dest: Path) -> Path:
+    dest.mkdir(parents=True, exist_ok=True)
+    with zipfile.ZipFile(zip_path) as zf:
+        for member in zf.namelist():
+            target = (dest / member).resolve()
+            if not str(target).startswith(str(dest.resolve())):
+                raise CapabilityError(f"unsafe path in dataset zip: {member!r}")
+        zf.extractall(dest)
+
+    def _has_images(d: Path) -> bool:
+        return any(p.suffix.lower() in _IMAGE_EXTS for p in d.iterdir() if p.is_file())
+
+    if _has_images(dest):
+        return dest
+    subdirs = [p for p in dest.iterdir() if p.is_dir()]
+    if len(subdirs) == 1 and _has_images(subdirs[0]):
+        return subdirs[0]
+    raise CapabilityError(
+        "dataset zip must contain images at the top level or in a single subfolder"
+    )
+
+
+def _validate_dataset(data_dir: Path, *, training_mode: str) -> int:
+    images = [
+        p for p in data_dir.iterdir() if p.is_file() and p.suffix.lower() in _IMAGE_EXTS
+    ]
+    if not images:
+        raise CapabilityError("dataset has no images")
+    out_images = [p for p in images if p.stem.endswith("_out")]
+    if training_mode == "edit":
+        if not out_images:
+            raise CapabilityError(
+                "edit training dataset must contain *_in / *_out image pairs"
+            )
+        count = 0
+        for out_path in out_images:
+            base = out_path.stem[:-4]
+            if not any((data_dir / f"{base}_in{ext}").exists() for ext in _IMAGE_EXTS):
+                raise CapabilityError(
+                    f"missing input image for {out_path.name}: expected {base}_in.*"
+                )
+            if not (data_dir / f"{base}_in.txt").exists():
+                raise CapabilityError(
+                    f"missing prompt file for {base}_in.*: expected {base}_in.txt"
+                )
+            count += 1
+        return count
+
+    if out_images:
+        raise CapabilityError(
+            "text training dataset must not use *_in / *_out edit pairs"
+        )
+    missing = [
+        p.name
+        for p in images
+        if not p.stem.startswith("preview")
+        and not (data_dir / f"{p.stem}.txt").exists()
+    ]
+    if missing:
+        raise CapabilityError(f"images missing a same-stem .txt caption: {missing}")
+    return len([p for p in images if not p.stem.startswith("preview")])

kiapi/capabilities/flux2/_operations/is_quantize_override.py

@@ -0,0 +1,16 @@
+"""Decide whether a request's quantize differs from the variant default.
+
+A True result means the resident model (quantized at its default) can't serve the
+request, so ``handle_*`` runs a one-off transient model instead.
+"""
+
+from .._settings import Flux2Settings
+from .._views.edit_request import EditRequest
+from .._views.generate_request import GenerateRequest
+
+
+def is_quantize_override(
+    settings: Flux2Settings, req: GenerateRequest | EditRequest, *, variant: str
+) -> bool:
+    default_quantize = settings.default_quantize.get(variant)
+    return req.quantize is not None and req.quantize != default_quantize

kiapi/capabilities/flux2/_operations/resolve_edit_params.py

@@ -0,0 +1,43 @@
+"""Merge an edit request with settings defaults into the complete EditParams."""
+
+import random
+
+from .._settings import Flux2Settings
+from .._views.edit_params import EditParams
+from .._views.edit_request import EditRequest
+
+
+def resolve_edit_params(
+    settings: Flux2Settings,
+    req: EditRequest,
+    *,
+    variant: str,
+    image_paths: list[str],
+) -> EditParams:
+    return EditParams(
+        model=variant,
+        prompt=req.prompt,
+        image_file_ids=[],
+        image_paths=image_paths,
+        image_strength=req.image_strength,
+        seed=req.seed if req.seed is not None else random.randint(0, 2**31 - 1),
+        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,
+        steps=req.steps
+        if req.steps is not None
+        else settings.default_steps.get(variant, 4),
+        guidance=(
+            req.guidance
+            if req.guidance is not None
+            else settings.default_guidance.get(variant, 1.0)
+        ),
+        quantize=(
+            req.quantize
+            if req.quantize is not None
+            else settings.default_quantize.get(variant)
+        ),
+        scheduler=req.scheduler,
+        format=req.format,
+        quality=req.quality,
+        loras=req.loras,
+    )

kiapi/capabilities/flux2/_operations/resolve_generate_params.py

@@ -0,0 +1,44 @@
+"""Merge a generate request with settings defaults into the complete GenerateParams."""
+
+import random
+
+from .._settings import Flux2Settings
+from .._views.generate_params import GenerateParams
+from .._views.generate_request import GenerateRequest
+
+
+def resolve_generate_params(
+    settings: Flux2Settings,
+    req: GenerateRequest,
+    *,
+    variant: str,
+    init_image_path: str | None,
+) -> GenerateParams:
+    return GenerateParams(
+        kind="img2img" if init_image_path else "txt2img",
+        model=variant,
+        prompt=req.prompt,
+        init_image_file_id=None,
+        init_image_path=init_image_path,
+        image_strength=req.image_strength,
+        seed=req.seed if req.seed is not None else random.randint(0, 2**31 - 1),
+        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,
+        steps=req.steps
+        if req.steps is not None
+        else settings.default_steps.get(variant, 4),
+        guidance=(
+            req.guidance
+            if req.guidance is not None
+            else settings.default_guidance.get(variant, 1.0)
+        ),
+        quantize=(
+            req.quantize
+            if req.quantize is not None
+            else settings.default_quantize.get(variant)
+        ),
+        scheduler=req.scheduler,
+        format=req.format,
+        quality=req.quality,
+        loras=req.loras,
+    )

kiapi/capabilities/flux2/_operations/resolve_lora_params.py

@@ -0,0 +1,23 @@
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+
+from kiapi.capabilities import get_file_path
+from kiapi.core.app import AppContext
+
+from .._schemas.lora_ref import LoraRef
+
+
+@dataclass
+class LoraParams:
+    paths: list[str] = field(default_factory=list)
+    scales: list[float] = field(default_factory=list)
+
+
+def resolve_lora_params(ctx: AppContext, loras: Iterable[LoraRef]) -> LoraParams:
+    params = LoraParams()
+
+    for lr in loras:
+        params.paths.append(get_file_path(ctx.file_store, lr.file, kind="lora"))
+        params.scales.append(lr.scale)
+
+    return params

kiapi/capabilities/flux2/_operations/resolve_train_params.py

@@ -0,0 +1,42 @@
+"""Merge a train request with settings defaults into the complete TrainParams."""
+
+from .._settings import Flux2Settings
+from .._views.train_params import TrainParams
+from .._views.train_request import TrainRequest
+
+
+def resolve_train_params(
+    settings: Flux2Settings,
+    req: TrainRequest,
+    *,
+    variant: str,
+) -> TrainParams:
+    steps = (
+        req.steps if req.steps is not None else settings.train_steps.get(variant, 40)
+    )
+    return TrainParams(
+        training_mode=req.training_mode,
+        seed=req.seed,
+        steps=steps,
+        guidance=req.guidance,
+        quantize=(
+            req.quantize
+            if req.quantize is not None
+            else settings.train_quantize.get(variant, 8)
+        ),
+        max_resolution=req.max_resolution,
+        low_ram=req.low_ram,
+        num_epochs=req.num_epochs,
+        batch_size=req.batch_size,
+        timestep_low=(
+            req.timestep_low
+            if req.timestep_low is not None
+            else settings.train_timestep_low.get(variant, 25)
+        ),
+        timestep_high=req.timestep_high if req.timestep_high is not None else steps,
+        learning_rate=req.learning_rate,
+        save_frequency=req.save_frequency if req.save_frequency is not None else 10**9,
+        lora_rank=req.lora_rank,
+        lora_include_ff=req.lora_include_ff,
+        lora_layers=req.lora_layers,
+    )

kiapi/capabilities/flux2/_operations/role_spec.py

@@ -0,0 +1,13 @@
+"""Derive a role-specific ModelSpec (its repo carries a ``#<role>`` marker).
+
+The edit role loads a different mflux class than generate, so edit acquires a
+resident model under a spec whose repo is tagged with the role.
+"""
+
+from kiapi.core.model import ModelSpec
+
+
+def role_spec(spec: ModelSpec, role: str) -> ModelSpec:
+    if role == "generate":
+        return spec
+    return spec.model_copy(update={"repo": f"{spec.repo}#{role}"})

kiapi/capabilities/flux2/_operations/split_repo.py

@@ -0,0 +1,11 @@
+"""Split a flux2 spec repo into ``(model_name, role)``.
+
+A role marker is appended to a spec's ``repo`` as ``<repo>#<role>`` so a single
+family registration can yield distinct resident models per role (generate/edit).
+A bare repo means the default ``generate`` role.
+"""
+
+
+def split_repo(repo: str) -> tuple[str, str]:
+    base, sep, role = repo.partition("#")
+    return base, role if sep else "generate"

kiapi/capabilities/flux2/_operations/validate_common.py

@@ -0,0 +1,35 @@
+from kiapi.capabilities import ValidationError
+
+from .._settings import Flux2Settings
+from .._views.edit_request import EditRequest
+from .._views.generate_request import GenerateRequest
+
+_MULTIPLE = 16
+
+
+def validate_common(
+    req: GenerateRequest | EditRequest, settings: Flux2Settings, *, 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, 4)
+    )
+    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")
+    if req.image_strength is not None and not 0 <= req.image_strength <= 1:
+        raise ValidationError("image_strength must be in 0..1")

kiapi/capabilities/flux2/_schemas/lora_ref.py

@@ -0,0 +1,27 @@
+"""A single LoRA adapter reference (file + scale) applied to a flux2 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/flux2/_settings.py

@@ -0,0 +1,181 @@
+"""FLUX.2 capability defaults + caps, read from ``KIAPI_FLUX2_`` env vars."""
+
+from typing import cast
+
+from pydantic import Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from pydantic_settings_manager import SettingsManager
+
+
+class Flux2Settings(BaseSettings):
+    """Settings for FLUX.2 model variants, generation defaults, and LoRA training."""
+
+    model_config = SettingsConfigDict(
+        env_prefix="KIAPI_FLUX2_",
+        extra="ignore",
+        protected_namespaces=(),
+    )
+
+    klein_9b_model: str = Field(
+        default="flux2-klein-9b",
+        title="Klein 9B model name",
+        description="Name registered in the model registry for the FLUX.2 Klein 9B variant.",
+    )
+
+    klein_base_4b_model: str = Field(
+        default="flux2-klein-base-4b",
+        title="Klein Base 4B model name",
+        description="Name registered in the model registry for the FLUX.2 Klein Base 4B variant.",
+    )
+
+    klein_base_9b_model: str = Field(
+        default="flux2-klein-base-9b",
+        title="Klein Base 9B model name",
+        description="Name registered in the model registry for the FLUX.2 Klein Base 9B variant.",
+    )
+
+    # --------------------------------------------------
+    # generate
+    # --------------------------------------------------
+
+    default_steps: dict[str, int] = Field(
+        default_factory=lambda: {
+            "klein-9b": 4,
+            "klein-base-4b": 40,
+            "klein-base-9b": 40,
+        },
+        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 klein-9b; values are step counts."
+        ),
+    )
+
+    default_guidance: dict[str, float] = Field(
+        default_factory=lambda: {
+            "klein-9b": 1.0,
+            "klein-base-4b": 1.0,
+            "klein-base-9b": 1.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],
+            {
+                "klein-9b": None,
+                "klein-base-4b": 8,
+                "klein-base-9b": 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."
+        ),
+    )
+
+    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.",
+    )
+
+    # --------------------------------------------------
+    # train
+    # --------------------------------------------------
+
+    train_default_model: str = Field(
+        default="klein-base-4b",
+        title="Default LoRA training model",
+        description=(
+            "FLUX.2 variant name used when a LoRA training request omits model.\n"
+            "Usually this should be a training-capable variant such as "
+            "klein-base-4b or klein-base-9b."
+        ),
+    )
+
+    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 FLUX.2 LoRA training job.",
+    )
+
+    train_steps: dict[str, int] = Field(
+        default_factory=lambda: {
+            "klein-base-4b": 40,
+            "klein-base-9b": 40,
+        },
+        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: {
+            "klein-base-4b": 25,
+            "klein-base-9b": 25,
+        },
+        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],
+            {
+                "klein-base-4b": 8,
+                "klein-base-9b": 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(Flux2Settings)

kiapi/capabilities/flux2/_views/edit_params.py

@@ -0,0 +1,36 @@
+"""The complete contract for one FLUX.2 multi-reference edit run.
+
+Built from settings + request by ``resolve_edit_params``; the model needs nothing
+else to produce the image and its metadata. ``image_paths`` are the locally
+resolved input files, kept out of the recorded metadata in favour of
+``image_file_ids``.
+"""
+
+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
+    image_file_ids: list[str]
+    image_paths: list[str]
+    image_strength: float | 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]