@ldavis9000aws/swarmui-generator
Version:
A Model Context Protocol server for SwarmUI image generation with TypeScript
261 lines • 10.4 kB
JavaScript
/**
* SwarmUI API Client
*
* Handles communication with the SwarmUI API, including session management,
* image generation, model listing, and other API operations.
*/
import axios from 'axios';
export class SwarmUIClient {
baseUrl;
client;
sessionId = null;
sessionRefreshTimer = null;
sessionRefreshInterval = 30 * 60 * 1000; // 30 minutes
/**
* Create a new SwarmUI API client
* @param baseUrl Base URL of the SwarmUI API server
*/
constructor(baseUrl) {
this.baseUrl = baseUrl.endsWith('/') ? baseUrl.slice(0, -1) : baseUrl;
this.client = axios.create({
baseURL: this.baseUrl,
timeout: 600000, // 10 minutes for image generation
headers: {
'Content-Type': 'application/json',
'Accept': 'application/json',
'User-Agent': 'SwarmUI-MCP-Server/1.0'
}
});
console.log(`SwarmUI client initialized with API URL: ${this.baseUrl}`);
}
/**
* Get a new session ID from the SwarmUI API
* @returns Session ID string
*/
async getNewSession() {
try {
console.log('Requesting new SwarmUI session...');
const response = await this.client.post('/API/GetNewSession', {});
const sessionId = response.data.session_id;
if (!sessionId) {
throw new Error('Failed to get session ID from SwarmUI API');
}
console.log(`New session established: ${sessionId}`);
if (this.sessionRefreshTimer) {
clearTimeout(this.sessionRefreshTimer);
}
this.sessionRefreshTimer = setTimeout(() => {
console.log('Session timeout reached, will request new session on next API call');
this.sessionId = null;
}, this.sessionRefreshInterval);
return sessionId;
}
catch (error) {
console.error('Session creation failed:', error);
throw new Error(`Failed to get session ID: ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Ensure we have a valid session ID
* @returns Valid session ID
*/
async ensureSession() {
if (!this.sessionId) {
this.sessionId = await this.getNewSession();
}
return this.sessionId;
}
/**
* Call the SwarmUI API
* @param endpoint API endpoint name (without the /API/ prefix)
* @param data Request data
* @param config Additional Axios request configuration
* @returns API response
*/
async callApi(endpoint, data = {}, config = {}) {
try {
const sessionId = await this.ensureSession();
const requestData = { ...data, session_id: sessionId };
console.log(`Calling SwarmUI API: ${endpoint} with data:`, JSON.stringify(requestData));
const response = await this.client.post(`/API/${endpoint}`, requestData, config);
if (typeof response.data === 'object' && response.data !== null) {
const responseObj = response.data;
if (responseObj.error_id === 'invalid_session_id') {
console.log('Session invalid, requesting new session and retrying');
this.sessionId = null;
return this.callApi(endpoint, data, config);
}
}
return response.data;
}
catch (error) {
if (axios.isAxiosError(error) && error.response?.data) {
const responseData = error.response.data;
console.error('SwarmUI API Error Response Data:', responseData);
if (typeof responseData === 'object' && responseData !== null) {
const responseObj = responseData;
if (responseObj.error_id === 'invalid_session_id') {
console.log('Session invalid (in error path), requesting new session and retrying');
this.sessionId = null;
return this.callApi(endpoint, data, config);
}
if (responseObj.error) {
throw new Error(`SwarmUI API error: ${responseObj.error}`);
}
}
}
throw new Error(`Error calling SwarmUI API: ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Download an image from the SwarmUI API
* @param imagePath Image path returned by the API
* @returns Image data as a Buffer
*/
async downloadImage(imagePath) {
try {
const normalizedPath = imagePath.startsWith('/') ? imagePath.substring(1) : imagePath;
console.log(`Downloading image from: ${this.baseUrl}/${normalizedPath}`);
const response = await axios.get(`${this.baseUrl}/${normalizedPath}`, {
responseType: 'arraybuffer'
});
return Buffer.from(response.data);
}
catch (error) {
console.error('Image download failed:', error);
throw new Error(`Failed to download image: ${error instanceof Error ? error.message : String(error)}`);
}
}
/**
* Download and encode an image to base64
* @param imagePath Image path returned by the API
* @returns Base64 encoded image string (without data URI prefix)
*/
async downloadAndEncodeImage(imagePath) {
const buffer = await this.downloadImage(imagePath);
return buffer.toString('base64');
}
/**
* List available models and folders
* @param options Options for listing models including path, depth, subtype, sortBy, sortReverse, allowRemote.
* @returns Object containing lists of available models (files) and folders
*/
async listModels(options = {}) {
const apiData = {
path: options.path === undefined ? "" : options.path,
depth: options.depth === undefined ? 1 : options.depth,
};
if (options.subtype !== undefined)
apiData.subtype = options.subtype;
if (options.sortBy !== undefined)
apiData.sortBy = options.sortBy;
if (options.sortReverse !== undefined)
apiData.sortReverse = options.sortReverse;
if (options.allowRemote !== undefined)
apiData.allowRemote = options.allowRemote; // Added allowRemote
console.log("Calling ListModels with data:", apiData);
// Ensure the type here matches the expected structure from the API, including 'folders'
const response = await this.callApi('ListModels', apiData);
return {
files: response.files || [],
folders: response.folders || []
};
}
/**
* List available schedulers
* @returns List of available schedulers
*/
async listSchedulers() {
const response = await this.callApi('ListSchedulers');
return response.samplers || [];
}
/**
* Generate an image from a text prompt
* @param prompt Text prompt
* @param options Generation options
* @returns Image generation response with paths to generated images
*/
async generateImage(prompt, options = {}) {
console.log(`Generating image for prompt: "${prompt}" with options:`, options);
const requestData = {
prompt,
negativeprompt: options.negativePrompt || '',
width: options.width,
height: options.height,
cfgscale: options.cfgScale,
steps: options.steps,
seed: options.seed,
model: options.model,
sampler: options.scheduler,
images: options.batchSize || 1,
donotsave: false,
};
Object.keys(requestData).forEach(key => requestData[key] === undefined && delete requestData[key]);
if (requestData.seed === undefined)
requestData.seed = -1;
if (requestData.images === undefined)
requestData.images = 1;
const response = await this.callApi('GenerateText2Image', requestData);
console.log('SwarmUI GenerateText2Image Raw Response:', response);
if (!response || response.error || (response.success !== undefined && !response.success)) {
let errorMessage = 'Unknown error';
if (response && response.error)
errorMessage = String(response.error);
else if (response && response.error_id)
errorMessage = `Error ID: ${response.error_id}`;
throw new Error(`Image generation failed: ${errorMessage}`);
}
let imagePaths = [];
if (response.images) {
if (Array.isArray(response.images)) {
imagePaths = response.images;
}
else if (typeof response.images === 'object' && response.images !== null) {
imagePaths = Object.values(response.images).map((imgDetails) => {
if (Array.isArray(imgDetails) && typeof imgDetails[0] === 'string')
return imgDetails[0];
if (typeof imgDetails === 'string')
return imgDetails;
console.warn('Unexpected image details format in response.images object:', imgDetails);
return null;
}).filter(path => path !== null);
}
}
return {
images: imagePaths,
success: true
};
}
/**
* Get the status of the text-to-image engine
* @returns Engine status
*/
async getEngineStatus() {
return this.callApi('GetCurrentStatus');
}
/**
* Check server health
* @returns True if server is healthy
*/
async checkHealth() {
try {
await this.callApi('Ping', {});
return true;
}
catch (error) {
console.error('Health check failed:', error);
return false;
}
}
/**
* Clean up resources
*/
destroy() {
if (this.sessionRefreshTimer) {
clearTimeout(this.sessionRefreshTimer);
this.sessionRefreshTimer = null;
}
console.log('SwarmUI client resources cleaned up');
}
}
//# sourceMappingURL=swarmui-client.js.map