kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/zimage/_models/zimage.py

@@ -0,0 +1,278 @@
+"""Handler for Z-Image generation/training via mflux."""
+
+import gc
+import json
+import time
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+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 .._settings import settings_manager
+from .._views.generate_params import GenerateParams
+from .._views.train_params import TrainParams
+
+FEATURES = {"text"}
+
+
+def _model_config(variant: str) -> Any:
+    from mflux.models.common.config import ModelConfig  # type: ignore
+
+    return ModelConfig.z_image() if variant == "base" else ModelConfig.z_image_turbo()
+
+
+def build_model(
+    spec: ModelSpec,
+    *,
+    quantize: int | None,
+    lora_paths: list[str] | None = None,
+    lora_scales: list[float] | None = None,
+) -> Any:
+    from mflux.models.z_image import ZImageTurbo  # type: ignore
+
+    return ZImageTurbo(
+        quantize=quantize,
+        model_path=spec.repo,
+        lora_paths=lora_paths or None,
+        lora_scales=lora_scales or None,
+        model_config=_model_config(spec.name),
+    )
+
+
+def load(spec: ModelSpec) -> SimpleNamespace:
+    quantize = settings_manager.get_settings().default_quantize.get(spec.name)
+    model = build_model(spec, quantize=quantize)
+    return SimpleNamespace(model=model, spec=spec, quantize=quantize)
+
+
+def warmup(payload: SimpleNamespace) -> None:
+    payload.model.generate_image(
+        seed=0, prompt="warmup", num_inference_steps=1, width=256, height=256
+    )
+
+
+def _store_image(
+    image: Any,
+    *,
+    fmt: str,
+    quality: int,
+    files: FileStore,
+    meta: dict[str, Any],
+) -> dict[str, Any]:
+    pil = getattr(image, "image", image)
+    ext = "jpg" if fmt == "jpeg" else fmt
+    tmp_dir = create_work_dir("image/zimage")
+    out_path = tmp_dir / f"image.{ext}"
+    save_kwargs: dict[str, Any] = {}
+    if fmt in ("jpeg", "webp"):
+        save_kwargs["quality"] = quality
+    if fmt == "jpeg":
+        pil = pil.convert("RGB")
+    pil.save(out_path, format=fmt.upper(), **save_kwargs)
+    if not out_path.exists():
+        raise RuntimeError("generation finished but no image file was produced")
+    rec = files.put_path(
+        out_path,
+        filename=f"image_{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 _generate(model: Any, params: GenerateParams, files: FileStore) -> dict[str, Any]:
+    t0 = time.time()
+    attach_mflux_progress(model)
+    image = model.generate_image(
+        seed=params.seed,
+        prompt=params.prompt,
+        num_inference_steps=params.steps,
+        width=params.width,
+        height=params.height,
+        guidance=params.guidance,
+        negative_prompt=params.negative_prompt,
+    )
+    total_s = round(time.time() - t0, 2)
+    meta = {
+        "model": params.model,
+        "prompt": params.prompt,
+        "params": params.model_dump(),
+        "width": params.width,
+        "height": params.height,
+        "timings": {"total_s": total_s},
+    }
+    return _store_image(
+        image, fmt=params.format, quality=params.quality, files=files, meta=meta
+    )
+
+
+def run_generate(
+    payload: SimpleNamespace, params: GenerateParams, files: FileStore
+) -> dict[str, Any]:
+    return _generate(payload.model, params, files)
+
+
+def run_generate_transient(
+    spec: ModelSpec,
+    params: GenerateParams,
+    files: FileStore,
+    *,
+    lora_paths: list[str],
+    lora_scales: list[float],
+) -> dict[str, Any]:
+    model = build_model(
+        spec, quantize=params.quantize, lora_paths=lora_paths, lora_scales=lora_scales
+    )
+    try:
+        return _generate(model, params, files)
+    finally:
+        _free(model)
+
+
+def _free(model: Any) -> None:
+    del model
+    gc.collect()
+    import mlx.core as mx
+
+    mx.clear_cache()
+
+
+_MFLUX_TRAIN_NAME = {"turbo": "z-image-turbo", "base": "z-image"}
+_BLOCK_RANGE = {"start": 0, "end": 30}
+
+
+def _default_lora_targets(rank: int, include_ff: bool) -> list[dict[str, Any]]:
+    attn = ["attention.to_q", "attention.to_k", "attention.to_v", "attention.to_out.0"]
+    ff = ["feed_forward.w1", "feed_forward.w2", "feed_forward.w3"] if include_ff else []
+    targets = [
+        {
+            "module_path": f"layers.{{block}}.{name}",
+            "blocks": dict(_BLOCK_RANGE),
+            "rank": rank,
+        }
+        for name in (*attn, *ff)
+    ]
+    targets += [
+        {"module_path": "cap_embedder.1", "rank": rank},
+        {"module_path": "all_final_layer.2-1.linear", "rank": rank},
+    ]
+    return targets
+
+
+def _build_train_config(
+    spec: ModelSpec,
+    params: TrainParams,
+    data_dirname: str,
+    out_dir: Path,
+) -> dict[str, Any]:
+    lora_layers = params.lora_layers or {
+        "targets": _default_lora_targets(params.lora_rank, params.lora_include_ff)
+    }
+
+    return {
+        "model": _MFLUX_TRAIN_NAME[spec.name],
+        "model_path": spec.repo,
+        "data": data_dirname,
+        "seed": params.seed,
+        "steps": params.steps,
+        "guidance": params.guidance,
+        "quantize": params.quantize,
+        "max_resolution": params.max_resolution,
+        "low_ram": params.low_ram,
+        "training_loop": {
+            "num_epochs": params.num_epochs,
+            "batch_size": params.batch_size,
+            "timestep_low": params.timestep_low,
+            "timestep_high": params.timestep_high,
+        },
+        "optimizer": {"name": "AdamW", "learning_rate": params.learning_rate},
+        "checkpoint": {
+            "save_frequency": params.save_frequency,
+            "output_path": str(out_dir),
+        },
+        "lora_layers": lora_layers,
+    }
+
+
+def train(
+    spec: ModelSpec,
+    params: TrainParams,
+    files: FileStore,
+    dataset_dir: str,
+) -> dict[str, Any]:
+    import zipfile
+
+    from mflux.models.common.training.runner import TrainingRunner  # type: ignore
+    from mflux.utils.exceptions import StopTrainingException  # type: ignore
+
+    workdir = Path(dataset_dir).parent
+    out_dir = workdir / "out"
+    config = _build_train_config(spec, params, Path(dataset_dir).name, out_dir)
+    config_path = workdir / "train.json"
+    config_path.write_text(json.dumps(config, indent=2))
+
+    t0 = time.time()
+    try:
+        _adapter, training_spec = TrainingRunner.train(
+            config_path=str(config_path), resume_path=None
+        )
+    except StopTrainingException:
+        from mflux.models.common.training.state.training_spec import (  # type: ignore[import-untyped]
+            TrainingSpec,
+        )
+
+        training_spec = TrainingSpec.resolve(
+            config_path=str(config_path), resume_path=None
+        )
+    total_s = round(time.time() - t0, 2)
+
+    ckpt_dir = Path(training_spec.checkpoint.output_path) / "checkpoints"
+    zips = sorted(ckpt_dir.glob("*_checkpoint.zip"))
+    if not zips:
+        raise RuntimeError(
+            f"training finished but no checkpoint was produced in {ckpt_dir}"
+        )
+    latest = zips[-1]
+
+    adapter_bytes = None
+    with zipfile.ZipFile(latest) as zf:
+        for name in zf.namelist():
+            if name.endswith("_adapter.safetensors"):
+                adapter_bytes = zf.read(name)
+                break
+    if adapter_bytes is None:
+        raise RuntimeError(f"no *_adapter.safetensors inside {latest.name}")
+
+    meta = {
+        "model": spec.name,
+        "kind": "lora_adapter",
+        "config": {
+            k: config[k]
+            for k in (
+                "model",
+                "seed",
+                "steps",
+                "guidance",
+                "quantize",
+                "max_resolution",
+                "low_ram",
+                "training_loop",
+                "optimizer",
+            )
+        },
+        "lora_targets": len(config["lora_layers"]["targets"]),
+        "checkpoint": latest.name,
+        "timings": {"total_s": total_s},
+    }
+    rec = files.put_bytes(
+        adapter_bytes,
+        filename=f"zimage_lora_{spec.name}_{int(time.time())}.safetensors",
+        content_type="application/octet-stream",
+        meta=meta,
+    )
+    return {"adapter_file_id": rec.file_id, "adapter_bytes": rec.size, **meta}

kiapi/capabilities/zimage/_operations/handle_generate.py

@@ -0,0 +1,38 @@
+"""Z-Image text-to-image service entry (worker-thread thunk body)."""
+
+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.generate_request import GenerateRequest
+from .is_quantize_override import is_quantize_override
+from .resolve_generate_params import resolve_generate_params
+from .resolve_lora_params import resolve_lora_params
+
+
+def handle_generate(
+    ctx: AppContext, req: GenerateRequest
+) -> tuple[JobResult, list[FileID]]:
+    settings = settings_manager.get_settings()
+    spec = model_registry.resolve("zimage", req.model)
+    ctx.ensure_model_ready(spec)
+    lora_params = resolve_lora_params(ctx, req.loras)
+    override = is_quantize_override(settings, req, variant=spec.name)
+    params = resolve_generate_params(settings, req, variant=spec.name)
+
+    if lora_params.paths or override:
+        ctx.memory_manager.reserve(spec.weight_gb + spec.peak_headroom_gb)
+        result = spec.module.run_generate_transient(
+            spec,
+            params,
+            ctx.file_store,
+            lora_paths=lora_params.paths,
+            lora_scales=lora_params.scales,
+        )
+    else:
+        payload = ctx.memory_manager.acquire(spec)
+        result = spec.module.run_generate(payload, params, ctx.file_store)
+
+    return result, [result["file_id"]]

kiapi/capabilities/zimage/_operations/handle_train.py

@@ -0,0 +1,78 @@
+"""Z-Image 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("zimage", req.model)
+    ctx.ensure_model_ready(spec)
+
+    dataset_path = get_file_path(ctx.file_store, req.dataset, kind="dataset")
+
+    workdir = create_work_dir("train/zimage")
+    try:
+        data_dir = _extract_dataset(dataset_path, workdir / "data")
+        n_images = _validate_dataset(data_dir)
+        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 (.jpg/.jpeg/.png/.webp) at the top level "
+        "or in a single subfolder"
+    )
+
+
+def _validate_dataset(data_dir: Path) -> 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")
+    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/zimage/_operations/is_quantize_override.py

@@ -0,0 +1,11 @@
+"""Decide whether a request's quantize differs from the variant default."""
+
+from .._settings import ZimageSettings
+from .._views.generate_request import GenerateRequest
+
+
+def is_quantize_override(
+    settings: ZimageSettings, req: GenerateRequest, *, variant: str
+) -> bool:
+    default_quantize = settings.default_quantize.get(variant)
+    return req.quantize is not None and req.quantize != default_quantize

kiapi/capabilities/zimage/_operations/resolve_generate_params.py

@@ -0,0 +1,39 @@
+"""Merge a generate request with settings defaults into GenerateParams."""
+
+import random
+
+from .._settings import ZimageSettings
+from .._views.generate_params import GenerateParams
+from .._views.generate_request import GenerateRequest
+
+
+def resolve_generate_params(
+    settings: ZimageSettings,
+    req: GenerateRequest,
+    *,
+    variant: str,
+) -> GenerateParams:
+    return GenerateParams(
+        model=variant,
+        prompt=req.prompt,
+        negative_prompt=req.negative_prompt,
+        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, 9),
+        guidance=(
+            req.guidance
+            if req.guidance is not None
+            else settings.default_guidance.get(variant)
+        ),
+        quantize=(
+            req.quantize
+            if req.quantize is not None
+            else settings.default_quantize.get(variant)
+        ),
+        format=req.format,
+        quality=req.quality,
+        loras=req.loras,
+    )

kiapi/capabilities/zimage/_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/zimage/_operations/resolve_train_params.py

@@ -0,0 +1,38 @@
+"""Merge a train request with settings defaults into TrainParams."""
+
+from .._settings import ZimageSettings
+from .._views.train_params import TrainParams
+from .._views.train_request import TrainRequest
+
+
+def resolve_train_params(
+    settings: ZimageSettings,
+    req: TrainRequest,
+    *,
+    variant: str,
+) -> TrainParams:
+    steps = req.steps if req.steps is not None else settings.train_steps.get(variant, 9)
+    return TrainParams(
+        model=variant,
+        seed=req.seed,
+        steps=steps,
+        guidance=req.guidance,
+        quantize=req.quantize
+        if req.quantize is not None
+        else settings.train_quantize.get(variant),
+        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, 0)
+        ),
+        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/zimage/_schemas/lora_ref.py

@@ -0,0 +1,27 @@
+"""A single LoRA adapter reference (file + scale) applied to a zimage 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. Train one via POST /v1/image/zimage/train."
+        ),
+    )
+    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."
+        ),
+    )