kiapi 0.1.0__py3-none-any.whl

kiapi/capabilities/ernie/_models/ernie.py

@@ -0,0 +1,356 @@
+"""Handler for ERNIE-Image generation/edit/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 .._views.edit_params import EditParams
+from .._views.generate_params import GenerateParams
+from .._views.train_params import TrainParams
+
+FEATURES = {"text", "image"}
+
+
+def _model_name(spec: ModelSpec) -> str:
+    return "ernie-image" if spec.name == "base" else "ernie-image-turbo"
+
+
+def _model_config(spec: ModelSpec) -> Any:
+    from mflux.models.common.config import ModelConfig  # type: ignore
+
+    return (
+        ModelConfig.ernie_image()
+        if spec.name == "base"
+        else ModelConfig.ernie_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.ernie_image import ErnieImage  # type: ignore
+
+    return ErnieImage(
+        quantize=quantize,
+        model_path=spec.repo,
+        lora_paths=lora_paths or None,
+        lora_scales=lora_scales or None,
+        model_config=_model_config(spec),
+    )
+
+
+def load(spec: ModelSpec) -> SimpleNamespace:
+    from .._settings import settings_manager
+
+    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/ernie")
+    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"ernie_{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,
+        negative_prompt=params.negative_prompt,
+        num_inference_steps=params.steps,
+        width=params.width,
+        height=params.height,
+        guidance=params.guidance,
+        scheduler=params.scheduler,
+    )
+    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 _edit(model: Any, params: EditParams, files: FileStore) -> dict[str, Any]:
+    t0 = time.time()
+    attach_mflux_progress(model)
+    image = model.generate_image(
+        seed=params.seed,
+        prompt=params.prompt,
+        negative_prompt=params.negative_prompt,
+        image_path=params.image_path,
+        image_strength=params.image_strength,
+        num_inference_steps=params.steps,
+        width=params.width,
+        height=params.height,
+        guidance=params.guidance,
+        scheduler=params.scheduler,
+    )
+    total_s = round(time.time() - t0, 2)
+    meta = {
+        "model": params.model,
+        "prompt": params.prompt,
+        "params": params.model_dump(exclude={"image_path"}),
+        "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_edit(
+    payload: SimpleNamespace, params: EditParams, files: FileStore
+) -> dict[str, Any]:
+    return _edit(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 run_edit_transient(
+    spec: ModelSpec,
+    params: EditParams,
+    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 _edit(model, params, files)
+    finally:
+        _free(model)
+
+
+def _free(model: Any) -> None:
+    del model
+    gc.collect()
+    import mlx.core as mx
+
+    mx.clear_cache()
+
+
+_BLOCK_RANGE = {"start": 0, "end": 36}
+
+
+def _default_lora_targets(rank: int, include_ff: bool) -> list[dict[str, Any]]:
+    attn = [
+        "self_attention.to_q",
+        "self_attention.to_k",
+        "self_attention.to_v",
+        "self_attention.to_out.0",
+    ]
+    targets = [
+        {
+            "module_path": f"layers.{{block}}.{name}",
+            "blocks": dict(_BLOCK_RANGE),
+            "rank": rank,
+        }
+        for name in attn
+    ]
+    targets.append({"module_path": "text_proj", "rank": rank})
+    if include_ff:
+        targets += [
+            {
+                "module_path": f"layers.{{block}}.mlp.{name}",
+                "blocks": dict(_BLOCK_RANGE),
+                "rank": rank,
+            }
+            for name in ("gate_proj", "up_proj", "linear_fc2")
+        ]
+        targets += [
+            {"module_path": "time_embedding.linear_1", "rank": rank},
+            {"module_path": "time_embedding.linear_2", "rank": rank},
+            {"module_path": "adaln_modulation", "rank": rank},
+            {"module_path": "final_norm.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": _model_name(spec),
+        "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"ernie_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/ernie/_operations/handle_edit.py

@@ -0,0 +1,39 @@
+"""ERNIE-Image single-image edit service entry (worker-thread thunk body)."""
+
+from kiapi.capabilities import 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 .._settings import settings_manager
+from .._views.edit_request import EditRequest
+from .is_quantize_override import is_quantize_override
+from .resolve_edit_params import resolve_edit_params
+from .resolve_lora_params import resolve_lora_params
+
+
+def handle_edit(ctx: AppContext, req: EditRequest) -> tuple[JobResult, list[FileID]]:
+    settings = settings_manager.get_settings()
+    spec = model_registry.resolve("ernie", req.model)
+    ctx.ensure_model_ready(spec)
+    image_path = get_file_path(ctx.file_store, req.image, kind="image")
+    lora_params = resolve_lora_params(ctx, req.loras)
+    override = is_quantize_override(settings, req, variant=spec.name)
+    params = resolve_edit_params(
+        settings, req, variant=spec.name, image_path=image_path
+    )
+
+    if lora_params.paths or override:
+        ctx.memory_manager.reserve(spec.weight_gb + spec.peak_headroom_gb)
+        result = spec.module.run_edit_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_edit(payload, params, ctx.file_store)
+    return result, [result["file_id"]]

kiapi/capabilities/ernie/_operations/handle_generate.py

@@ -0,0 +1,38 @@
+"""ERNIE-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("ernie", 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/ernie/_operations/handle_train.py

@@ -0,0 +1,77 @@
+"""ERNIE-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("ernie", req.model)
+    ctx.ensure_model_ready(spec)
+
+    dataset_path = get_file_path(ctx.file_store, req.dataset, kind="dataset")
+
+    workdir = create_work_dir("train/ernie")
+    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 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/ernie/_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 ErnieSettings
+from .._views.edit_request import EditRequest
+from .._views.generate_request import GenerateRequest
+
+
+def is_quantize_override(
+    settings: ErnieSettings, 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/ernie/_operations/resolve_edit_params.py

@@ -0,0 +1,44 @@
+"""Merge an edit request with settings defaults into the complete EditParams."""
+
+import random
+
+from .._settings import ErnieSettings
+from .._views.edit_params import EditParams
+from .._views.edit_request import EditRequest
+
+
+def resolve_edit_params(
+    settings: ErnieSettings,
+    req: EditRequest,
+    *,
+    variant: str,
+    image_path: str,
+) -> EditParams:
+    return EditParams(
+        model=variant,
+        prompt=req.prompt,
+        negative_prompt=req.negative_prompt,
+        image_file_id="",
+        image_path=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, 8),
+        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/ernie/_operations/resolve_generate_params.py

@@ -0,0 +1,40 @@
+"""Merge a generate request with settings defaults into the complete GenerateParams."""
+
+import random
+
+from .._settings import ErnieSettings
+from .._views.generate_params import GenerateParams
+from .._views.generate_request import GenerateRequest
+
+
+def resolve_generate_params(
+    settings: ErnieSettings,
+    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, 8),
+        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/ernie/_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/ernie/_operations/resolve_train_params.py

@@ -0,0 +1,43 @@
+"""Merge a train request with settings defaults into the complete TrainParams."""
+
+from .._settings import ErnieSettings
+from .._views.train_params import TrainParams
+from .._views.train_request import TrainRequest
+
+
+def resolve_train_params(
+    settings: ErnieSettings,
+    req: TrainRequest,
+    *,
+    variant: str,
+) -> TrainParams:
+    steps = req.steps if req.steps is not None else settings.train_steps.get(variant, 8)
+    return TrainParams(
+        seed=req.seed,
+        steps=steps,
+        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.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, 1)
+        ),
+        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,
+    )