ComfyUI.js - ComfyUI Helper Class in JavaScript for ComfyUI API

ComfyUI.js

// ComfyUI.js
// Simple ComfyUI Helper Class in JavaScript for ComfyUI API

// Usage:
/*
const ComfyUI = require('./classes/ComfyUI');

// Define the nodes that you want to modify in your workflow
const nodes = {
  seeds: ['4', '12'],          // Node IDs that contain seeds
  checkpoint: ['2'],           // Node ID for model checkpoint
  prompts: ['3'],             // Node ID for text prompts
  base64_image_string: ['8'],  // Node ID for image input (if needed)
  api_save: ['15'],           // Node ID for saving output
  api_queue: ['15'],        // Same as api_save for auto-queueing
  //... You can add any node overrides you like! Browse your workflow to find them
};

// Callback for when images are saved
const onSaveCallback = async ({ message }) => {
  try {
    const files = message.data.output.images || [].concat(message.data.output.gifs || []);
    
    if (files.length > 0) {
      console.log('Files generated:', files);
      // Handle the generated files here
    }
  } catch (err) {
    console.error('Save callback error:', err);
  }
};

// Callback for when queue is ready for next item
const onQueueCallback = () => {
  console.log('Ready for next queue item');
  // Add logic to queue next item
};

// Initialize ComfyUI
const comfyui = new ComfyUI({
  comfyUIServerURL: 'https://your-comfyui-server.com',
  nodes,
  onSaveCallback,
  onQueueCallback
});

// Example workflow usage
async function generateImage() {
  try {
    // Your workflow JSON and API configuration
    const workflow = { }; // Your workflow JSON
    const workflowAPI = { }; // Your workflow API JSON

    // Queue an image generation
    await comfyui.queue({
      workflow,
      workflowAPI,
      prompt: "a beautiful landscape with mountains",
      checkpoint: "realisticVisionV40_v40VAE.safetensors"
    });

  } catch (err) {
    console.error('Generation error:', err);
  }
}

// Run the example
generateImage();
*/

const crypto = require('node:crypto');
const WebSocket = require('ws');

class ComfyUI {
  constructor({ comfyUIServerURL, nodes, onSaveCallback, onQueueCallback }) {
    // Note that connections will have to open and close for each request
    // This is due to the nature of the callbacks

    console.log('Starting ComfyUI.');

    // Init
    this.comfyUI = null;
    this.clientId = crypto.randomUUID();
    this.comfyUIServerURL = comfyUIServerURL;

    // Nodes allow you to change workflow params, see queue function
    this.nodes = nodes;

    // Callbacks
    if (onSaveCallback) {
      this.onSaveCallback = onSaveCallback;
    }

    if (onQueueCallback) {
      this.onQueueCallback = onQueueCallback;
    }

    // Init Queue Remaining
    this.queueRemaining = 0;

    this.connect();
  }

  connect() {
    console.log('Connecting to ComfyUI server...');
    const socketURL = `${this.comfyUIServerURL.replace('https://', 'wss://')}/ws?clientId=${this.clientId}`;

    this.comfyUI = new WebSocket(socketURL);

    // Connection Opened
    this.comfyUI.on('open', (data) => {
      console.log('ComfyUI server opened.');
    });

    // Connection Closed
    this.comfyUI.on('close', (data) => {
      console.log(`ComfyUI server closed: ${data}`);
    });

    // Message
    this.comfyUI.on('message', async (data) => {
      const message = JSON.parse(data);

      // if (message.type === 'status') {
      //   console.log(util.inspect(message, false, null, true));
      // }

      if (message.type === 'status') {
        this.queueRemaining = message.data.status.exec_info.queue_remaining;
      }

      if (message.data?.prompt_id) {
        // Matched File
        if (message.type === 'executed' && this.nodes.api_save.includes(message.data.node)) {
          // Save File
          if (this.onSaveCallback) {
            this.onSaveCallback({ message });
          }
        }

        // Queue Next (AutoQueue with api_queue)
        if (this.nodes['api_queue'] && message.type === 'executing'
          && this.nodes['api_queue'].includes(message.data.node)) {
          if (this.onQueueCallback) {
            this.onQueueCallback();
          }
        }

        // Done
        if (message.type === 'executing' && !message.data.node) {
          // End of inference
          if (this.queueRemaining === 0) {
            // Close websocket (Can't keep open due to callbacks getting crossed)
            this.comfyUI.close();
          } else {
            console.log('Queue Remaining:', this.queueRemaining);
          }
        }
      }
    });

    // Error
    this.comfyUI.on('error', (err) => {
      console.error(err);
      console.error(`Websocket Error with Client ${this.clientId}`);

      // Close Websocket
      this.disconnect();
    });
  }

  disconnect() {
    if (this.comfyUI) {
      this.comfyUI.close();
    }
  }

  queue({
    workflow,
    workflowAPI,
    prompt,
    base64ImageString,
    checkpoint = 'realisticVisionV40_v40VAE.safetensors',
  }) {
    return new Promise(async (resolve, reject) => {
      try {
        console.log(`Queuing ComfyUI on ${this.comfyUIServerURL}`);

        // TODO: Ping to see if serverURL is open

        const randomSeed = Math.floor(Math.random() * 18446744073709552000);
        const workflowApiTemp = { ...workflowAPI };

        for (const nodeType in this.nodes) {
          this.nodes[nodeType].forEach((nodeId) => {
            if (nodeType === 'seeds') {
              workflowApiTemp[nodeId].inputs.seed = randomSeed;
            } else if (nodeType === 'prompts') {
              workflowApiTemp[nodeId].inputs.text = prompt;
            } else if (nodeType === 'checkpoint') {
              workflowApiTemp[nodeId].inputs.ckpt_name = checkpoint;
            } else if (nodeType === 'base64_image_string') {
              workflowApiTemp[nodeId].inputs.image = base64ImageString;
            } 
            // Add more node modifiers
          });
        }

        const options = {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({
            prompt: workflowAPI,
            extra_data: {
              extra_pnginfo: {
                workflow,
              },
            },
            client_id: this.clientId,
          }),
        };

        const response = await fetch(`${this.comfyUIServerURL}/prompt`, options).then((res) => resolve(res.json()));
        resolve(response);
      } catch (err) {
        console.error(err);
        reject(err);
      }
    });
  }

  clearQueue() {
    return new Promise(async (resolve, reject) => {
      try {
        const options = {
          method: 'POST',
          headers: {
            'Content-Type': 'application/json',
          },
          body: JSON.stringify({
            clear: true,
          }),
        };

        await fetch(`${this.comfyUIServerURL}/queue`, options);
        resolve();

        console.log(`Queue cleared on server ${this.comfyUIServerURL}.`);
      } catch (err) {
        console.error(err);
        reject(err);
      }
    });
  }

  getFile({ filename, subfolder, type }) {
    return new Promise(async (resolve, reject) => {
      try {
        const data = {
          filename,
          subfolder,
          type,
        };

        const urlString = `${this.comfyUIServerURL}/view?${new url.URLSearchParams(data)}`;

        console.log(`Retrieving File ${urlString}.`);

        const response = await fetch(urlString);
        resolve(response.arrayBuffer());

      } catch (err) {
        console.error(err);
        reject(err);
      }
    });
  }

}

module.exports = ComfyUI;

// EOF