LaiTool/src/main/Original/MJOriginalImageGenerate.js

import { DEFINE_STRING } from '../../define/define_string'
import { AsyncQueue } from '../quene'
import { PublicMethod } from '../Public/publicMethod'
import { ImageStyleDefine } from '../../define/iamgeStyleDefine'
import { DiscordSimple } from '../discord/discordSimple'
import { DiscordWorker } from '../discord/discordWorker'
import { Tools } from '../tools'
import path from 'path'
import sharp from 'sharp'
import { define } from '../../define/define'
import { AwesomeHelp } from 'awesome-js'
import { checkStringValueAddSuffix, errorMessage, successMessage } from '../generalTools'
import { ImageSetting } from '../../define/setting/imageSetting'
import { DiscordAPI } from '../../api/discordApi'
import { GPT } from '../Public/GPT'
import { TagDefine } from '../../define/tagDefine'
import { cloneDeep } from 'lodash'
import { LOGGER_DEFINE } from '../../define/logger_define'
import { MJImageType } from '../../define/enum/mjEnum'
import { MJSettingService } from '../../define/db/service/SoftWare/mjSettingService'
const { v4: uuidv4 } = require('uuid')

/**
 * MJ原创生图的类
 */
export class MJOriginalImageGenerate {
  constructor(global) {
    this.global = global
    this.pm = new PublicMethod(global)
    this.discordWorker = new DiscordWorker()
    this.tools = new Tools()
    this.discordAPI = new DiscordAPI()
    this.gpt = new GPT(global)
    this.tagDefine = new TagDefine(global)
  }
  /**
   * 返回指定的人物到前端
   * @param {*} data
   */
  sendChangeMessage(data, message_name = DEFINE_STRING.DISCORD.MAIN_DISCORD_MESSAGE_CHANGE) {
    this.global.newWindow[0].win.webContents.send(message_name, data)
  }

  /**
   * 通过文本自动匹配数据
   * @param {*} value
   */
  async AutoMatchUser(value) {
    try {
      value = JSON.parse(value)
      // 获取所有的角色数据，包括别名
      // 获取所有的角色数据
      let character_tags = await this.tagDefine.getTagDataByTypeAndProperty(
        'dynamic',
        'character_tags'
      )
      if (character_tags.code == 0) {
        return errorMessage(
          '获取角色数据错误，错误信息如下：' + character_tags.message,
          LOGGER_DEFINE.ORIGINAL_AUTO_MATCH_USER
        )
      }
      character_tags = character_tags.data
      if (character_tags.length == 0) {
        return errorMessage('请先添加角色数据', LOGGER_DEFINE.ORIGINAL_AUTO_MATCH_USER)
      }
      let character_tags_data = []
      for (let i = 0; i < character_tags.length; i++) {
        let item = character_tags[i]
        // 这边还要判断是不是显示（只有显示的才hi自动匹配）
        if (!item.hasOwnProperty('isShow') || !item.isShow) {
          continue
        }

        let temp_name = [item.label]
        // 判断当前的数是不是存在别名
        if (item.children && item.children.length > 0) {
          for (let j = 0; j < item.children.length; j++) {
            const element = item.children[j]
            temp_name.push(element.label)
          }
        }
        character_tags_data.push({
          key: item.key,
          value: item.value,
          names: temp_name
        })
      }
      if (character_tags_data.length == 0) {
        return errorMessage(
          '当前没有显示的角色数据，请先选择哪些是要显示的角色数据',
          LOGGER_DEFINE.ORIGINAL_AUTO_MATCH_USER
        )
      }

      for (let i = 0; i < value.length; i++) {
        const element = value[i]

        let res_data = {
          code: 1,
          id: element.id, // 当前 data 的ID
          match_character: []
        }

        // 获取当前的字幕数据
        let temp_sub = []
        for (let j = 0; element.suValue && j < element.suValue.length; j++) {
          const element = array[j]
          temp_sub.push(element.srt_value)
        }
        let word = ''
        if (temp_sub.length == 0) {
          word = element.after_gpt
        } else {
          word = temp_sub.join(',')
        }

        let match_keys = []

        // 开始循环判断，只要又一个满足就跳出新婚换
        for (let j = 0; j < character_tags_data.length; j++) {
          const item = character_tags_data[j]
          let names = AwesomeHelp.makeSensitiveMap(item.names)
          // 开始判断
          let name_res = AwesomeHelp.checkSensitiveWord(word, false, names)
          if (name_res.size > 0) {
            match_keys.push(item.key)
          }
        }

        // 判断是不是又匹配到的数据
        if (match_keys.length > 0) {
          // 进行数据的处理，通过对应的key，获取对应的数据，将所有的数进行返回
          for (let i = 0; i < match_keys.length; i++) {
            const item = match_keys[i]
            let index = character_tags.findIndex((x) => x.key == item)
            if (index == -1) {
              continue
            }
            let temp_item_data = cloneDeep(character_tags[index])
            if (temp_item_data.children) {
              delete temp_item_data.children
            }
            res_data.match_character.push(temp_item_data)
          }
        }

        // 开始往前端传递数据
        this.sendChangeMessage(res_data, DEFINE_STRING.MJ.MACTH_USER_RETURN)
      }
      return successMessage(null, '人物标签自动匹配完成', LOGGER_DEFINE.ORIGINAL_AUTO_MATCH_USER)
    } catch (error) {
      return errorMessage(
        '通过文本自动匹配数据错误，错误信息如下：' + error.message,
        LOGGER_DEFINE.ORIGINAL_AUTO_MATCH_USER
      )
    }
  }

  /**
   * 初始化MJ设置
   */
  async InitMjSetting() {
    let mjSetting_res = await ImageSetting.GetDefineConfigJsonByProperty(
      JSON.stringify(['img_base', 'mj_config', false, null])
    )
    if (mjSetting_res.code == 0 || !mjSetting_res.data) {
      throw new Error('请先添加MJ配置')
    }
    let mjSetting = mjSetting_res.data
    return mjSetting
  }

  /**
   * 初始化MJ API的URL
   */
  async InitMJAPIUrl(id) {
    let mj_api = (await this.gpt.GetGPTBusinessOption('all', (value) => value.mj_url)).data
    let mj_api_url_index = mj_api.findIndex((item) => item.value == id)
    if (mj_api_url_index == -1) {
      throw new Error('没有找到对应的MJ API的配置，请先检查配置')
    }
    return mj_api[mj_api_url_index]
  }

  /**
   * 下载指定的图片地址并且分割
   * @param {*} value
   */
  async DownloadImageUrlAndSplit(value) {
    try {
      console.log(value)
      value = JSON.parse(value)
      let element = value[0]
      let iamge_url = value[1]
      let image_path = ''
      if (value.length > 2) {
        image_path = value[2]
      } else {
        image_path = path.join(
          global.config.project_path,
          `data\\MJOriginalImage\\${element.id}.png`
        )
      }
      // 判断是不是一个链接
      const urlRegex = /^(http|https):\/\/[^ "]+$/
      if (!urlRegex.test(iamge_url)) {
        throw new Error('指定的图片地址不是一个链接')
      }

      // 这边开始下载对应的图片
      await this.tools.downloadFileUrl(iamge_url, image_path)

      // 将下载的图片进行分割
      let split_res = await this.ImageSplit(JSON.stringify([image_path, element.name]))
      if (split_res.code == 0) {
        throw new Error(split_res.message)
      }

      element.image_click = iamge_url
      element.subImagePath = split_res.data.subImagePath
      element.outImagePath = split_res.data.outImagePath
      element['image_path'] = image_path

      return {
        code: 1,
        data: element
      }
    } catch (error) {
      return {
        code: 0,
        message: '下载指定的图片地址并且分割错误，错误信息如下：' + error.message
      }
    }
  }

  /**
   * 获取已经生图完成的数据，并获取图片
   * @param {*} value
   * @returns
   */
  async GetGeneratedMJImageAndSplit(value) {
    try {
      value = JSON.parse(value)
      let mjSetting = await this.InitMjSetting()
      let request_model = mjSetting.requestModel
      let result = []
      // 浏览器生图模式
      if (request_model == 'browser_mj') {
        let param = []
        // 循环数据，直传需要的数据
        for (let i = 0; i < value.length; i++) {
          const element = value[i]
          // 一般进度大于 50 会出现图片，
          if (!element.mj_message) {
            continue
          }
          if (element.mj_message.progress && element.mj_message.progress == 100) {
            // 判断 image_path 是不是存在。
            if (item.mj_message.image_id && !element.mj_message.image_path) {
              // 通过当前的image_id获取图片
              param.push({
                id: element.id,
                image_id: element.mj_message.image_id,
                name: element.name
              })
            }
          }
        }

        // 判断窗口是不是开启
        let discordWin = await this.discordWorker.CheckDiscordWindowIsOpenAndLoad()
        // 执行采集图片的脚本
        // 开始写入
        let discordSimple = new DiscordSimple(discordWin)
        // 开始执行脚本
        result = await discordSimple.ExecuteScript(
          define.discordScript,
          `GetGeneratedMJImageAndSplit(${JSON.stringify(param)})`
        )
      } else if (request_model == 'api_mj') {
        let mj_api = await this.InitMJAPIUrl(mjSetting.apiSetting.mjApiUrl)
        let once_get_task = mj_api.mj_url.once_get_task

        // 请求
        for (let i = 0; i < value.length; i++) {
          const element = value[i]
          if (element.mj_message.progress == 100) {
            continue
          }
          if (element.mj_message.progress.status == 'success') {
            continue
          }

          let task_res = await this.discordAPI.GetMJAPITaskByID(
            element.mj_message.message_id,
            once_get_task,
            mjSetting.apiSetting.apiKey
          )
          if (task_res.code == 0) {
            task_res['id'] = element.id
            task_res['mj_api_url'] = mjSetting.apiSetting.mjApiUrl
            this.sendChangeMessage()
          }
          // 判断进度是不是百分百
          if (task_res.progress != 100) {
            continue
          }

          result.push({
            id: element.id,
            image_id: null,
            result: task_res.image_click,
            name: element.name
          })
        }
      }

      let res = []
      // 判断返回的数据是不是一个字符串
      if (typeof result == 'string') {
        result = JSON.parse(result)
      }

      // 将返回的数据进行分割
      for (let i = 0; i < result.length; i++) {
        const element = result[i]
        let image_path = path.join(
          global.config.project_path,
          `data\\MJOriginalImage\\${uuidv4()}.png`
        )
        let ds = await this.DownloadImageUrlAndSplit(
          JSON.stringify([element, element.result, image_path])
        )
        if (ds.code == 0) {
          throw new Error(ds.message)
        }
        // 修改数据。
        ds.data['progress'] = 100
        ds.data['status'] = 'success'
        res.push(ds.data)
      }
      // 全部分割完毕，返回
      return successMessage(res)
    } catch (error) {
      return errorMessage('获取已经生图完成的数据，并获取图片错误，错误信息如下' + error.message)
    }
  }

  // MJ生成的图片分割
  async ImageSplit(value) {
    try {
      value = JSON.parse(value)
      let inputPath = value[0]
      let r_name = value[1]
      let outputDir = path.join(this.global.config.project_path, `data\\MJOriginalImage`)
      const metadata = await sharp(inputPath).metadata()
      const smallWidth = metadata.width / 2
      const smallHeight = metadata.height / 2
      let times = new Date().getTime()
      let imgs = []
      let first_p = path.join(this.global.config.project_path, `tmp\\output_crop_00001\\${r_name}`)

      for (let i = 0; i < 4; i++) {
        const xOffset = i % 2 === 0 ? 0 : smallWidth
        const yOffset = Math.floor(i / 2) * smallHeight
        let out_file = path.join(outputDir, `/${r_name}_${times}_${i}.png`)
        await sharp(inputPath)
          .extract({
            left: xOffset,
            top: yOffset,
            width: smallWidth,
            height: smallHeight
          })
          .resize(smallWidth, smallHeight)
          .toFile(out_file)

        imgs.push(out_file)

        // 将第一个图片复制一个到指定的位置
        if (i == 0) {
          await this.tools.copyFileOrDirectory(out_file, first_p)
          // 复制一份到input
          let input_p = path.join(this.global.config.project_path, `tmp\\input_crop\\${r_name}`)
          await this.tools.copyFileOrDirectory(out_file, input_p)
        }
      }

      return {
        code: 1,
        data: {
          subImagePath: imgs,
          outImagePath: first_p
        }
      }
    } catch (err) {
      return {
        code: 0,
        message: 'MJ图片切割错误，错误信息如下' + err.message
      }
    }
  }

  /**
   * 调用API生图的方法
   * @param {*} element
   * @param {*} mjSetting
   */
  async MJImagineRequest(
    element,
    mjSetting,
    prompt,
    tasK_id = null,
    batch = null,
    request_model = 'api_mj'
  ) {
    try {
      if (mjSetting.apiSetting == null) {
        throw new Error('没有API设置，请先设置API设置')
      }
      let apiUrl
      if (mjSetting.requestModel == MJImageType.API_MJ) {
        // 获取当前的API url
        apiUrl = await this.InitMJAPIUrl(mjSetting.apiSetting.mjApiUrl)
      } else if (mjSetting.requestModel == MJImageType.REMOTE_MJ) {
        apiUrl = {
          mj_url: {
            imagine: define.remotemj_api + 'mj/submit/imagine',
            once_get_task: define.remotemj_api + 'mj/task/${id}/fetch'
          }
        }
      } else {
        throw new Error('未知的生图模式，请检查配置')
      }
      let imagine_url = apiUrl.mj_url.imagine
      let once_get_task = apiUrl.mj_url.once_get_task

      let task_count = mjSetting.taskCount ? mjSetting.taskCount : 3
      let mj_speed = mjSetting.apiSetting.mjSpeed ? mjSetting.apiSetting.mjSpeed : 'relaxed'
      let res
      // 判断当前的API是哪个
      if (imagine_url.includes('mjapi.deepwl.net')) {
        // DrawAPI(MJ)
        let data = {
          prompt: prompt,
          mode: mj_speed == 'fast' ? 'FAST' : 'RELAX'
        }
        let headers = {
          Authorization: mjSetting.apiSetting.apiKey
        }
        res = await this.discordAPI.mjApiImagine(imagine_url, data, headers)

        if (res.code == 24) {
          throw new Error('提示词包含敏感词，请修改后重试')
        }
      } else if (
        imagine_url.includes('api.ephone.ai') ||
        imagine_url.includes('https://laitool.net')
      ) {
        // ePhoneAPI
        let headers = {
          Authorization: mjSetting.apiSetting.apiKey
        }
        let data = {
          prompt: prompt,
          botType: 'MID_JOURNEY',
          accountFilter: {
            modes: [mj_speed == 'fast' ? 'FAST' : 'RELAX']
          }
        }
        res = await this.discordAPI.mjApiImagine(imagine_url, data, headers)
      } else if (
        imagine_url.includes(define.remotemj_api) &&
        request_model == MJImageType.REMOTE_MJ
      ) {
        // 代理模式
        let headers = {
          'mj-api-secret': define.API
        }
        let data = {
          prompt: prompt,
          botType: 'MID_JOURNEY',
          accountFilter: {
            remark: this.global.machineId
          }
        }
        res = await this.discordAPI.mjApiImagine(imagine_url, data, headers)
      }

      this.global.mjGenerateQuene.setCurrentCreateItem(null)
      // 错误检查
      if (res.code == 23) {
        // 任务队列已满，及结束该任务，然后开始下一个任务，并将该任务重新排序
        this.global.mjGenerateQuene.removeTaskProgress((taskProgress) => {
          return taskProgress.filter((item) => item?.id != element.id)
        })
        this.sendChangeMessage({
          code: 0,
          status: 'error',
          message: '任务队列已满，任务结束会重新排序，并等待3分钟，开始后面的任务',
          id: element.id
        })
        await this.tools.delay(40000)
        // 重新将当前任务加入队列
        this.global.mjGenerateQuene.enqueue(
          async () => {
            this.global.mjGenerateQuene.setCurrentCreateItem(element)
            await this.MJImagineRequest(element, mjSetting, prompt)
          },
          tasK_id,
          batch
        )

        this.global.mjGenerateQuene.startNextTask(task_count)
        return
      }

      if (res.code != 1 && res.code != 22) {
        // 未知错误，将当前任务删除，开始下一个任务
        this.global.mjGenerateQuene.removeTaskProgress((taskProgress) => {
          return taskProgress.filter((item) => item?.id != element.id)
        })
        this.global.mjGenerateQuene.startNextTask(task_count)
        this.sendChangeMessage({
          code: 0,
          status: 'error',
          message: '未知错误，可联系管理员排查，' + res.description,
          id: element.id
        })
        return
      }
      // 创建成功，开始下一个
      this.sendChangeMessage({
        code: 1,
        type: 'created',
        category: 'api_mj',
        message_id: res.result,
        image_click: null,
        image_show: null,
        id: element.id,
        progress: 0,
        mj_api_url: mjSetting.apiSetting.mjApiUrl
      })
      // 开始监听当前ID是不是的生图任务完成
      // 这边设置一个循环监听，每隔一段时间去请求一次
      let timeoutId
      let startInterval = () => {
        timeoutId = setTimeout(async () => {
          // 执行你的操作
          let task_res = await this.discordAPI.GetMJAPITaskByID(
            res.result,
            once_get_task,
            mjSetting.apiSetting.apiKey
          )
          console.log(task_res)
          // 判断他的状态是不是成功
          if (task_res.code == 0) {
            // 将但钱任务删除
            this.global.mjGenerateQuene.removeTaskProgress((taskProgress) => {
              return taskProgress.filter((item) => item?.id != element.id)
            })
            // 停止当前循环
            clearTimeout(timeoutId)
          } else {
            if (task_res.progress == 100) {
              // 将但钱任务删除
              this.global.mjGenerateQuene.removeTaskProgress((taskProgress) => {
                return taskProgress.filter((item) => item?.id != element.id)
              })
              task_res.type = 'finished'
              // 下载对应的图片
              let image_path = path.join(
                this.global.config.project_path,
                `data\\MJOriginalImage\\${task_res.message_id}.png`
              )
              // 这边开始下载对应的图片
              await this.tools.downloadFileUrl(task_res.image_click, image_path)
              task_res['image_path'] = image_path
              // 开始下一个任务
              this.global.mjGenerateQuene.startNextTask(task_count)
            } else {
              // 当获取的图片的进度小于100的时候，继续监听
              startInterval()
            }
          }
          task_res['id'] = element.id
          task_res['mj_api_url'] = mjSetting.apiSetting.mjApiUrl
          this.sendChangeMessage(task_res)
        }, 5000)
      }
      startInterval()
      this.global.mjGenerateQuene.startNextTask(task_count)
    } catch (error) {
      this.sendChangeMessage({
        code: 0,
        status: 'error',
        message: error.message,
        id: element.id
      })
      throw new Error('MJ API 出图错误，错误信息如下：' + error.message)
    }
  }

  /**
   * MJ 原创生图
   * @param {*} value
   */
  async OriginalMJImageGenerate(value) {
    try {
      let data = value[0]
      if (value[1]) {
        data = JSON.parse(data)
      }
      let show_global_message = value[2]
      let batch = DEFINE_STRING.QUEUE_BATCH.MJ_ORIGINAL_GENERATE_IMAGE

      // 判断存放的文件夹是不是存在，不存在的话创建
      let outputDir = path.join(this.global.config.project_path, `data\\MJOriginalImage`)
      await this.tools.checkFolderExistsOrCreate(outputDir)
      let fileExist = await this.tools.checkExists(outputDir)
      if (!fileExist) {
        await this.tools.createDirectory(outputDir)
      }
      // 判断该当前tmp\output_crop_00001文件夹是不是存在，不存在创建
      let output_crop_00001 = path.join(this.global.config.project_path, `tmp\\output_crop_00001`)
      await this.tools.checkFolderExistsOrCreate(output_crop_00001)

      // 获取MJ配置，从数据库中
      let _mjSettingService = await MJSettingService.getInstance()
      let mjSettings = _mjSettingService.GetMJSettingTreeData()
      if (mjSettings.code == 0) {
        throw new Error(mjSettings.message)
      }
      let mjSetting = mjSettings.data

      // 检查this.global中是不是又mj队列，没有的话创建一个
      if (!this.global.mjGenerateQuene) {
        this.global.mjGenerateQuene = new AsyncQueue(this.global, 1, true)
      }

      // 替换风格的逻辑
      let current_task = null

      for (let i = 0; i < data.length; i++) {
        const element = data[i]
        let tasK_id = `${batch}_${element.name}_${element.id}`

        let old_prompt = element.prompt
        // 拼接提示词
        // 图生图的链接
        // 获取风格词 + 命令后缀
        let prompt = old_prompt + (mjSetting.imageSuffix ? mjSetting.imageSuffix : '')

        // 判断当前生图模式
        let request_model = mjSetting.requestModel
        switch (request_model) {
          case 'api_mj':
            this.global.mjGenerateQuene.enqueue(
              async () => {
                this.global.mjGenerateQuene.setCurrentCreateItem(element)
                await this.MJImagineRequest(element, mjSetting, prompt, tasK_id, batch)
              },
              tasK_id,
              batch
            )
            break

          case MJImageType.REMOTE_MJ:
            this.global.mjGenerateQuene.enqueue(
              async () => {
                this.global.mjGenerateQuene.setCurrentCreateItem(element)
                await this.MJImagineRequest(
                  element,
                  mjSetting,
                  prompt,
                  tasK_id,
                  batch,
                  request_model
                )
              },
              tasK_id,
              batch
            )
            break

          case 'browser_mj':
            this.global.mjGenerateQuene.enqueue(
              async () => {
                try {
                  this.global.mjGenerateQuene.setCurrentCreateItem(element)
                  // 开始进行mj生图
                  current_task = element.name
                  // 判断窗口是不是开启
                  let discordW = await this.discordWorker.CheckDiscordWindowIsOpenAndLoad()
                  // 开始写入
                  let discordSimple = new DiscordSimple(discordW, mjSetting)
                  await discordSimple.WritePromptToInput(prompt)
                  // 发送命令完成（删除当前正在执行。开始下一个任务）
                } catch (error) {
                  throw error
                }
              },
              tasK_id,
              batch
            )
          default:
            break
        }
      }

      // 判断该当前正在执行的人物队列数（小于设置的数量，开始一个任务）
      this.global.mjGenerateQuene.startNextTask(mjSetting.taskCount ? mjSetting.taskCount : 3)

      this.global.requestQuene.setBatchCompletionCallback(batch, (failedTasks) => {
        if (failedTasks.length > 0) {
          let message = `
                MJ生图任务都已完成。
                但是以下任务执行失败：
                `
          failedTasks.forEach(({ taskId, error }) => {
            message += `${taskId}-, \n 错误信息: ${error}` + '\n'
          })

          this.global.newWindow[0].win.webContents.send(
            DEFINE_STRING.SHOW_MESSAGE_DIALOG,
            errorMessage(message)
          )
        } else {
          if (show_global_message) {
            this.global.newWindow[0].win.webContents.send(
              DEFINE_STRING.SHOW_MESSAGE_DIALOG,
              successMessage(null, '所有MJ生图任务完成')
            )
          }
        }
      })
      return successMessage(null)
    } catch (error) {
      return errorMessage('MJ生图错误，错误信息如下' + error.message)
    }
  }
}