hanakura 0.1.0__py3-none-any.whl

hanakura/__init__.py

@@ -0,0 +1,33 @@
+"""Hanakura: download and manage Stable Diffusion models.
+
+To embed the web UI and API in another application:
+
+    from hanakura import HanakuraServer, ModelRoot
+
+    hub = HanakuraServer(model_roots=[ModelRoot("/srv/models", layout="comfyui")], port=0)
+    print(hub.start())   # http://127.0.0.1:54123
+"""
+
+from typing import TYPE_CHECKING, Any
+
+from hanakura.version import VERSION
+
+if TYPE_CHECKING:
+    from hanakura.embed import HanakuraServer, ModelRoot, serve
+
+__all__ = ["VERSION", "HanakuraServer", "ModelRoot", "serve"]
+
+_LAZY = {"HanakuraServer", "ModelRoot", "serve"}
+
+
+def __getattr__(name: str) -> Any:
+    """Import the server lazily, so ``hanakura version`` does not pay for FastAPI."""
+    if name in _LAZY:
+        from hanakura import embed
+
+        return getattr(embed, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(__all__)

hanakura/__main__.py

@@ -0,0 +1,6 @@
+"""Allow ``python -m hanakura``."""
+
+from hanakura.cli.app import main
+
+if __name__ == "__main__":
+    main()

hanakura/api/__init__.py

@@ -0,0 +1 @@
+"""FastAPI wrapper over ``hanakura.core``."""

hanakura/api/app.py

@@ -0,0 +1,94 @@
+"""``create_app(services)``: the FastAPI app with routers, the socket, and the web UI."""
+
+import logging
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+
+from hanakura.api.errors import install_error_handlers
+from hanakura.api.openapi import HanakuraAPI
+from hanakura.api.paths import validate_public_base_url
+from hanakura.api.routers import app_info, auth, downloads, hubs, library, settings, sources
+from hanakura.api.security import SecurityMiddleware
+from hanakura.api.sockets import SocketBridge
+from hanakura.api.static import SPAStaticFiles, web_dist_dir
+from hanakura.core.context import Services
+from hanakura.version import VERSION
+
+logger = logging.getLogger(__name__)
+
+
+def normalize_prefix(prefix: str | None) -> str:
+    """``/hub/`` and ``hub`` both become ``/hub``; nothing becomes ``""``."""
+    cleaned = (prefix or "").strip().strip("/")
+    return f"/{cleaned}" if cleaned else ""
+
+
+def create_app(
+    services: Services,
+    bound_host: str | None = None,
+    bound_port: int | None = None,
+    extra_hosts: set[str] | None = None,
+    start_downloads: bool = True,
+    serve_ui: bool = True,
+    api_prefix: str | None = None,
+    public_base_url: str | None = None,
+) -> FastAPI:
+    """Build the application.
+
+    ``api_prefix`` moves everything this package serves — the API, the socket and the web UI —
+    under one path, so a host application can mount it beside its own routes without a clash.
+    The web UI needs no change: it derives its base URL from the URL its own script was loaded
+    from, which already carries the prefix.
+
+    When mounted, the host must enter this app's ``router.lifespan_context`` on its serving
+    event loop and close the services afterwards. ``public_base_url`` supplies a trusted
+    browser-facing URL for OAuth behind a proxy; it does not add a second route prefix.
+    """
+    prefix = normalize_prefix(api_prefix)
+    public_base_url = validate_public_base_url(public_base_url)
+    socket_bridge = SocketBridge(services.events, lambda: services.settings.settings.server.access_token)
+
+    @asynccontextmanager
+    async def lifespan(_app: FastAPI) -> AsyncIterator[None]:
+        await socket_bridge.start()
+        if start_downloads:
+            services.downloads.start()
+        try:
+            yield
+        finally:
+            services.downloads.shutdown()
+            await socket_bridge.stop()
+
+    app = HanakuraAPI(title="Hanakura", version=VERSION, lifespan=lifespan, docs_url=f"{prefix}/docs", redoc_url=None, openapi_url=f"{prefix}/openapi.json")
+    app.state.services = services
+    app.state.bound_port = bound_port
+    app.state.api_prefix = prefix
+    app.state.public_base_url = public_base_url
+    install_error_handlers(app)
+
+    for module in (app_info, settings, auth, sources, hubs, downloads, library):
+        app.include_router(module.router, prefix=f"{prefix}/api")
+    app.mount(f"{prefix}/ws", socket_bridge.asgi_app(prefix), name="socket")
+
+    origins = services.settings.settings.server.allowed_origins
+    if origins:
+        app.add_middleware(CORSMiddleware, allow_origins=origins, allow_methods=["*"], allow_headers=["*"], allow_credentials=True)
+    app.add_middleware(
+        SecurityMiddleware,
+        bound_host=lambda: bound_host or services.settings.settings.server.host,
+        allowed_origins=lambda: services.settings.settings.server.allowed_origins,
+        access_token=lambda: services.settings.settings.server.access_token,
+        extra_hosts=extra_hosts,
+        prefix=prefix,
+    )
+
+    # Registered last, so it only receives paths nothing else matched.
+    dist = web_dist_dir()
+    if serve_ui and (dist / "index.html").is_file():
+        app.mount(f"{prefix}/", SPAStaticFiles(directory=dist, html=True), name="ui")
+    elif serve_ui:
+        logger.warning("No web UI build found at %s; serving the API only", dist)
+    return app

hanakura/api/deps.py

@@ -0,0 +1,14 @@
+"""FastAPI dependencies. Services come from ``app.state`` rather than a global locator."""
+
+from typing import Annotated
+
+from fastapi import Depends, Request
+
+from hanakura.core.context import Services
+
+
+def get_services(request: Request) -> Services:
+    return request.app.state.services
+
+
+ServicesDep = Annotated[Services, Depends(get_services)]

hanakura/api/errors.py

@@ -0,0 +1,41 @@
+"""Map domain exceptions onto one JSON error shape."""
+
+from typing import Any
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from pydantic import Field
+
+from hanakura.core.errors import HanakuraError
+from hanakura.core.record import Record
+
+
+class ErrorResponse(Record):
+    code: str
+    message: str
+    detail: dict[str, Any] = Field(default_factory=dict)
+
+
+def error_response(status: int, code: str, message: str, detail: dict[str, Any] | None = None, headers: dict[str, str] | None = None) -> JSONResponse:
+    return JSONResponse(status_code=status, content={"code": code, "message": message, "detail": detail or {}}, headers=headers)
+
+
+async def _domain_error(_request: Request, exc: Exception) -> JSONResponse:
+    assert isinstance(exc, HanakuraError)
+    headers = None
+    retry_after = exc.detail.get("retry_after") if isinstance(exc.detail, dict) else None
+    if exc.http_status == 429 and retry_after is not None:
+        headers = {"Retry-After": str(int(retry_after))}
+    return error_response(exc.http_status, exc.code, exc.message, exc.detail, headers)
+
+
+def install_error_handlers(app: FastAPI) -> None:
+    app.add_exception_handler(HanakuraError, _domain_error)
+
+
+# Declared on routes so the generated client knows the error shape.
+ERROR_RESPONSES: dict[int | str, dict[str, Any]] = {
+    400: {"model": ErrorResponse},
+    404: {"model": ErrorResponse},
+    409: {"model": ErrorResponse},
+}

hanakura/api/openapi.py

@@ -0,0 +1,29 @@
+"""Add the event models to the OpenAPI schema, so the web UI types socket payloads from it."""
+
+from typing import Any
+
+from fastapi import FastAPI
+from fastapi.openapi.utils import get_openapi
+
+from hanakura.core.events.models import EventBase
+
+
+class HanakuraAPI(FastAPI):
+    """FastAPI with socket event models included in its OpenAPI schema."""
+
+    def openapi(self) -> dict[str, Any]:
+        if self.openapi_schema:
+            return self.openapi_schema
+        schema = get_openapi(title=self.title, version=self.version, routes=self.routes, description=self.description)
+        components = schema.setdefault("components", {}).setdefault("schemas", {})
+        events: dict[str, dict[str, str]] = {}
+        for event_cls in EventBase.get_events():
+            model_schema = event_cls.model_json_schema(ref_template="#/components/schemas/{model}")
+            for name, sub in model_schema.pop("$defs", {}).items():
+                components.setdefault(name, sub)
+            components[event_cls.__name__] = model_schema
+            events[event_cls.__event_name__] = {"$ref": f"#/components/schemas/{event_cls.__name__}"}
+        components["ServerEvents"] = {"type": "object", "properties": events, "required": sorted(events)}
+        schema["components"]["schemas"] = dict(sorted(components.items()))
+        self.openapi_schema = schema
+        return schema

hanakura/api/paths.py

@@ -0,0 +1,33 @@
+"""Deployment paths shared by mounted apps, authentication and Socket.IO."""
+
+from urllib.parse import urlsplit
+
+from starlette import _utils
+from starlette.requests import Request
+from starlette.types import Scope
+
+
+def route_path(scope: Scope) -> str:
+    """Older Starlette strips mount prefixes; newer versions leave that to the router."""
+    resolve = getattr(_utils, "get_route_path", None)
+    return resolve(scope) if resolve is not None else scope.get("path", "")
+
+
+def public_prefix(request: Request) -> str:
+    base = getattr(request.app.state, "public_base_url", None)
+    if base:
+        return urlsplit(base).path.rstrip("/")
+    root = request.scope.get("root_path", "").rstrip("/")
+    return root + (getattr(request.app.state, "api_prefix", "") or "")
+
+
+def validate_public_base_url(value: str | None) -> str | None:
+    """Only the host application may supply this URL; never derive it from forwarded headers."""
+    if value is None:
+        return None
+    url = urlsplit(value)
+    if url.scheme not in ("http", "https") or not url.hostname or url.username is not None or url.password is not None or url.query or url.fragment:
+        raise ValueError("public_base_url must be an absolute HTTP(S) URL without credentials, query or fragment")
+    if any(c.isspace() for c in value) or "\\" in value:
+        raise ValueError("public_base_url contains invalid characters")
+    return value.rstrip("/")

hanakura/api/routers/__init__.py

@@ -0,0 +1 @@
+"""One router per resource, under ``/api/v1``."""

hanakura/api/routers/app_info.py

@@ -0,0 +1,59 @@
+"""App version, health and vocabulary."""
+
+from fastapi import APIRouter, Request
+
+from hanakura.api.deps import ServicesDep
+from hanakura.core.detection.models import BASE_MODELS, MODEL_KINDS
+from hanakura.core.library import LibraryService
+from hanakura.core.library.models import TrashLocation
+from hanakura.core.record import Record
+from hanakura.version import VERSION
+
+router = APIRouter(prefix="/v1/app", tags=["app"])
+
+
+class AppVersion(Record):
+    version: str
+
+
+class Health(Record):
+    status: str
+    auth_required: bool
+
+
+class LabeledValue(Record):
+    value: str
+    label: str
+
+
+class AppMeta(Record):
+    kinds: list[str]
+    base_models: list[LabeledValue]
+    layouts: dict[str, dict[str, str]]
+    roots_locked: bool
+    """True when a host application supplies the model folders; the interface then hides adding,
+    editing and removing them."""
+    trash: TrashLocation
+    api_prefix: str
+
+
+@router.get("/version", operation_id="get_app_version")
+def get_version() -> AppVersion:
+    return AppVersion(version=VERSION)
+
+
+@router.get("/health", operation_id="get_health")
+def get_health(services: ServicesDep) -> Health:
+    return Health(status="ok", auth_required=bool(services.settings.settings.server.access_token))
+
+
+@router.get("/meta", operation_id="get_app_meta")
+def get_meta(request: Request, services: ServicesDep) -> AppMeta:
+    return AppMeta(
+        kinds=MODEL_KINDS,
+        base_models=[LabeledValue(value=k, label=v) for k, v in BASE_MODELS.items()],
+        layouts=LibraryService.layouts(),
+        roots_locked=services.library.roots_locked,
+        trash=services.library.trash_location(),
+        api_prefix=getattr(request.app.state, "api_prefix", "") or "",
+    )

hanakura/api/routers/auth.py

@@ -0,0 +1,126 @@
+"""Civitai authentication: connect, return from the provider, status, disconnect.
+
+The browser never receives an access or refresh token. Only the callback is reachable without
+the application's access token, because the browser arrives from Civitai with no header to
+carry it; that request is still tied to the transaction and the browser that started it.
+"""
+
+import logging
+from urllib.parse import urlsplit
+
+from fastapi import APIRouter, Request, Response
+from fastapi.responses import RedirectResponse
+
+from hanakura.api.deps import ServicesDep
+from hanakura.api.errors import ERROR_RESPONSES
+from hanakura.api.paths import public_prefix
+from hanakura.core.auth.models import AuthStart, CivitaiAuthStatus, MethodRequest, StartRequest
+from hanakura.core.errors import HanakuraError, ValidationError
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter(prefix="/v1/auth/civitai", tags=["auth"], responses=ERROR_RESPONSES)
+
+CALLBACK_PATH = "/api/v1/auth/civitai/callback"
+BINDING_COOKIE = "hanakura_oauth"
+COOKIE_MAX_AGE = 600
+
+
+def callback_url(request: Request, services: ServicesDep) -> str:
+    """The callback to use, taken from the configured allowlist or built for the bound address.
+
+    The address is never taken from the Host header or forwarding headers: a callback has to
+    match one registered with Civitai exactly, and an attacker must not be able to choose it.
+    """
+    base = getattr(request.app.state, "public_base_url", None)
+    allowed = services.settings.settings.auth.civitai.redirect_uris
+    if base:
+        callback = f"{base}{CALLBACK_PATH}"
+        if allowed and callback not in allowed:
+            raise ValidationError("public_base_url does not match a registered auth.civitai.redirect_uris callback")
+        return callback
+    prefix = public_prefix(request)
+    path = f"{prefix}{CALLBACK_PATH}"
+    if allowed:
+        origin = f"{request.url.scheme}://{request.url.netloc}"
+        match = next((uri for uri in allowed if (urlsplit(uri).scheme, urlsplit(uri).netloc) == (request.url.scheme, request.url.netloc) and urlsplit(uri).path == path), None)
+        if match is None:
+            raise ValidationError(
+                f"No callback is registered for {origin}. Add the exact URL to auth.civitai.redirect_uris and register it with Civitai, or use a manual API token."
+            )
+        return match
+    host = services.settings.settings.server.host
+    port = getattr(request.app.state, "bound_port", None) or services.settings.settings.server.port
+    host = "127.0.0.1" if host in ("0.0.0.0", "::", "localhost") else host
+    if ":" in host and not host.startswith("["):
+        host = f"[{host}]"
+    return f"http://{host}:{port}{path}"
+
+
+@router.post("/start", operation_id="start_civitai_auth")
+def start(request: Request, response: Response, services: ServicesDep, body: StartRequest | None = None) -> AuthStart:
+    """Begin an authorization and return the URL for the user's browser."""
+    return_to = (body.return_to if body else None) or None
+    if return_to and (urlsplit(return_to).scheme or return_to.startswith("//")):
+        raise ValidationError("return_to must be a path inside this application")
+    start_result, secret = services.auth.start(callback_url(request, services), return_to)
+    response.set_cookie(
+        BINDING_COOKIE,
+        secret,
+        max_age=COOKIE_MAX_AGE,
+        httponly=True,
+        samesite="lax",
+        # Set only over HTTPS: a loopback deployment is plain HTTP and would drop the cookie.
+        secure=urlsplit(getattr(request.app.state, "public_base_url", None) or str(request.url)).scheme == "https",
+        path=f"{public_prefix(request)}/api/v1/auth/civitai",
+    )
+    return start_result
+
+
+@router.get("/callback", operation_id="civitai_auth_callback", include_in_schema=False)
+def callback(
+    request: Request, services: ServicesDep, code: str | None = None, state: str | None = None, error: str | None = None, error_description: str | None = None
+) -> Response:
+    """Where Civitai sends the browser back. Never shows a token, and keeps none in the URL."""
+    binding = request.cookies.get(BINDING_COOKIE)
+    target = "/#/settings"
+    status = "error"
+    if error:
+        logger.info("Civitai authorization was refused: %s", error)
+        if state:
+            services.auth.transactions.discard(state)
+    elif not code or not state:
+        logger.info("A Civitai callback arrived without a code or state")
+    else:
+        try:
+            target = services.auth.complete(state, code, binding)
+            status = "connected"
+        except HanakuraError as e:
+            logger.info("Completing the Civitai authorization failed: %s", e)
+
+    prefix = public_prefix(request)
+    if prefix and target.startswith("/") and not target.startswith(f"{prefix}/"):
+        target = f"{prefix}{target}"
+    separator = "&" if "?" in target else "?"
+    response = RedirectResponse(f"{target}{separator}civitai={status}", status_code=303)
+    response.delete_cookie(BINDING_COOKIE, path=f"{prefix}/api/v1/auth/civitai")
+    response.headers["Cache-Control"] = "no-store"
+    response.headers["Referrer-Policy"] = "no-referrer"
+    return response
+
+
+@router.get("/status", operation_id="get_civitai_auth_status")
+def status(services: ServicesDep) -> CivitaiAuthStatus:
+    return services.auth.status()
+
+
+@router.post("/disconnect", operation_id="disconnect_civitai_auth")
+def disconnect(services: ServicesDep) -> CivitaiAuthStatus:
+    """Revoke and forget the OAuth credentials. The manual token is left untouched."""
+    return services.auth.disconnect()
+
+
+@router.post("/method", operation_id="set_civitai_auth_method")
+def set_method(services: ServicesDep, body: MethodRequest) -> CivitaiAuthStatus:
+    """Choose which credential to use. Nothing else ever changes this."""
+    return services.auth.set_method(body.method)

hanakura/api/routers/downloads.py

@@ -0,0 +1,54 @@
+"""The download queue."""
+
+from fastapi import APIRouter, status
+
+from hanakura.api.deps import ServicesDep
+from hanakura.api.errors import ERROR_RESPONSES
+from hanakura.core.downloads.models import DownloadCreate, DownloadJob
+
+router = APIRouter(prefix="/v1/downloads", tags=["downloads"], responses=ERROR_RESPONSES)
+
+
+@router.get("", operation_id="list_downloads")
+def list_downloads(services: ServicesDep) -> list[DownloadJob]:
+    return services.downloads.list_jobs()
+
+
+@router.post("", operation_id="create_download", status_code=status.HTTP_201_CREATED)
+def create_download(services: ServicesDep, body: DownloadCreate) -> DownloadJob:
+    return services.downloads.create(body)
+
+
+@router.post("/clear-finished", operation_id="clear_finished_downloads")
+def clear_finished(services: ServicesDep) -> list[int]:
+    return services.downloads.clear_finished()
+
+
+@router.get("/{job_id}", operation_id="get_download")
+def get_download(services: ServicesDep, job_id: int) -> DownloadJob:
+    return services.downloads.get(job_id)
+
+
+@router.post("/{job_id}/pause", operation_id="pause_download")
+def pause_download(services: ServicesDep, job_id: int) -> DownloadJob:
+    return services.downloads.pause(job_id)
+
+
+@router.post("/{job_id}/resume", operation_id="resume_download")
+def resume_download(services: ServicesDep, job_id: int) -> DownloadJob:
+    return services.downloads.resume(job_id)
+
+
+@router.post("/{job_id}/cancel", operation_id="cancel_download")
+def cancel_download(services: ServicesDep, job_id: int) -> DownloadJob:
+    return services.downloads.cancel(job_id)
+
+
+@router.post("/{job_id}/restart", operation_id="restart_download")
+def restart_download(services: ServicesDep, job_id: int) -> DownloadJob:
+    return services.downloads.restart(job_id)
+
+
+@router.delete("/{job_id}", operation_id="delete_download", status_code=status.HTTP_204_NO_CONTENT)
+def delete_download(services: ServicesDep, job_id: int) -> None:
+    services.downloads.remove(job_id)

hanakura/api/routers/hubs.py

@@ -0,0 +1,37 @@
+"""Hugging Face and ModelScope."""
+
+from fastapi import APIRouter, Query
+
+from hanakura.api.deps import ServicesDep
+from hanakura.api.errors import ERROR_RESPONSES
+from hanakura.core.hubs.models import HubInfo, HubPage, HubQuery, RepoDetail, RepoFile
+
+router = APIRouter(prefix="/v1/hubs", tags=["hubs"], responses=ERROR_RESPONSES)
+
+
+@router.get("", operation_id="list_hubs")
+def list_hubs(services: ServicesDep) -> list[HubInfo]:
+    return services.hubs.list_hubs()
+
+
+@router.get("/{hub}/repos", operation_id="search_repos")
+def search_repos(
+    services: ServicesDep,
+    hub: str,
+    query: str = "",
+    sort: str | None = None,
+    limit: int = Query(default=30, ge=1, le=100),
+    cursor: str | None = None,
+) -> HubPage:
+    return services.hubs.search(hub, HubQuery(query=query, sort=sort, limit=limit, cursor=cursor))
+
+
+# Repository ids contain a slash, so they are path parameters of type path.
+@router.get("/{hub}/repos/{repo_id:path}/files", operation_id="list_repo_files")
+def list_repo_files(services: ServicesDep, hub: str, repo_id: str, revision: str | None = None) -> list[RepoFile]:
+    return services.hubs.list_files(hub, repo_id, revision)
+
+
+@router.get("/{hub}/repos/{repo_id:path}", operation_id="get_repo")
+def get_repo(services: ServicesDep, hub: str, repo_id: str, revision: str | None = None) -> RepoDetail:
+    return services.hubs.get_repo(hub, repo_id, revision)