api.py

from __future__ import annotations
from aiohttp import web
from typing import NamedTuple
from pathlib import Path
import json
import traceback
import re
import logging

from comfy import model_detection, supported_models
import comfy.utils
import folder_paths
import server

from .translation import available_languages, translate
from .krita import WorkflowExchange

input_block_name = "model.diffusion_model.input_blocks.0.0.weight"

model_names = {
    "SD15": "sd15",
    "SD20": "sd20",
    "SD21UnclipL": "sd21",
    "SD21UnclipH": "sd21",
    "SDXLRefiner": "sdxl",
    "SDXL": "sdxl",
    "SSD1B": "ssd1b",
    "SVD_img2vid": "svd",
    "Stable_Cascade_B": "cascade-b",
    "Stable_Cascade_C": "cascade-c",
    "SD3": "sd3",
    "AuraFlow": "aura-flow",
    "HunyuanDiT": "hunyuan-dit",
    "HunyuanDiT1": "hunyuan-dit",
    "Flux": "flux",
    "FluxSchnell": "flux-schnell",
}


class FakeTensor(NamedTuple):
    shape: tuple

    @staticmethod
    def from_dict(d):
        try:
            return FakeTensor(tuple(d["shape"]))
        except KeyError:
            return d


def inspect_diffusion_model(filename: str, model_type: str, is_checkpoint: bool):
    try:
        # Read header of safetensors file
        path = folder_paths.get_full_path(model_type, filename)
        header = comfy.utils.safetensors_header(path)
        if header:
            cfg = json.loads(header.decode("utf-8"))

            # Build a fake "state_dict" from the header info to avoid reading the full weights
            for key in cfg:
                if not key == "__metadata__":
                    cfg[key] = FakeTensor.from_dict(cfg[key])

            # Reuse Comfy's model detection
            prefix = model_detection.unet_prefix_from_state_dict(cfg)
            if not is_checkpoint:
                temp_sd = comfy.utils.state_dict_prefix_replace(cfg, {prefix: ""}, filter_keys=True)
                if len(temp_sd) > 0:
                    cfg = temp_sd
                prefix = ""
            try:  # latest ComfyUI takes 2 args
                unet_config = model_detection.detect_unet_config(cfg, prefix)
            except TypeError as e:  # older ComfyUI versions take 3 args
                raise TypeError(f"{e} when calling detect_unet_config - old version of ComfyUI?")

            # Get input count to detect inpaint models
            if input_block := cfg.get(input_block_name, None):
                input_count = input_block.shape[1]
            else:
                input_count = 4

            # Find a matching base model depending on unet config
            base_model = model_detection.model_config_from_unet_config(unet_config)
            if base_model is None:
                return {"base_model": "unknown"}

            base_model_class = base_model.__class__
            base_model_name = model_names.get(base_model_class.__name__, "unknown")
            return {
                "base_model": base_model_name,
                "is_inpaint": base_model_name in ["sd15", "sdxl"] and input_count > 4,
                "is_refiner": base_model_class is supported_models.SDXLRefiner,
            }
        return {"base_model": "unknown"}
    except Exception as e:
        # traceback.print_exc()
        return {"base_model": "unknown", "error": f"Failed to detect base model: {e}"}


def inspect_models(model_type: str):
    try:
        is_checkpoint = model_type == "checkpoints"
        info = {
            filename: inspect_diffusion_model(filename, model_type, is_checkpoint)
            for filename in folder_paths.get_filename_list(model_type)
        }
        return web.json_response(info)
    except Exception as e:
        traceback.print_exc()
        return web.json_response(dict(error=str(e)), status=500)


def has_invalid_folder_name(folder_name: str):
    valid_names = list(folder_paths.folder_names_and_paths.keys())
    if folder_name not in valid_names:
        return web.json_response(
            dict(error=f"Invalid folder path, must be one of {', '.join(valid_names)}"),
            status=400,
        )
    return None


def has_invalid_filename(filename: str):
    if not filename.lower().endswith((".sft", ".safetensors")):
        return web.json_response(dict(error="File extension must be .safetensors"), status=400)
    if not filename or not filename.strip() or len(filename) > 255:
        return web.json_response(dict(error="Invalid filename"), status=400)
    if any(char in filename for char in ["..", "/", "\\", "\n", "\r", "\t", "\0"]):
        return web.json_response(dict(error="Invalid filename"), status=400)
    if filename.startswith(".") or not re.match(r"^[a-zA-Z0-9_\-. ]+$", filename):
        return web.json_response(dict(error="Invalid filename"), status=400)
    return None


_server: server.PromptServer | None = getattr(server.PromptServer, "instance", None)
if _server is not None:

    _workflow_exchange = WorkflowExchange(_server)

    @_server.routes.get("/api/etn/model_info/{folder_name}")
    async def model_info(request: web.Request):
        folder_name = request.match_info.get("folder_name", "checkpoints")
        if error := has_invalid_folder_name(folder_name):
            return error
        return inspect_models(folder_name)

    @_server.routes.get("/api/etn/model_info")
    async def api_model_info(request):
        return inspect_models("checkpoints")

    @_server.routes.get("/etn/model_info")
    async def api_model_info(request):
        return inspect_models("checkpoints")

    @_server.routes.get("/api/etn/languages")
    async def languages(request):
        try:
            result = [dict(name=name, code=code) for code, name in available_languages()]
            return web.json_response(result)
        except Exception as e:
            return web.json_response(dict(error=str(e)), status=500)

    @_server.routes.get("/api/etn/translate/{lang}/{text}")
    async def translate_text(request):
        try:
            language = request.match_info.get("lang", "en")
            text = request.match_info.get("text", "")
            result = translate(f"lang:{language} {text}")
            return web.json_response(result)
        except Exception as e:
            return web.json_response(dict(error=str(e)), status=500)

    @_server.routes.put("/api/etn/upload/{folder_name}/{filename}")
    async def upload(request: web.Request):
        folder_name = request.match_info.get("folder_name", "")
        if error := has_invalid_folder_name(folder_name):
            return error

        filename = request.match_info.get("filename", "")
        if error := has_invalid_filename(filename):
            return error

        try:
            if folder_paths.get_full_path(folder_name, filename) is not None:
                return web.json_response(dict(status="cached"), status=200)

            folder = Path(folder_paths.folder_names_and_paths[folder_name][0][0])
            total_size = int(request.headers.get("Content-Length", "0"))
            logging.info(f"Uploading {filename} ({total_size/(1024**2):.1f} MB) to {folder} folder")

            with open(folder / filename, "wb") as f:
                async for chunk, _ in request.content.iter_chunks():
                    f.write(chunk)

            return web.json_response(dict(status="success"), status=201)
        except Exception as e:
            return web.json_response(dict(error=str(e)), status=500)

    async def _handle_workflow_request(request: web.Request, handler, *arg_keys):
        try:
            data = await request.json()
            args = [data[key] for key in arg_keys]
            await handler(*args)
            return web.json_response(dict(status="success"), status=200)
        except KeyError as e:
            return web.json_response(dict(error=str(e)), status=400)
        except Exception as e:
            return web.json_response(dict(error=str(e)), status=500)

    @_server.routes.post("/api/etn/workflow/publish")
    async def publish_workflow(request: web.Request):
        return await _handle_workflow_request(
            request, _workflow_exchange.publish, "name", "client_id", "workflow"
        )

    @_server.routes.post("/api/etn/workflow/subscribe")
    async def subscribe_workflow(request: web.Request):
        return await _handle_workflow_request(request, _workflow_exchange.subscribe, "client_id")

    @_server.routes.post("/api/etn/workflow/unsubscribe")
    async def unsubscribe_workflow(request: web.Request):
        return await _handle_workflow_request(request, _workflow_exchange.unsubscribe, "client_id")