api.py

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

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

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

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

model_names = {
    "SD15": "sd15",
    "SD20": "sd20",
    "SD21UnclipL": "sd21",
    "SD21UnclipH": "sd21",
    "SDXLRefiner": "sdxl-refiner",
    "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",
    "FluxInpaint": "flux",
    "FluxSchnell": "flux-schnell",
    "GenmoMochi": "mochi",
    "LTXV": "ltxv",
    "HunyuanVideo": "hunyuan-video",
    "CosmosT2V": "cosmos",
    "CosmosI2V": "cosmos",
    "CosmosT2IPredict2": "cosmos-predict2",
    "CosmosI2VPredict2": "cosmos-predict2",
    "ZImage": "z-image",
    "Lumina2": "lumina2",
    "WAN21_T2V": "wan21",
    "WAN21_I2V": "wan21",
    "WAN21_FunControl2V": "wan21-fun",
    "WAN21_Vace": "wan21-vace",
    "WAN21_Camera": "wan21-camera",
    "HiDream": "hi-dream",
    "Chroma": "chroma",
    "ACEStep": "ace-step",
    "Omnigen2": "omnigen2",
    "QwenImage": "qwen-image",
    "Flux2": "flux2",
}

gguf_architectures = {
    "sd1": "sd15",
    "qwen_image": "qwen-image",
}


class FakeTensor(NamedTuple):
    shape: tuple

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


def inspect_safetensors(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:
                cfg = comfy.utils.state_dict_prefix_replace(cfg, {prefix: ""}, filter_keys=False)
                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 = None
            model_type = None
            model_quant = None

            # Check if it's a Nunchaku SVDQ model by inspecting metadata
            raw_name = detect_svdq(cfg)
            if raw_name:
                model_quant = "svdq"
            # Otherwise try ComfyUI's model detection
            elif unet_config is not None:
                base_model = model_detection.model_config_from_unet_config(unet_config)
                if base_model:
                    raw_name = base_model.__class__.__name__
                    if raw_name == "SDXL":
                        model_type = base_model.model_type(cfg).name.lower().replace("_", "-")
                    if raw_name == "Flux2":
                        hidden_size = unet_config.get("hidden_size", 0)
                        model_type = {3072: "klein-4b", 4096: "klein-9b"}.get(hidden_size, "dev")

            if not raw_name:
                return {"base_model": "unknown"}

            base_model_name = model_names.get(raw_name, "unknown")
            result: dict[str, Any] = {"base_model": base_model_name}
            result["is_inpaint"] = (
                base_model_name in ["sd15", "sdxl"] and input_count > 4
            ) or raw_name == "FluxInpaint"
            if model_quant:
                result["quant"] = model_quant
            if model_type:
                result["type"] = model_type
            elif "T2I" in raw_name:
                result["type"] = "t2i"
            elif "I2V" in raw_name:
                result["type"] = "i2v"
            elif "T2V" in raw_name:
                result["type"] = "t2v"
            elif "Control2V" in raw_name:
                result["type"] = "control2v"
            return result
        return {"base_model": "unknown"}
    except Exception as e:
        print("[comfyui-tooling-nodes] Error inspecting file", filename)
        traceback.print_exc()
        return {"base_model": "unknown", "error": f"Failed to detect base model: {e}"}


def detect_svdq(cfg: dict) -> str | None:
    if md := cfg.get("__metadata__"):
        if comfy_config := md.get("comfy_config"):
            if isinstance(comfy_config, str):
                comfy_config = json.loads(comfy_config)
            if model_class := comfy_config.get("model_class"):
                return model_class

        match md.get("model_class"):
            case "NunchakuFluxTransformer2dModel":
                return "Flux"
            case "NunchakuQwenImageTransformer2DModel":
                return "QwenImage"
            case "NunchakuZImageTransformer2DModel":
                return "ZImage"
    return None


def inspect_gguf(filename: str, model_type: str):
    try:
        import gguf
    except ImportError:
        return {"base_model": "unknown", "error": "GGUF module not found"}

    try:
        path = folder_paths.get_full_path(model_type, filename)
        if path is None:
            raise Exception(f"Could not find full path for {model_type}/{filename}")

        reader = gguf.GGUFReader(path)
        arch_field = reader.get_field("general.architecture")
        if arch_field is not None:
            if len(arch_field.types) != 1 or arch_field.types[0] != gguf.GGUFValueType.STRING:
                raise TypeError(
                    f"Bad type for GGUF general.architecture key: expected string, got {arch_field.types!r}"
                )
            arch_str = str(arch_field.parts[arch_field.data[-1]], encoding="utf-8")
        else:  # stable-diffusion.cpp, requires conversion. not handled for now
            return {"base_model": "flux", "is_inpaint": False}

        if arch_str == "flux" and any(
            t.name.startswith("distilled_guidance_layer")
            for t in itertools.islice(reader.tensors, 5)
        ):
            arch_str = "chroma"

        # Detect Z-Image (modified Lumina2)
        if arch_str == "lumina2":
            for t in reader.tensors:
                if t.name == "cap_embedder.1.bias" and t.shape[0] == 3840:
                    arch_str = "z-image"
                    break

        # Detect Flux variants
        result_type = None
        if arch_str == "flux":
            for t in reader.tensors:
                if t.name.startswith("distilled_guidance_layer"):
                    arch_str = "chroma"
                    break
                elif t.name == "double_stream_modulation_img.lin.weight":
                    arch_str = "flux2"
                    if t.shape[0] == 3072:
                        result_type = "klein-4b"
                    elif t.shape[0] == 4096:
                        result_type = "klein-9b"
                    break

        result = {
            "base_model": gguf_architectures.get(arch_str, arch_str),
            "is_inpaint": False,
        }
        if result_type is not None:
            result["type"] = result_type
        try:
            if file_type := reader.get_field("general.file_type"):
                result["quant"] = file_type.contents().lower()
        except Exception:
            result["quant"] = "gguf"
        return result

    except Exception as e:
        # traceback.print_exc()
        return {"base_model": "unknown", "error": f"Failed to detect base model: {e}"}


def inspect_diffusion_model(filename: str, model_type: str, is_checkpoint: bool):
    if filename.endswith(".gguf"):
        return inspect_gguf(filename, model_type)
    return inspect_safetensors(filename, model_type, is_checkpoint)


def inspect_models(model_type: str, params: dict[str, str]):
    try:
        try:
            files = folder_paths.get_filename_list(model_type)
        except KeyError:
            return web.json_response({"error": f"Model folder not found: {model_type}"})
        limit = int(params.get("limit", "1000"))
        offset = int(params.get("offset", "0"))
        files_range = files[offset : offset + limit]
        is_checkpoint = model_type == "checkpoints"
        info = {
            filename: inspect_diffusion_model(filename, model_type, is_checkpoint)
            for filename in files_range
        }
        if "limit" in params:
            info["_meta"] = dict(offset=offset, count=len(files_range), total=len(files))
        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


async def image_sender(data: bytes):
    mem = memoryview(data)
    csize = 2**14
    for i in range(0, len(mem), csize):
        yield mem[i : i + csize]


_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")
        error = has_invalid_folder_name(folder_name)
        if error is not None:
            return error
        return inspect_models(folder_name, request.rel_url.query)

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

    @_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.get("/api/etn/image/{id}")
    async def get_image(request: web.Request):
        try:
            id = request.match_info.get("id", "")
            data, content_type = image_cache.get(id)
            if data is None or content_type is None:
                return web.json_response(dict(error="Image not found"), status=404)
            response = web.Response(
                body=image_sender(data),
                content_type=content_type,
                headers={"Content-Length": str(len(data))},
            )
            return response
        except Exception as e:
            return web.json_response(dict(error=str(e)), status=500)

    @_server.routes.put("/api/etn/image/{id}")
    async def put_image(request: web.Request):
        try:
            id = request.match_info.get("id", "")
            if id in image_cache:
                return web.json_response(dict(status="cached"), status=200)

            content_type = request.headers.get("Content-Type", "application/octet-stream")
            data = bytearray()
            async for chunk, _ in request.content.iter_chunks():
                data.extend(chunk)

            image_cache.insert(id, bytes(data), content_type)
            return web.json_response(dict(status="success"), status=201)
        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", "")
        error = has_invalid_folder_name(folder_name)
        if error is not None:
            return error

        filename = request.match_info.get("filename", "")
        error = has_invalid_filename(filename)
        if error is not None:
            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")