shuzhiren-comfyui/任务队列后端/outside/outPlatforms/comfyui.js

88 lines
2.6 KiB
JavaScript

import dotenv from 'dotenv';
dotenv.config();
const DISPATCH_TYPES = {
RUNNINGHUB: 'runninghub',
MESSAGEDISPATCHER: 'messageDispatcher'
};
export function getGenerateUrl(dispatchType = DISPATCH_TYPES.RUNNINGHUB) {
if (dispatchType === DISPATCH_TYPES.MESSAGEDISPATCHER) {
return process.env.MESSAGE_DISPATCHER_URL;
}
return process.env.RunningHub_URL;
}
export function getGenerateHeader(apikey, dispatchType = DISPATCH_TYPES.RUNNINGHUB, jwtToken = null) {
if (dispatchType === DISPATCH_TYPES.MESSAGEDISPATCHER) {
return {
'Content-Type': 'application/json',
'Authorization': `Bearer ${jwtToken || ''}`
};
}
return {
'Content-Type': 'application/json',
'Host': 'www.runninghub.cn'
};
}
export function getGenerateBody(task, dispatchType = DISPATCH_TYPES.RUNNINGHUB, jwtToken = null) {
const taskData = JSON.parse(task.payload);
if (dispatchType === DISPATCH_TYPES.MESSAGEDISPATCHER) {
const payload = { ...taskData, apiKey: jwtToken, webhookUrl: process.env.CALLBACK_URL };
console.log('[comfyui - messageDispatcher] 请求体:', payload);
return JSON.stringify(payload);
}
const payload = { ...taskData, apiKey: task.apikey, webhookUrl: process.env.CALLBACK_URL };
return JSON.stringify(payload);
}
export function getQueryUrl() {
return process.env.CALLBACK_URL;
}
export function getTaskStatus() {
if (response.task_status === 'SUCCESS') return true;
}
export async function getSuccessTasks(response, dispatchType = DISPATCH_TYPES.RUNNINGHUB) {
if (dispatchType === DISPATCH_TYPES.MESSAGEDISPATCHER) {
try {
const res = await response.json();
console.log('[comfyui - messageDispatcher] 响应:\n', res);
if (res.success === true && res.data && res.data.requestId) {
return { msg: 'success', code: 0, data: { taskId: res.data.requestId } };
} else {
return { message: res, type: 2 };
}
} catch (error) {
console.error('[comfyui - messageDispatcher] 解析响应失败:', error);
return { message: error.message, type: 2 };
}
}
const res = await response.json();
console.log('[comfyui - runninghub]:\n', res);
if (res.msg === 'success' && res.code === 0) {
return res.data.taskId;
} else {
return { message: res, type: 2 };
}
}
export async function getTaskResult(response) {
const res = await JSON.parse(response);
const files = [];
if (res.msg === 'success' && res.code === 0) {
for (const file of res.data)
files.push(file.fileUrl);
return { files: files[0], type: 1 };
} else {
return { message: res.msg, type: 2 };
}
}
export { DISPATCH_TYPES };