smartbotic/nodes/ai/comfyui-prompt.js

/**
 * @node comfyui-prompt
 * @name ComfyUI Prompt
 * @category ai
 * @version 1.0.0
 * @description Submit an API-format workflow graph to ComfyUI, wait for it to finish, and return the generated image
 * @icon image
 */

const configSchema = {
    type: 'object',
    properties: {
        baseUrl: {
            type: 'string',
            title: 'ComfyUI Base URL',
            description: 'Base URL of the ComfyUI server',
            default: 'http://127.0.0.1:8188'
        },
        workflowJson: {
            type: 'string',
            title: 'Workflow Graph',
            description: 'ComfyUI API-format graph, as produced by Save (API Format) in ComfyUI. Placeholders use double angle brackets around a dot path, and a seed placeholder yields a random seed. Double braces are deliberately not used, because the engine already interpolates those in node config.',
            format: 'json',
            openUrlFrom: 'baseUrl',
            importFrom: 'comfyui',
            default: ''
        },
        outputNodeId: {
            type: 'string',
            title: 'Output Node ID',
            description: 'Graph node whose images are returned, normally the SaveImage node',
            default: '9'
        },
        imageIndex: {
            type: 'number',
            title: 'Image Index',
            description: 'Which image to return when the output node produced several',
            default: 0
        },
        pollIntervalMs: {
            type: 'number',
            title: 'Poll Interval (ms)',
            description: 'How often to check whether the job has finished',
            default: 3000
        },
        timeoutMs: {
            type: 'number',
            title: 'Generation Timeout (ms)',
            description: 'Give up waiting after this long. A cold model load can take minutes.',
            default: 600000
        },
        returnImage: {
            type: 'boolean',
            title: 'Return Image Data',
            description: 'Download the result and return it base64-encoded. Disable to return only the filename.',
            default: true
        },
        retryCount: {
            type: 'number',
            title: 'Retries',
            description: 'Extra attempts after a failed submission or generation',
            default: 0
        },
        retryDelayMs: {
            type: 'number',
            title: 'Retry Delay (ms)',
            description: 'Pause before the first retry. Doubles on each further attempt.',
            default: 3000
        },
        retryMaxDelayMs: {
            type: 'number',
            title: 'Max Retry Delay (ms)',
            description: 'Upper bound for the exponential backoff',
            default: 30000
        },
        skipOnError: {
            type: 'boolean',
            title: 'Skip On Error',
            description: 'Return success false instead of failing the workflow. Useful inside a loop.',
            default: false
        }
    },
    required: ['baseUrl', 'workflowJson']
};

const inputSchema = {
    type: 'object',
    properties: {
        data: {
            type: 'any',
            description: 'Values available to placeholder substitution in the graph'
        }
    }
};

const outputSchema = {
    type: 'object',
    properties: {
        success: { type: 'boolean', description: 'False when generation failed and Skip On Error is enabled' },
        error: { type: 'string', description: 'Failure reason when success is false' },
        imageBase64: { type: 'string', description: 'Generated image, base64-encoded' },
        filename: { type: 'string', description: 'Filename ComfyUI assigned to the result' },
        subfolder: { type: 'string', description: 'ComfyUI output subfolder' },
        mimeType: { type: 'string', description: 'MIME type of the returned image' },
        size: { type: 'number', description: 'Image size in bytes' },
        promptId: { type: 'string', description: 'ComfyUI prompt id for the job' },
        elapsedMs: { type: 'number', description: 'Time from submission to completion' },
        attempts: { type: 'number', description: 'How many attempts were made' }
    }
};

function getPath(root, path) {
    if (!root || !path) {
        return undefined;
    }
    const parts = String(path).split('.');
    let current = root;
    for (let i = 0; i < parts.length; i++) {
        if (current === null || typeof current !== 'object') {
            return undefined;
        }
        current = current[parts[i]];
    }
    return current;
}

// The graph is a JSON document, so substituted values must be JSON-escaped or a
// prompt containing a quote or newline would produce an unparseable graph.
function jsonFragment(value) {
    if (value === undefined || value === null) {
        return '';
    }
    if (typeof value === 'number' || typeof value === 'boolean') {
        return String(value);
    }
    const encoded = JSON.stringify(String(value));
    return encoded.substring(1, encoded.length - 1);
}

// Placeholders are <<path>>, deliberately not {{path}}: the workflow engine
// interpolates {{...}} in node config before the node ever runs, which would
// blank these out. Substitution uses split/join rather than String.replace
// because this QuickJS build mishandles literal braces in regex patterns and
// silently drops the replacement.
function findPlaceholders(text) {
    const found = [];
    let rest = String(text);
    while (true) {
        const start = rest.indexOf('<<');
        if (start === -1) { break; }
        const end = rest.indexOf('>>', start + 2);
        if (end === -1) { break; }
        const key = rest.substring(start + 2, end).trim();
        if (key.length > 0 && found.indexOf(key) === -1) {
            found.push(key);
        }
        rest = rest.substring(end + 2);
    }
    return found;
}

// Inside a loop the engine hands a node the loop item rather than the upstream
// node's output, and the upstream result is nested somewhere under it. Resolve a
// path against the obvious roots first, then search the tree for it.
function deepResolve(root, path, depth) {
    if (!root || typeof root !== 'object' || depth > 6) {
        return undefined;
    }

    const direct = getPath(root, path);
    if (direct !== undefined) {
        return direct;
    }

    const keys = Object.keys(root);
    for (let i = 0; i < keys.length; i++) {
        const found = deepResolve(root[keys[i]], path, depth + 1);
        if (found !== undefined) {
            return found;
        }
    }
    return undefined;
}

function substitute(template, data) {
    let out = String(template);
    const keys = findPlaceholders(out);

    for (let i = 0; i < keys.length; i++) {
        const key = keys[i];
        let value = deepResolve(data, key, 0);

        if (value === undefined && key === 'seed') {
            value = Math.floor(Math.random() * 1000000000000000);
        }
        if (value === undefined) {
            continue;
        }

        out = out.split('<<' + key + '>>').join(jsonFragment(value));
    }

    return out;
}

function httpJson(method, url, body) {
    const options = { method: method, url: url, timeout: 60000 };
    if (body) {
        options.headers = { 'Content-Type': 'application/json' };
        options.body = body;
    }
    const response = smartbotic.http.request(options);
    if (response.status < 200 || response.status >= 300) {
        const detail = typeof response.data === 'string' ? response.data : JSON.stringify(response.data);
        throw new Error('ComfyUI HTTP ' + response.status + ': ' + String(detail).substring(0, 300));
    }
    return typeof response.data === 'string' ? JSON.parse(response.data) : response.data;
}

function runOnce(config, data, base) {
    const graphText = substitute(config.workflowJson, data);

    // An unresolved placeholder is the usual cause of a parse failure, and the raw
    // JSON error does not say which one, so name them explicitly.
    const unresolved = findPlaceholders(graphText);
    if (unresolved.length > 0) {
        const available = data && typeof data === 'object' ? Object.keys(data).join(', ') : '(none)';
        throw new Error('Unresolved placeholders in workflow graph: <<' + unresolved.join('>>, <<') +
                        '>>. Paths are relative to the previous node output, whose top-level keys are: ' + available);
    }

    let graph;
    try {
        graph = JSON.parse(graphText);
    } catch (e) {
        // Show where it broke: an empty substitution leaves "key": , which is the
        // usual cause and is impossible to spot from the parser message alone.
        let where = graphText.indexOf(': ,');
        if (where === -1) { where = graphText.indexOf(':,'); }
        if (where === -1) { where = 0; }
        const snippet = graphText.substring(Math.max(0, where - 80), where + 80);
        throw new Error('Workflow graph is not valid JSON after substitution: ' + e.message +
                        ' | near: ' + snippet);
    }

    const payload = graph.prompt ? graph : { prompt: graph };
    payload.client_id = 'smartbotic-' + (config.clientSuffix || 'node');

    const submitted = httpJson('POST', base + '/prompt', JSON.stringify(payload));
    if (!submitted || !submitted.prompt_id) {
        throw new Error('ComfyUI did not return a prompt_id');
    }
    if (submitted.node_errors && Object.keys(submitted.node_errors).length > 0) {
        throw new Error('ComfyUI rejected the graph: ' + JSON.stringify(submitted.node_errors).substring(0, 300));
    }

    const promptId = submitted.prompt_id;
    const interval = Number(config.pollIntervalMs) || 3000;
    const timeout = Number(config.timeoutMs) || 600000;
    const deadline = Date.now() + timeout;
    const startedAt = Date.now();

    let history = null;
    while (Date.now() < deadline) {
        smartbotic.utils.sleep(interval);
        const entry = httpJson('GET', base + '/history/' + promptId, null);
        if (entry && entry[promptId]) {
            history = entry[promptId];
            break;
        }
    }

    if (!history) {
        throw new Error('ComfyUI generation timed out after ' + timeout + 'ms (prompt ' + promptId + ')');
    }

    if (history.status && history.status.status_str && history.status.status_str !== 'success') {
        throw new Error('ComfyUI generation failed: ' + history.status.status_str);
    }

    const outputs = history.outputs || {};
    const nodeId = String(config.outputNodeId || '9');
    const nodeOutput = outputs[nodeId];
    if (!nodeOutput || !nodeOutput.images || nodeOutput.images.length === 0) {
        throw new Error('Output node ' + nodeId + ' produced no images. Available: ' + Object.keys(outputs).join(', '));
    }

    const index = Number(config.imageIndex) || 0;
    const image = nodeOutput.images[index] || nodeOutput.images[0];

    const result = {
        promptId: promptId,
        filename: image.filename || '',
        subfolder: image.subfolder || '',
        elapsedMs: Date.now() - startedAt,
        imageBase64: '',
        mimeType: 'image/png',
        size: 0
    };

    if (config.returnImage === false) {
        return result;
    }

    const viewUrl = base + '/view?filename=' + encodeURIComponent(result.filename) +
                    '&subfolder=' + encodeURIComponent(result.subfolder) + '&type=output';
    const download = smartbotic.http.request({ method: 'GET', url: viewUrl, timeout: 120000 });
    if (download.status < 200 || download.status >= 300) {
        throw new Error('Failed to download generated image: HTTP ' + download.status);
    }

    // dataBase64 carries the raw bytes; response.data is UTF-8 validated and
    // would destroy the PNG.
    if (typeof download.dataBase64 !== 'string' || download.dataBase64.length === 0) {
        throw new Error('Runtime returned no binary data for the generated image');
    }

    result.imageBase64 = download.dataBase64;
    result.size = Math.floor(result.imageBase64.replace(/=+$/, '').length * 3 / 4);

    const headers = download.headers || {};
    if (headers['content-type']) {
        result.mimeType = String(headers['content-type']).split(';')[0].trim();
    }

    return result;
}

module.exports = {
    configSchema,
    inputSchema,
    outputSchema,

    async execute(config, input, context) {
        if (!config.workflowJson || String(config.workflowJson).trim().length === 0) {
            throw new Error('A ComfyUI workflow graph is required');
        }

        // Resolve against the whole input: the upstream result may be nested under
        // the loop item rather than sitting at input.data.
        const data = input;
        const base = String(config.baseUrl).replace(/\/+$/, '');
        const attempts = 1 + (Number(config.retryCount) > 0 ? Number(config.retryCount) : 0);

        let lastError = '';
        let used = 0;

        for (let attempt = 1; attempt <= attempts; attempt++) {
            used = attempt;
            try {
                const result = runOnce(config, data, base);
                return {
                    success: true,
                    error: '',
                    imageBase64: result.imageBase64,
                    filename: result.filename,
                    subfolder: result.subfolder,
                    mimeType: result.mimeType,
                    size: result.size,
                    promptId: result.promptId,
                    elapsedMs: result.elapsedMs,
                    attempts: used
                };
            } catch (err) {
                lastError = err && err.message ? err.message : String(err);
                smartbotic.log.warn('comfyui-prompt: attempt ' + attempt + ' of ' + attempts + ' failed: ' + lastError);
                if (attempt < attempts) {
                    const baseDelay = Number(config.retryDelayMs) || 3000;
                    const cap = Number(config.retryMaxDelayMs) || 30000;
                    let delay = baseDelay * Math.pow(2, attempt - 1);
                    if (delay > cap) { delay = cap; }
                    smartbotic.utils.sleep(delay);
                }
            }
        }

        if (config.skipOnError !== true) {
            throw new Error('comfyui-prompt failed after ' + used + ' attempt(s): ' + lastError);
        }

        smartbotic.log.warn('comfyui-prompt: skipping after ' + used + ' attempt(s)');
        return {
            success: false,
            error: lastError,
            imageBase64: '',
            filename: '',
            subfolder: '',
            mimeType: '',
            size: 0,
            promptId: '',
            elapsedMs: 0,
            attempts: used
        };
    }
};