Source code for novelai_image_mcp.tools.enhance

"""MCP tools: upscale, Director tools, and ControlNet annotation."""

from __future__ import annotations

import base64
from pathlib import Path
from typing import Annotated, Any, Literal

from fastmcp.dependencies import Depends
from fastmcp.exceptions import ToolError
from fastmcp.utilities.types import Image
from pydantic import Field

from ..deps import (
    AppContext,
    get_app_context,
    get_effective_features,
    get_route_scoped_output,
    get_routing_context,
)
from ..features import EXPERIMENTAL_TOOLS
from ..nai import (
    ControlNetModel,
    DirectorTool,
    Emotion,
    EmotionLevel,
    Model,
    NovelAIImage,
)
from ..output import route_output_dir, save_image
from ..routing import RoutingContext
from ..server import mcp
from ._errors import translate_errors
from ._meta import IMAGE_WRITE_ANNOTATIONS

#: NovelAI's line-art sharpening parameter is clamped to this range.
_MAX_DEFRY = 10


def _save_and_return(
    image: NovelAIImage,
    *,
    name: str,
    output_dir: str | Path,
    routing: RoutingContext,
) -> list[Any]:
    """Persist a single image and return an ImageContent block + saved path.

    FastMCP auto-converts its `Image` helper into an `ImageContent` block
    when returned from a tool, so returning `Image(...)` needs no manual
    conversion step. When routing is active the returned text attributes the
    saved path to the request's route.
    """
    path = save_image(image.data, name=name, output_dir=output_dir)
    location = f" [route={routing.route}]" if routing.enabled else ""
    return [
        Image(data=image.data, format="png"),
        f"Saved image{location}: {path}",
    ]




@mcp.tool(
    title="Upscale an image (4x)",
    tags={"novelai", "enhancement"},
    annotations=IMAGE_WRITE_ANNOTATIONS,
)
@translate_errors
async def upscale_image(
    image: str,
    factor: Literal[4] = 4,
    routing: RoutingContext = Depends(get_routing_context),
    route_scoped_output: bool = Depends(get_route_scoped_output),
    app: AppContext = Depends(get_app_context),
) -> list[Any]:
    """Upscale an image 4x using NovelAI's standalone upscaler.

    ``image`` is a base64-encoded PNG/JPEG. NovelAI's standalone upscaler
    offers a single 4x step — the API exposes no scaling factor — so
    ``factor`` only accepts ``4``. The upscaler runs on the V5 model
    configured by ``NOVELAI_UPSCALE_MODEL`` (default
    ``nai-diffusion-5-full``) and consumes Anlas based on the source
    resolution.
    """
    settings = app.settings
    client = app.client
    result = await client.upscale(
        base64.b64decode(image), factor=factor, model=Model(settings.upscale_model)
    )
    return _save_and_return(
        result,
        name="upscale",
        output_dir=route_output_dir(
            settings.output_dir, routing, route_scoped=route_scoped_output
        ),
        routing=routing,
    )





@mcp.tool(
    title="Apply a Director tool",
    tags={"novelai", "enhancement"},
    annotations=IMAGE_WRITE_ANNOTATIONS,
)
@translate_errors
async def director_tool(
    tool: DirectorTool,
    image: str,
    prompt: str = "",
    defry: Annotated[int, Field(ge=0, le=_MAX_DEFRY)] = 0,
    emotion: Emotion | None = None,
    emotion_level: EmotionLevel = EmotionLevel.NORMAL,
    routing: RoutingContext = Depends(get_routing_context),
    route_scoped_output: bool = Depends(get_route_scoped_output),
    features: frozenset[str] = Depends(get_effective_features),
    app: AppContext = Depends(get_app_context),
) -> list[Any]:
    """Apply a NovelAI Director tool to an image.

    ``tool`` is one of: ``lineart``, ``sketch``, ``bg-removal``,
    ``declutter``, ``colorize``, ``emotion``. ``image`` is a
    base64-encoded PNG/JPEG. The ``emotion`` tool additionally requires an
    ``emotion`` name (e.g. ``happy``, ``sad``) and accepts an
    ``emotion_level`` (0-5, where 0 is normal intensity and 5 is weakest).
    ``prompt`` guides ``colorize`` and ``emotion``; ``defry``
    (0-10) sharpens line art.

    ``emotion`` and ``emotion_level`` are **experimental**: they are served only to
    requests that opt in through the insiders header (``X-MCP-Insiders:
    experimental-tools=true``). Without that opt-in, passing either one fails
    with a ``ToolError``; a non-default ``emotion_level`` alone counts as passing it.
    Both parameters stay visible in this tool's schema either way, because a
    tool schema is published when the server starts and cannot vary per
    request.
    """
    settings = app.settings
    client = app.client
    director = DirectorTool(tool)
    emotion_enum = Emotion(emotion) if emotion is not None else None
    level = EmotionLevel(emotion_level)
    if (emotion_enum is not None or level != EmotionLevel.NORMAL) and (
        EXPERIMENTAL_TOOLS not in features
    ):
        raise ToolError(
            "emotion and emotion_level are experimental and are not enabled for "
            "this request; opt in with the X-MCP-Insiders header "
            "(experimental-tools=true)"
        )
    if director is DirectorTool.EMOTION and emotion_enum is None:
        raise ValueError(
            "emotion tool requires an emotion name; expected one of: "
            + ", ".join(item.value for item in Emotion)
        )
    result = await client.director(
        director,
        base64.b64decode(image),
        prompt=prompt,
        defry=defry,
        emotion=emotion_enum,
        emotion_level=level,
    )
    return _save_and_return(
        result,
        name=f"director-{director.value}",
        output_dir=route_output_dir(
            settings.output_dir, routing, route_scoped=route_scoped_output
        ),
        routing=routing,
    )





@mcp.tool(
    title="Annotate an image (ControlNet)",
    tags={"novelai", "enhancement"},
    annotations=IMAGE_WRITE_ANNOTATIONS,
)
@translate_errors
async def annotate_image(
    image: str,
    model: ControlNetModel,
    routing: RoutingContext = Depends(get_routing_context),
    route_scoped_output: bool = Depends(get_route_scoped_output),
    app: AppContext = Depends(get_app_context),
) -> list[Any]:
    """Annotate an image with a ControlNet preprocessor.

    ``image`` is a base64-encoded PNG/JPEG. ``model`` is one of:
    ``hed`` (palette swap), ``midas`` (form lock / depth),
    ``fake_scribble`` (scribbler), ``mlsd`` (building control),
    ``uniformer`` (landscaper). The returned image is the annotation (e.g. a
    line-art map) suitable for use as a ControlNet condition in a subsequent
    generation.
    """
    settings = app.settings
    client = app.client
    controlnet = ControlNetModel(model)
    result = await client.annotate(base64.b64decode(image), controlnet)
    return _save_and_return(
        result,
        name=f"annotate-{controlnet.value}",
        output_dir=route_output_dir(
            settings.output_dir, routing, route_scoped=route_scoped_output
        ),
        routing=routing,
    )



__all__ = ["annotate_image", "director_tool", "upscale_image"]