comfy-api-addons/__init__.py

from server import PromptServer
from aiohttp import web, ClientSession
import os
import inspect
import json

"""ComfyUI custom node: Mobile UI Workflow Endpoints

This custom node registers HTTP routes that allow a standalone client to:
- List workflow JSON files on the ComfyUI machine
- Fetch the raw workflow JSON
- Fetch an API-compatible prompt JSON (if already API prompt OR if a converter endpoint is available)
- Convert a posted workflow JSON payload to an API prompt JSON

Routes:
- GET     /mobileui/info
- GET     /mobileui/workflows
- GET     /mobileui/workflows/{name:.+}
- GET     /mobileui/workflows_api/{name:.+}
- POST    /mobileui/convert
- OPTIONS /mobileui/{tail:.*}

Config:
- COMFY_WORKFLOWS_DIR   : single directory
- COMFY_WORKFLOWS_DIRS  : comma-separated list of directories
- MOBILEUI_CORS_ORIGIN  : if set, adds Access-Control-Allow-Origin for these routes

Notes:
- This node only serves files ending in .json.
- Includes path traversal protection.
"""

NODE_VERSION = "0.3.0"


def _comfy_root_dir() -> str:
  # server.py lives under the ComfyUI root in typical installs
  server_py = inspect.getfile(PromptServer)
  return os.path.dirname(os.path.abspath(server_py))


def _resolve_workflow_dirs() -> list[str]:
  dirs: list[str] = []

  env_single = os.environ.get("COMFY_WORKFLOWS_DIR")
  env_multi = os.environ.get("COMFY_WORKFLOWS_DIRS")

  if env_multi:
    dirs.extend([d.strip() for d in env_multi.split(",") if d.strip()])
  if env_single:
    dirs.append(env_single)

  root = _comfy_root_dir()
  candidates = [
    os.path.join(root, "workflows"),
    os.path.join(root, "user", "workflows"),
    os.path.join(root, "examples"),
    os.path.join(root, "examples", "workflows"),
    # popular location used by ComfyUI-Custom-Scripts
    os.path.join(root, "pysssss-workflows"),
  ]
  dirs.extend(candidates)

  normalized: list[str] = []
  seen: set[str] = set()
  for d in dirs:
    if not d:
      continue
    absd = os.path.abspath(d)
    if os.path.isdir(absd) and absd not in seen:
      seen.add(absd)
      normalized.append(absd)

  return normalized


_BASES = _resolve_workflow_dirs()


def _cors_origin() -> str | None:
  # If you start ComfyUI with --enable-cors-header this may be redundant,
  # but having per-route CORS makes the node easier to drop-in.
  return os.environ.get("MOBILEUI_CORS_ORIGIN")


def _add_cors_headers(resp: web.StreamResponse):
  origin = _cors_origin()
  if not origin:
    return resp

  resp.headers["Access-Control-Allow-Origin"] = origin
  resp.headers["Vary"] = "Origin"
  resp.headers["Access-Control-Allow-Methods"] = "GET,POST,OPTIONS"
  resp.headers["Access-Control-Allow-Headers"] = "Content-Type"
  resp.headers["Access-Control-Allow-Credentials"] = "true"
  return resp


def _list_workflows_with_meta() -> list[dict]:
  items: list[dict] = []
  seen: set[str] = set()

  for base in _BASES:
    for dirpath, _dirnames, filenames in os.walk(base):
      for fname in filenames:
        if not fname.lower().endswith(".json"):
          continue
        full = os.path.join(dirpath, fname)
        rel = os.path.relpath(full, base)
        rel_no_ext = os.path.splitext(rel)[0]
        rel_norm = rel_no_ext.replace(os.sep, "/")
        if rel_norm in seen:
          continue

        seen.add(rel_norm)

        try:
          st = os.stat(full)
          items.append(
            {
              "name": rel_norm,
              "bytes": int(st.st_size),
              "mtime": int(st.st_mtime),
            }
          )
        except Exception:
          items.append({"name": rel_norm, "bytes": None, "mtime": None})

  items.sort(key=lambda x: x.get("name") or "")
  return items


def _safe_join_and_validate(base: str, rel: str) -> str | None:
  # Prevent path traversal attacks.
  full = os.path.abspath(os.path.join(base, rel))
  base_abs = os.path.abspath(base)
  try:
    if os.path.commonpath([full, base_abs]) != base_abs:
      return None
  except Exception:
    return None
  return full


def _find_workflow_file(name: str) -> str | None:
  filename = name if name.lower().endswith(".json") else f"{name}.json"
  for base in _BASES:
    full = _safe_join_and_validate(base, filename)
    if not full:
      continue
    if os.path.isfile(full):
      return full
  return None


def _looks_like_api_prompt(payload) -> bool:
  # Heuristic: a ComfyUI API prompt is usually an object keyed by node-id strings,
  # with values containing at least class_type.
  if not isinstance(payload, dict) or not payload:
    return False
  sample_key = next(iter(payload.keys()))
  if not isinstance(sample_key, str):
    return False
  sample_val = payload.get(sample_key)
  return isinstance(sample_val, dict) and ("class_type" in sample_val)


def _converter_url(request) -> str:
  # Use the same host/scheme the client used.
  return f"{request.scheme}://{request.host}/workflow/convert"


async def _convert_via_converter_endpoint(request, workflow_json):
  url = _converter_url(request)
  async with ClientSession() as session:
    async with session.post(url, json=workflow_json) as r:
      if r.status != 200:
        text = await r.text()
        raise RuntimeError(f"conversion failed ({r.status}): {text}")
      return await r.json()


@PromptServer.instance.routes.get("/mobileui/info")
async def _mobileui_info(_request):
  payload = {
    "name": "comfyui_mobileui_workflow_endpoints",
    "version": NODE_VERSION,
    "workflow_dirs": _BASES,
    "endpoints": {
      "workflows": "/mobileui/workflows",
      "workflow": "/mobileui/workflows/{name}",
      "workflow_api": "/mobileui/workflows_api/{name}",
      "convert": "/mobileui/convert",
    },
  }
  resp = web.json_response(payload)
  return _add_cors_headers(resp)


@PromptServer.instance.routes.get("/mobileui/workflows")
async def _mobileui_list_workflows(_request):
  # Returns [{ name, bytes, mtime }, ...]
  items = _list_workflows_with_meta()
  resp = web.json_response(items)
  return _add_cors_headers(resp)


@PromptServer.instance.routes.get("/mobileui/workflows/")
async def _mobileui_list_workflows_slash(request):
  return await _mobileui_list_workflows(request)


@PromptServer.instance.routes.get("/mobileui/workflows/{name:.+}")
async def _mobileui_get_workflow(request):
  name = request.match_info["name"].lstrip("/\\")
  path = _find_workflow_file(name)
  if not path:
    resp = web.Response(status=404, text="workflow not found")
    return _add_cors_headers(resp)

  resp = web.FileResponse(path)
  return _add_cors_headers(resp)


@PromptServer.instance.routes.get("/mobileui/workflows_api/{name:.+}")
async def _mobileui_get_workflow_api(request):
  name = request.match_info["name"].lstrip("/\\")
  path = _find_workflow_file(name)
  if not path:
    resp = web.Response(status=404, text="workflow not found")
    return _add_cors_headers(resp)

  # Load JSON
  try:
    with open(path, "r", encoding="utf-8") as f:
      data = json.load(f)
  except Exception:
    resp = web.Response(status=400, text="invalid json")
    return _add_cors_headers(resp)

  # Already API prompt
  if _looks_like_api_prompt(data):
    resp = web.json_response(data)
    return _add_cors_headers(resp)

  # Attempt conversion via SethRobinson's /workflow/convert endpoint if installed.
  try:
    converted = await _convert_via_converter_endpoint(request, data)
    resp = web.json_response(converted)
    return _add_cors_headers(resp)
  except Exception as e:
    resp = web.Response(
      status=501,
      text=(
        "converter endpoint not available or conversion failed. "
        "Install comfyui-workflow-to-api-converter-endpoint to enable /workflow/convert. "
        f"Details: {e}"
      ),
    )
    return _add_cors_headers(resp)


@PromptServer.instance.routes.post("/mobileui/convert")
async def _mobileui_convert(request):
  # Accepts posted JSON body and returns an API prompt JSON.
  try:
    data = await request.json()
  except Exception:
    resp = web.Response(status=400, text="invalid json")
    return _add_cors_headers(resp)

  if _looks_like_api_prompt(data):
    resp = web.json_response(data)
    return _add_cors_headers(resp)

  try:
    converted = await _convert_via_converter_endpoint(request, data)
    resp = web.json_response(converted)
    return _add_cors_headers(resp)
  except Exception as e:
    resp = web.Response(
      status=501,
      text=(
        "converter endpoint not available or conversion failed. "
        "Install comfyui-workflow-to-api-converter-endpoint to enable /workflow/convert. "
        f"Details: {e}"
      ),
    )
    return _add_cors_headers(resp)


@PromptServer.instance.routes.options("/mobileui/{tail:.*}")
async def _mobileui_options(_request):
  resp = web.Response(status=204)
  return _add_cors_headers(resp)