614 lines
19 KiB
TypeScript
614 lines
19 KiB
TypeScript
import { OptionRealmService } from '@/define/db/service/optionService'
|
||
import { OptionKeyName } from '@/define/enum/option'
|
||
import { SettingModal } from '@/define/model/setting'
|
||
import { TaskModal } from '@/define/model/task'
|
||
import { optionSerialization } from '../option/optionSerialization'
|
||
import {
|
||
CheckFileOrDirExist,
|
||
CheckFolderExistsOrCreate,
|
||
CopyFileOrFolder
|
||
} from '@/define/Tools/file'
|
||
import fs from 'fs'
|
||
import { ValidateJson } from '@/define/Tools/validate'
|
||
import axios from 'axios'
|
||
import { isEmpty } from 'lodash'
|
||
import { BookBackTaskStatus, BookTaskStatus, OperateBookType } from '@/define/enum/bookEnum'
|
||
import { SendReturnMessage } from '@/public/generalTools'
|
||
import { Book } from '@/define/model/book/book'
|
||
import { MJAction } from '@/define/enum/mjEnum'
|
||
import { ImageGenerateMode } from '@/define/data/mjData'
|
||
import path from 'path'
|
||
import { getProjectPath } from '../option/optionCommonService'
|
||
import { SDServiceHandle } from './sdServiceHandle'
|
||
|
||
export class ComfyUIServiceHandle extends SDServiceHandle {
|
||
constructor() {
|
||
super()
|
||
}
|
||
|
||
ComfyUIImageGenerate = async (task: TaskModal.Task) => {
|
||
try {
|
||
await this.InitSDBasic()
|
||
let comfyUISettingCollection = await this.GetComfyUISetting()
|
||
let bookTaskDetail = await this.bookTaskDetailService.GetBookTaskDetailDataById(
|
||
task.bookTaskDetailId as string
|
||
)
|
||
if (bookTaskDetail == null) {
|
||
throw new Error('未找到对应的小说分镜')
|
||
}
|
||
let book = await this.bookService.GetBookDataById(bookTaskDetail.bookId as string)
|
||
if (book == null) {
|
||
throw new Error('未找到对应的小说')
|
||
}
|
||
let bookTask = await this.bookTaskService.GetBookTaskDataById(
|
||
bookTaskDetail.bookTaskId as string
|
||
)
|
||
if (bookTask == null) {
|
||
throw new Error('未找到对应的小说任务')
|
||
}
|
||
|
||
// 调用方法合并提示词
|
||
let mergeRes = await this.MergeSDPrompt(
|
||
task.bookTaskDetailId as string,
|
||
OperateBookType.BOOKTASKDETAIL
|
||
)
|
||
if (mergeRes.code == 0) {
|
||
throw new Error(mergeRes.message)
|
||
}
|
||
// 获取提示词
|
||
bookTaskDetail.prompt = mergeRes.data[0].prompt
|
||
|
||
let prompt = bookTaskDetail.prompt
|
||
let negativePrompt = comfyUISettingCollection.comfyuiSimpleSetting.negativePrompt
|
||
|
||
// 开始组合请求体
|
||
let body = await this.GetComfyUIAPIBody(
|
||
prompt ?? '',
|
||
negativePrompt ?? '',
|
||
comfyUISettingCollection.comfyuiSelectedWorkflow.workflowPath
|
||
)
|
||
|
||
// 开始发送请求
|
||
let resData = await this.SubmitComfyUIImagine(body, comfyUISettingCollection)
|
||
|
||
// 修改任务状态
|
||
await this.bookTaskDetailService.ModifyBookTaskDetailById(task.bookTaskDetailId as string, {
|
||
status: BookTaskStatus.IMAGE
|
||
})
|
||
this.taskListService.UpdateTaskStatus({
|
||
id: task.id as string,
|
||
status: BookBackTaskStatus.RUNNING
|
||
})
|
||
SendReturnMessage(
|
||
{
|
||
code: 1,
|
||
message: '任务已提交',
|
||
id: task.bookTaskDetailId as string,
|
||
data: {
|
||
status: 'submited',
|
||
message: '任务已提交',
|
||
id: task.bookTaskDetailId as string
|
||
} as any
|
||
},
|
||
task.messageName as string
|
||
)
|
||
|
||
await this.FetchImageTask(
|
||
task,
|
||
resData.prompt_id,
|
||
book,
|
||
bookTask,
|
||
bookTaskDetail,
|
||
comfyUISettingCollection
|
||
)
|
||
} catch (error: any) {
|
||
let errorMsg = 'ComfyUI 生图失败，失败信息如下：' + error.toString()
|
||
this.taskListService.UpdateTaskStatus({
|
||
id: task.id as string,
|
||
status: BookBackTaskStatus.FAIL,
|
||
errorMessage: errorMsg
|
||
})
|
||
await this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
|
||
task.bookTaskDetailId as string,
|
||
{
|
||
mjApiUrl: '',
|
||
progress: 0,
|
||
category: ImageGenerateMode.ComfyUI,
|
||
imageClick: '',
|
||
imageShow: '',
|
||
messageId: '',
|
||
action: MJAction.IMAGINE,
|
||
status: 'error',
|
||
message: errorMsg
|
||
}
|
||
)
|
||
|
||
SendReturnMessage(
|
||
{
|
||
code: 0,
|
||
message: errorMsg,
|
||
id: task.bookTaskDetailId as string,
|
||
data: {
|
||
status: 'error',
|
||
message: errorMsg,
|
||
id: task.bookTaskDetailId
|
||
}
|
||
},
|
||
task.messageName as string
|
||
)
|
||
throw error
|
||
}
|
||
}
|
||
|
||
//#region 获取ComfyUI的设置
|
||
/**
|
||
* 获取ComfyUI的设置
|
||
* @returns
|
||
*/
|
||
private async GetComfyUISetting(): Promise<SettingModal.ComfyUISettingCollection> {
|
||
let result = {} as SettingModal.ComfyUISettingCollection
|
||
let optionRealmService = await OptionRealmService.getInstance()
|
||
let comfyuiSimpleSettingOption = optionRealmService.GetOptionByKey(
|
||
OptionKeyName.SD.ComfyUISimpleSetting
|
||
)
|
||
result['comfyuiSimpleSetting'] = optionSerialization<SettingModal.ComfyUISimpleSettingModel>(
|
||
comfyuiSimpleSettingOption,
|
||
'设置 -> ComfyUI 设置'
|
||
)
|
||
|
||
let comfyuiWorkFlowSettingOption = optionRealmService.GetOptionByKey(
|
||
OptionKeyName.SD.ComfyUIWorkFlowSetting
|
||
)
|
||
|
||
let comfyuiWorkFlowList = optionSerialization<SettingModal.ComfyUIWorkFlowSettingModel[]>(
|
||
comfyuiWorkFlowSettingOption,
|
||
'设置 -> ComfyUI 设置'
|
||
)
|
||
result['comfyuiWorkFlowSetting'] = comfyuiWorkFlowList
|
||
|
||
if (comfyuiWorkFlowList.length <= 0) {
|
||
throw new Error('ComfyUI的工作流设置为空，请检查是否正确设置！！')
|
||
}
|
||
|
||
// 获取选中的工作流
|
||
let selectedWorkflow = comfyuiWorkFlowList.find(
|
||
(item) => item.id == result.comfyuiSimpleSetting.selectedWorkflow
|
||
)
|
||
if (selectedWorkflow == null) {
|
||
throw new Error('未找到选中的工作流，请检查是否正确设置！！')
|
||
}
|
||
|
||
// 判断工作流对应的文件是不是存在
|
||
if (!(await CheckFileOrDirExist(selectedWorkflow.workflowPath))) {
|
||
throw new Error('本地未找到选中的工作流文件地址，请检查是否正确设置！！')
|
||
}
|
||
result['comfyuiSelectedWorkflow'] = selectedWorkflow
|
||
|
||
return result
|
||
}
|
||
|
||
//#endregion
|
||
|
||
//#region 组合ComfyUI的请求体
|
||
/**
|
||
* 组合ComfyUI的请求体
|
||
* @param prompt 正向提示词
|
||
* @param negativePrompt 反向提示词
|
||
* @param workflowPath 工作流地址
|
||
*/
|
||
private async GetComfyUIAPIBody(
|
||
prompt: string,
|
||
negativePrompt: string,
|
||
workflowPath: string
|
||
): Promise<string> {
|
||
let jsonContentString = await fs.promises.readFile(workflowPath, 'utf-8')
|
||
|
||
if (!ValidateJson(jsonContentString)) {
|
||
throw new Error('工作流文件内容不是有效的JSON格式，请检查是否正确设置！！')
|
||
}
|
||
|
||
let jsonContent = JSON.parse(jsonContentString)
|
||
// 判断是否是对象
|
||
if (jsonContent !== null && typeof jsonContent === 'object' && !Array.isArray(jsonContent)) {
|
||
// 遍历对象属性
|
||
for (const key in jsonContent) {
|
||
let element = jsonContent[key]
|
||
if (element && element.class_type === 'CLIPTextEncode') {
|
||
if (element._meta?.title === '正向提示词') {
|
||
jsonContent[key].inputs.text = prompt
|
||
}
|
||
if (element._meta?.title === '反向提示词') {
|
||
jsonContent[key].inputs.text = negativePrompt
|
||
}
|
||
}
|
||
if (element && element.class_type === 'KSampler') {
|
||
const crypto = require('crypto')
|
||
const buffer = crypto.randomBytes(8)
|
||
let seed = BigInt('0x' + buffer.toString('hex'))
|
||
|
||
jsonContent[key].inputs.seed = seed.toString()
|
||
} else if (element && element.class_type === 'KSamplerAdvanced') {
|
||
const crypto = require('crypto')
|
||
const buffer = crypto.randomBytes(8)
|
||
let seed = BigInt('0x' + buffer.toString('hex'))
|
||
|
||
jsonContent[key].inputs.noise_seed = seed.toString()
|
||
}
|
||
}
|
||
} else {
|
||
throw new Error('工作流文件内容不是有效的JSON对象格式，请检查是否正确设置！！')
|
||
}
|
||
let result = JSON.stringify({
|
||
prompt: jsonContent
|
||
})
|
||
return result
|
||
}
|
||
|
||
//#endregion
|
||
|
||
//#region 提交ComfyUI生成图片任务
|
||
|
||
private async SubmitComfyUIImagine(
|
||
body: string,
|
||
comfyUISettingCollection: SettingModal.ComfyUISettingCollection
|
||
): Promise<any> {
|
||
let url = comfyUISettingCollection.comfyuiSimpleSetting.requestUrl?.replace(
|
||
'localhost',
|
||
'127.0.0.1'
|
||
)
|
||
if (url.endsWith('/')) {
|
||
url = url + 'api/prompt'
|
||
} else {
|
||
url = url + '/api/prompt'
|
||
}
|
||
var config = {
|
||
method: 'post',
|
||
url: url,
|
||
headers: {
|
||
'User-Agent': 'Apifox/1.0.0 (https://apifox.com)',
|
||
'Content-Type': 'application/json'
|
||
},
|
||
data: body
|
||
}
|
||
|
||
let res = await axios(config)
|
||
let resData = res.data
|
||
// 判断是不是失败
|
||
if (resData.error) {
|
||
let errorNode = ''
|
||
if (resData.node_errors) {
|
||
for (const key in resData.node_errors) {
|
||
errorNode += key + ', '
|
||
}
|
||
}
|
||
let msg = '错误信息：' + resData.error.message + '错误节点：' + errorNode
|
||
throw new Error(msg)
|
||
}
|
||
// 没有错误 判断是不是成功
|
||
if (resData.prompt_id && !isEmpty(resData.prompt_id)) {
|
||
// 成功
|
||
return resData
|
||
} else {
|
||
throw new Error('未知错误，未获取到请求ID，请检查是否正确设置！！')
|
||
}
|
||
}
|
||
|
||
//#endregion
|
||
|
||
//#region 获取出图任务
|
||
|
||
async FetchImageTask(
|
||
task: TaskModal.Task,
|
||
promptId: string,
|
||
book: Book.SelectBook,
|
||
bookTask: Book.SelectBookTask,
|
||
bookTaskDetail: Book.SelectBookTaskDetail,
|
||
comfyUISettingCollection: SettingModal.ComfyUISettingCollection
|
||
) {
|
||
while (true) {
|
||
try {
|
||
let resData = await this.GetComfyUIImageTask(promptId, comfyUISettingCollection)
|
||
|
||
// 判断他的状态是不是成功
|
||
if (resData.status == 'error') {
|
||
// 生图失败
|
||
await this.bookTaskDetailService.ModifyBookTaskDetailById(
|
||
task.bookTaskDetailId as string,
|
||
{
|
||
status: BookTaskStatus.IMAGE_FAIL
|
||
}
|
||
)
|
||
let errorMsg = `MJ生成图片失败，失败信息如下：${resData.message}`
|
||
this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
|
||
task.bookTaskDetailId as string,
|
||
{
|
||
mjApiUrl: comfyUISettingCollection.comfyuiSimpleSetting.requestUrl,
|
||
progress: 100,
|
||
category: ImageGenerateMode.ComfyUI,
|
||
imageClick: '',
|
||
imageShow: '',
|
||
messageId: promptId,
|
||
action: MJAction.IMAGINE,
|
||
status: 'error',
|
||
message: errorMsg
|
||
}
|
||
)
|
||
this.taskListService.UpdateTaskStatus({
|
||
id: task.id as string,
|
||
status: BookBackTaskStatus.FAIL,
|
||
errorMessage: errorMsg
|
||
})
|
||
SendReturnMessage(
|
||
{
|
||
code: 0,
|
||
message: errorMsg,
|
||
id: task.bookTaskDetailId as string,
|
||
data: {
|
||
status: 'error',
|
||
message: errorMsg,
|
||
id: task.bookTaskDetailId
|
||
}
|
||
},
|
||
task.messageName as string
|
||
)
|
||
return
|
||
} else if (resData.status == 'in_progress') {
|
||
// 生图中
|
||
this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
|
||
task.bookTaskDetailId as string,
|
||
{
|
||
mjApiUrl: comfyUISettingCollection.comfyuiSimpleSetting.requestUrl,
|
||
progress: 0,
|
||
category: ImageGenerateMode.ComfyUI,
|
||
imageClick: '',
|
||
imageShow: '',
|
||
messageId: promptId,
|
||
action: MJAction.IMAGINE,
|
||
status: 'running',
|
||
message: '任务正在执行中'
|
||
}
|
||
)
|
||
|
||
SendReturnMessage(
|
||
{
|
||
code: 1,
|
||
message: 'running',
|
||
id: task.bookTaskDetailId as string,
|
||
data: {
|
||
status: 'running',
|
||
message: '任务正在执行中',
|
||
id: task.bookTaskDetailId
|
||
}
|
||
},
|
||
task.messageName as string
|
||
)
|
||
} else {
|
||
let res = await this.DownloadFileUrl(
|
||
resData.imageNames,
|
||
comfyUISettingCollection,
|
||
book,
|
||
bookTask,
|
||
bookTaskDetail
|
||
)
|
||
let projectPath = await getProjectPath()
|
||
// 修改数据库数据
|
||
// 修改数据库
|
||
await this.bookTaskDetailService.ModifyBookTaskDetailById(bookTaskDetail.id as string, {
|
||
outImagePath: path.relative(projectPath, res.outImagePath),
|
||
subImagePath: res.subImagePath.map((item) => path.relative(projectPath, item))
|
||
})
|
||
|
||
this.taskListService.UpdateTaskStatus({
|
||
id: task.id as string,
|
||
status: BookBackTaskStatus.DONE
|
||
})
|
||
|
||
this.bookTaskDetailService.UpdateBookTaskDetailMjMessage(
|
||
task.bookTaskDetailId as string,
|
||
{
|
||
mjApiUrl: comfyUISettingCollection.comfyuiSimpleSetting.requestUrl,
|
||
progress: 100,
|
||
category: ImageGenerateMode.ComfyUI,
|
||
imageClick: '',
|
||
imageShow: '',
|
||
messageId: promptId,
|
||
action: MJAction.IMAGINE,
|
||
status: 'success',
|
||
message: 'ComfyUI 生成图片成功'
|
||
}
|
||
)
|
||
|
||
SendReturnMessage(
|
||
{
|
||
code: 1,
|
||
message: 'ComfyUI 生成图片成功',
|
||
id: task.bookTaskDetailId as string,
|
||
data: {
|
||
status: 'success',
|
||
message: 'ComfyUI 生成图片成功',
|
||
id: task.bookTaskDetailId,
|
||
outImagePath: res.outImagePath + '?t=' + new Date().getTime(),
|
||
subImagePath: res.subImagePath.map((item) => item + '?t=' + new Date().getTime())
|
||
}
|
||
},
|
||
task.messageName as string
|
||
)
|
||
break
|
||
}
|
||
await new Promise((resolve) => setTimeout(resolve, 3000))
|
||
} catch (error) {
|
||
throw error
|
||
}
|
||
}
|
||
}
|
||
|
||
//#endregion
|
||
|
||
//#region 获取comfyui出图任务
|
||
|
||
/**
|
||
* 获取ComfyUI出图任务
|
||
* @param promptId
|
||
* @param comfyUISettingCollection
|
||
*/
|
||
private async GetComfyUIImageTask(
|
||
promptId: string,
|
||
comfyUISettingCollection: SettingModal.ComfyUISettingCollection
|
||
): Promise<any> {
|
||
if (isEmpty(promptId)) {
|
||
throw new Error('未获取到请求ID，请检查是否正确设置！！')
|
||
}
|
||
if (isEmpty(comfyUISettingCollection.comfyuiSimpleSetting.requestUrl)) {
|
||
throw new Error('未获取到ComfyUI的请求地址，请检查是否正确设置！！')
|
||
}
|
||
|
||
let url = comfyUISettingCollection.comfyuiSimpleSetting.requestUrl?.replace(
|
||
'localhost',
|
||
'127.0.0.1'
|
||
)
|
||
if (url.endsWith('/')) {
|
||
url = url + 'api/history'
|
||
} else {
|
||
url = url + '/api/history'
|
||
}
|
||
|
||
var config = {
|
||
method: 'get',
|
||
url: `${url}/${promptId}`,
|
||
headers: {
|
||
'User-Agent': 'Apifox/1.0.0 (https://apifox.com)'
|
||
}
|
||
}
|
||
|
||
let res = await axios.request(config)
|
||
let resData = res.data
|
||
// 判断状态是失败还是成功
|
||
let data = resData[promptId]
|
||
if (data == null) {
|
||
// 还在执行中 或者是任务不存在
|
||
return {
|
||
progress: 0,
|
||
status: 'in_progress',
|
||
message: '任务正在执行中'
|
||
}
|
||
}
|
||
let completed = data.status?.completed
|
||
let outputs = data.outputs
|
||
if (completed && outputs) {
|
||
let imageNames: string[] = []
|
||
for (const key in outputs) {
|
||
let outputNode = outputs[key]
|
||
if (outputNode && outputNode?.images && outputNode?.images.length > 0) {
|
||
for (let i = 0; i < outputNode?.images.length; i++) {
|
||
const element = outputNode?.images[i]
|
||
imageNames.push(element.filename as string)
|
||
}
|
||
}
|
||
}
|
||
return {
|
||
progress: 100,
|
||
status: 'success',
|
||
imageNames: imageNames
|
||
}
|
||
} else {
|
||
return {
|
||
progress: 0,
|
||
status: 'error',
|
||
message: '生图失败，详细失败信息看启动器控制台'
|
||
}
|
||
}
|
||
}
|
||
|
||
//#endregion
|
||
|
||
//#region 请求下载图片
|
||
|
||
/**
|
||
* 请求下载对应的图片
|
||
* @param url
|
||
* @param path
|
||
*/
|
||
private async DownloadFileUrl(
|
||
imageNames: string[],
|
||
comfyUISettingCollection: SettingModal.ComfyUISettingCollection,
|
||
book: Book.SelectBook,
|
||
bookTask: Book.SelectBookTask,
|
||
bookTaskDetail: Book.SelectBookTaskDetail
|
||
): Promise<{
|
||
outImagePath: string
|
||
subImagePath: string[]
|
||
}> {
|
||
let url = comfyUISettingCollection.comfyuiSimpleSetting.requestUrl?.replace(
|
||
'localhost',
|
||
'127.0.0.1'
|
||
)
|
||
if (url.endsWith('/')) {
|
||
url = url + 'api/view'
|
||
} else {
|
||
url = url + '/api/view'
|
||
}
|
||
let outImagePath = ''
|
||
let subImagePath: string[] = []
|
||
for (let i = 0; i < imageNames.length; i++) {
|
||
const imageName = imageNames[i]
|
||
|
||
var config = {
|
||
method: 'get',
|
||
url: `${url}?filename=${imageName}&nocache=${Date.now()}`,
|
||
headers: {
|
||
'User-Agent': 'Apifox/1.0.0 (https://apifox.com)'
|
||
},
|
||
responseType: 'arraybuffer' as 'arraybuffer' // 明确指定类型
|
||
}
|
||
|
||
let res = await axios.request(config)
|
||
|
||
// 检查响应状态和类型
|
||
console.log(`图片下载状态: ${res.status}, 内容类型: ${res.headers['content-type']}`)
|
||
|
||
// 确保得到的是图片数据
|
||
if (!res.headers['content-type']?.includes('image/')) {
|
||
console.error(`响应不是图片: ${res.headers['content-type']}`)
|
||
continue
|
||
}
|
||
|
||
let resData = res.data
|
||
console.log(resData)
|
||
|
||
let subImageFolderPath = path.join(
|
||
bookTask.imageFolder as string,
|
||
`subImage/${bookTaskDetail.name}`
|
||
)
|
||
await CheckFolderExistsOrCreate(subImageFolderPath)
|
||
let outputFolder = bookTask.imageFolder as string
|
||
await CheckFolderExistsOrCreate(outputFolder)
|
||
let inputFolder = path.join(book.bookFolderPath as string, 'tmp/input')
|
||
await CheckFolderExistsOrCreate(inputFolder)
|
||
|
||
// 包含info信息的图片地址
|
||
let infoImgPath = path.join(subImageFolderPath, `${new Date().getTime()}_${i}.png`)
|
||
|
||
// 直接将二进制数据写入文件
|
||
await fs.promises.writeFile(infoImgPath, Buffer.from(resData))
|
||
|
||
if (i == 0) {
|
||
// 复制到对应的文件夹里面
|
||
let outPath = path.join(outputFolder, `${bookTaskDetail.name}.png`)
|
||
await CopyFileOrFolder(infoImgPath, outPath)
|
||
outImagePath = outPath
|
||
}
|
||
subImagePath.push(infoImgPath as string)
|
||
}
|
||
console.log(outImagePath)
|
||
console.log(subImagePath)
|
||
|
||
// 将获取的数据返回
|
||
return {
|
||
outImagePath: outImagePath,
|
||
subImagePath: subImagePath
|
||
}
|
||
}
|
||
|
||
//#endregion
|
||
}
|