LaiTool/src/main/Service/d3.ts

import { BookServiceBasic } from "./ServiceBasic/bookServiceBasic";
import { define } from '../../define/define'
import fs from 'fs'
import { GptService } from "./GPT/gpt";
import { isEmpty } from "lodash";
import path from 'path'
import { CheckFolderExistsOrCreate, CopyFileOrFolder } from "../../define/Tools/file";
import { Base64ToFile, GetImageBase64 } from "../../define/Tools/image";
import { BookBackTaskStatus } from "../../define/enum/bookEnum";
import { MJAction, MJImageType } from "../../define/enum/mjEnum";
import axios from "axios";
import { TaskModal } from "@/model/task";
export class D3Opt {
  bookServiceBasic: BookServiceBasic
  gptService: GptService
  constructor() {
    this.gptService = new GptService()
    this.bookServiceBasic = new BookServiceBasic()
  }

  /**
   * 获取SD的设置，之后要删掉，改为数据库
   */
  private async GetSDSetting() {
    let sdSetting = JSON.parse(await fs.promises.readFile(define.sd_setting, 'utf-8'))
    return sdSetting
  }

  //#region D3生图

  /**
 * 异步发送D3 API图像请求
 * 该函数用于向指定的D3 API发送图像生成请求，通过提供模型、提示和图像大小等参数来定制生成的图像
 * 
 * @param url {string} - API的基础URL字符串，指定了发送请求的目标地址
 * @param key {string} - API密钥字符串，用于认证请求，使函数能够访问API服务
 * @param body {object} - 包含模型、提示和图像大小的对象
 *        - model {string} 模型字段，指定用于生成图像的AI模型
 *        - prompt {string} 提示字段，提供生成图像的文本描述或提示
 *        - size {string} 大小字段，定义生成图像的尺寸
 * @returns 无返回值，但函数内部可能包含对API的异步请求操作
 */
  async D3APIImageRequest(url: string, key: string, body: { model: string; prompt: string; size: string; }) {
    let response = await axios.post(url, { ...body, n: 1 }, {
      headers: {
        Authorization: 'Bearer ' + key
      }
    })
    if (response.data && response.data.data && response.data.data.length > 0) {
      return response.data.data[0].url
    } else {
      return undefined
    }
  }

  /**
 * 异步生成D3图像
 * @param task 任务模型
 * @returns Promise<void>
 */
  async D3ImageGenerate(task: TaskModal.Task): Promise<void> {
    let sdSetting = undefined
    try {
      console.log("D3ImageGenerate", task)
      let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(task.bookTaskDetailId);
      let bookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskDetail.bookTaskId);
      let book = await this.bookServiceBasic.GetBookDataById(bookTask.bookId);
      sdSetting = await this.GetSDSetting()
      await this.gptService.RefreshGptSetting();

      let prompt = bookTaskDetail.prompt;
      let requestUrl = this.gptService.gptUrl
      let uil = new URL(requestUrl);
      let url = `${uil.protocol}//${uil.hostname}` + "/v1/images/generations"

      if (!isEmpty(sdSetting.webui.prompt)) {
        prompt = sdSetting.webui.prompt + ', ' + prompt
      }
      let size = `${sdSetting.webui.width}x${sdSetting.webui.height}`
      // 这边需要判断下size是不是合法的
      if (!['1024x1024', '1024x1792', '1792x1024'].includes(size)) {
        throw new Error('D3 生成图片的尺寸不合法，只支持 1024x1024、1024x1792、1792x1024')
      }
      let model = 'dall-e-3'
      // 一次请求生成一张 多个请求

      let SdOriginalImage = path.join(book.bookFolderPath, 'data/SdOriginalImage');
      await CheckFolderExistsOrCreate(SdOriginalImage);
      let outputFolder = bookTask.imageFolder;
      await CheckFolderExistsOrCreate(outputFolder);
      let inputFolder = path.join(book.bookFolderPath, 'tmp/input')
      await CheckFolderExistsOrCreate(inputFolder);

      let outImagePath = ''
      let subImagePath = []

      let batchSize = sdSetting.setting.batch_size; 7
      for (let i = 0; i < batchSize; i++) {
        const element = batchSize;
        let imageUrl = await this.D3APIImageRequest(url, this.gptService.gptApiKey, {
          model: model,
          prompt: prompt,
          size: size
        })
        // 这边开始处理返回的数据
        if (isEmpty(imageUrl)) {
          throw new Error('D3 生图返回的图片地址为空')
        }
        // 下载指定的文件
        let base64 = await GetImageBase64(imageUrl)
        // 将base64 写出
        let imgPath = path.join(SdOriginalImage, `${bookTaskDetail.name}_${new Date().getTime()}_${i}.png`)
        await Base64ToFile(base64, imgPath);

        // 写出去
        if (bookTask.name == 'output_00001') {
          // 复制一个到input
          let inputImgPath = path.join(inputFolder, `${bookTaskDetail.name}.png`)
          await CopyFileOrFolder(imgPath, inputImgPath)
        }
        if (i == 0) {
          // 复制到对应的文件夹里面
          let outPath = path.join(outputFolder, `${bookTaskDetail.name}.png`)
          await CopyFileOrFolder(imgPath, outPath)
          outImagePath = outPath
        }
        subImagePath.push(imgPath)
      }

      // 结束 开始返回
      // 修改数据库
      await this.bookServiceBasic.UpdateBookTaskDetail(bookTaskDetail.id, {
        outImagePath: path.relative(define.project_path, outImagePath),
        subImagePath: subImagePath.map((item) => path.relative(define.project_path, item))
      })
      await this.bookServiceBasic.UpdateTaskStatus({
        id: task.id,
        status: BookBackTaskStatus.DONE
      });
      let resp = {
        mjApiUrl: url,
        progress: 100,
        category: MJImageType.FLUX_API,
        imageClick: subImagePath.join(','),
        imageShow: subImagePath.join(','),
        messageId: "",
        action: MJAction.IMAGINE,
        status: "success",
        subImagePath: subImagePath,
        outImagePath: outImagePath,
        message: "D3 生成图片成功"
      }
      await this.bookServiceBasic.UpdateBookTaskDetailMjMessage(task.bookTaskDetailId, resp)
      global.newWindow[0].win.webContents.send(task.messageName, {
        code: 1,
        message: "D3 生成图片成功",
        data: {
          ...resp,
          id: bookTaskDetail.id
        }
      })

    } catch (error) {
      let errorMsg = "D3 生成图片失败，错误信息如下：" + error.toString()
      await this.bookServiceBasic.UpdateBookTaskDetailMjMessage(task.bookTaskDetailId, {
        mjApiUrl: "",
        progress: 0,
        category: MJImageType.FLUX_API,
        imageClick: "",
        imageShow: "",
        messageId: "",
        action: MJAction.IMAGINE,
        status: "error",
        message: errorMsg,
      })

      global.newWindow[0].win.webContents.send(task.messageName, {
        code: 0,
        message: errorMsg,
        data: {
          status: 'error',
          message: errorMsg,
          id: task.bookTaskDetailId
        }
      })
      throw error
    }
  }
}

//#endregion