comfy-discord-web/web/routers/workflow_router.py

"""
GET /api/workflow — current workflow info
GET /api/workflow/inputs — dynamic NodeInput list
GET /api/workflow/files — list files in workflows/
POST /api/workflow/upload — upload a workflow JSON
POST /api/workflow/load — load a workflow from workflows/
GET /api/workflow/models?type=checkpoints|loras — available models (60s TTL cache)
"""
from __future__ import annotations

import json
import logging
import time
from pathlib import Path
from typing import Optional

from fastapi import APIRouter, Depends, File, Form, HTTPException, UploadFile

from web.auth import require_auth
from web.deps import get_comfy, get_config, get_inspector, get_user_registry

router = APIRouter()
logger = logging.getLogger(__name__)

_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
_WORKFLOWS_DIR = _PROJECT_ROOT / "workflows"

# Simple in-memory TTL cache for models
_models_cache: dict = {}
_MODELS_TTL = 60.0


@router.get("")
async def get_workflow(user: dict = Depends(require_auth)):
    """Return basic info about the currently loaded workflow."""
    user_label: str = user["sub"]
    registry = get_user_registry()

    if registry:
        template = registry.get_workflow_template(user_label)
        last_wf = registry.get_state_manager(user_label).get_last_workflow_file()
    else:
        # Fallback to global state when registry is unavailable
        comfy = get_comfy()
        if comfy is None:
            raise HTTPException(503, "Workflow manager not available")
        template = comfy.workflow_manager.get_workflow_template()
        last_wf = comfy.state_manager.get_last_workflow_file()

    return {
        "loaded": template is not None,
        "node_count": len(template) if template else 0,
        "last_workflow_file": last_wf,
    }


@router.get("/inputs")
async def get_workflow_inputs(user: dict = Depends(require_auth)):
    """Return dynamic NodeInput list (common + advanced) for the current workflow."""
    user_label: str = user["sub"]
    inspector = get_inspector()
    if inspector is None:
        raise HTTPException(503, "Workflow components not available")

    registry = get_user_registry()
    if registry:
        template = registry.get_workflow_template(user_label)
        overrides = registry.get_state_manager(user_label).get_overrides()
    else:
        comfy = get_comfy()
        if comfy is None:
            raise HTTPException(503, "Workflow components not available")
        template = comfy.workflow_manager.get_workflow_template()
        overrides = comfy.state_manager.get_overrides()

    if template is None:
        return {"common": [], "advanced": []}

    inputs = inspector.inspect(template)
    result = []
    for ni in inputs:
        val = overrides.get(ni.key, ni.current_value)
        result.append({
            "key": ni.key,
            "label": ni.label,
            "input_type": ni.input_type,
            "current_value": val,
            "node_class": ni.node_class,
            "node_title": ni.node_title,
            "is_common": ni.is_common,
        })
    common = [r for r in result if r["is_common"]]
    advanced = [r for r in result if not r["is_common"]]
    return {"common": common, "advanced": advanced}


@router.get("/files")
async def list_workflow_files(_: dict = Depends(require_auth)):
    """List .json files in the workflows/ folder."""
    _WORKFLOWS_DIR.mkdir(parents=True, exist_ok=True)
    files = sorted(p.name for p in _WORKFLOWS_DIR.glob("*.json"))
    return {"files": files}


@router.post("/upload")
async def upload_workflow(
    file: UploadFile = File(...),
    _: dict = Depends(require_auth),
):
    """Upload a workflow JSON to the workflows/ folder."""
    _WORKFLOWS_DIR.mkdir(parents=True, exist_ok=True)
    filename = file.filename or "workflow.json"
    if not filename.endswith(".json"):
        filename += ".json"
    data = await file.read()
    try:
        json.loads(data)  # validate JSON
    except json.JSONDecodeError as exc:
        raise HTTPException(400, f"Invalid JSON: {exc}")
    dest = _WORKFLOWS_DIR / filename
    dest.write_bytes(data)
    return {"ok": True, "filename": filename}


@router.post("/load")
async def load_workflow(filename: str = Form(...), user: dict = Depends(require_auth)):
    """Load a workflow from the workflows/ folder into the user's isolated state."""
    wf_path = _WORKFLOWS_DIR / filename
    if not wf_path.exists():
        raise HTTPException(404, f"Workflow file '{filename}' not found")
    try:
        with open(wf_path, "r", encoding="utf-8") as f:
            workflow = json.load(f)
    except Exception as exc:
        raise HTTPException(500, str(exc))

    registry = get_user_registry()
    if registry:
        registry.set_workflow(user["sub"], workflow, filename)
    else:
        # Fallback: update global state when registry unavailable
        comfy = get_comfy()
        if comfy is None:
            raise HTTPException(503, "ComfyUI not available")
        comfy.workflow_manager.set_workflow_template(workflow)
        comfy.state_manager.clear_overrides()
        comfy.state_manager.set_last_workflow_file(filename)

    inspector = get_inspector()
    node_count = len(workflow)
    inputs_count = len(inspector.inspect(workflow)) if inspector else 0
    return {
        "ok": True,
        "filename": filename,
        "node_count": node_count,
        "inputs_count": inputs_count,
    }


@router.get("/models")
async def get_models(type: str = "checkpoints", _: dict = Depends(require_auth)):
    """Return available model names from ComfyUI (60s TTL cache)."""
    global _models_cache
    now = time.time()
    cache_key = type
    cached = _models_cache.get(cache_key)
    if cached and (now - cached["ts"]) < _MODELS_TTL:
        return {"type": type, "models": cached["models"]}

    comfy = get_comfy()
    if comfy is None:
        raise HTTPException(503, "ComfyUI not available")
    models = await comfy.get_models(type)
    _models_cache[cache_key] = {"models": models, "ts": now}
    return {"type": type, "models": models}