LaiTool/src/main/Service/Flux/flux.ts

import { define } from '../../../define/define'
import fs from "fs";
import { ImageStyle } from "../Book/imageStyle";
import { BookServiceBasic } from "../ServiceBasic/bookServiceBasic";
import { isEmpty } from 'lodash';
import axios from 'axios';
const fspromise = fs.promises
import path from 'path'
import { CheckFolderExistsOrCreate, CopyFileOrFolder, DeleteFileExifData } 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 { GptService } from '../GPT/gpt';

export class FluxOpt {
  gptService: GptService
  bookServiceBasic: BookServiceBasic
  constructor() {
    this.bookServiceBasic = new BookServiceBasic()
    this.gptService = new GptService()
  }


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


  //#region flux forge 生图
  /**
   * 使用flux forge生图
   * @param task 
   */
  async FluxForgeImage(task: TaskModal.Task): Promise<void> {
    let sdSetting = undefined
    try {
      // 开始生图
      sdSetting = await this.GetSDSetting()

      let bookTaskDetail = await this.bookServiceBasic.GetBookTaskDetailDataById(task.bookTaskDetailId);
      let bookTask = await this.bookServiceBasic.GetBookTaskDataById(bookTaskDetail.bookTaskId);
      let book = await this.bookServiceBasic.GetBookDataById(bookTask.bookId);

      let prompt = bookTaskDetail.prompt;
      let url = sdSetting.setting.webui_api_url
      if (url.endsWith('/')) {
        url = url + "sdapi/v1/txt2img"
      } else {
        url = url + "/sdapi/v1/txt2img"
      }
      if (!isEmpty(sdSetting.webui.prompt)) {
        prompt = sdSetting.webui.prompt + ', ' + prompt
      }

      // 判断当前是不是有开修脸修手
      let ADetailer = {
        args: sdSetting.adetailer
      }

      // 种子默认 -1，随机
      let seed = -1;
      let body = {
        scheduler: 'Simple',
        prompt: prompt,
        seed: seed,
        sampler_name: sdSetting.webui.sampler_name,
        // 提示词相关性
        cfg_scale: sdSetting.webui.cfg_scale,
        distilled_cfg_scale: 3.5,
        width: sdSetting.webui.width,
        height: sdSetting.webui.height,
        batch_size: sdSetting.setting.batch_size,
        steps: sdSetting.webui.steps,
        save_images: false,
        tiling: false,
        override_settings_restore_afterwards: true
      }

      if (bookTaskDetail.adetailer) {
        body['alwayson_scripts'] = {
          "ADetailer": ADetailer
        }
      }

      const response = await axios.post(url, body)
      // TODO 对SD图片种子的一些处理，待定
      // let info = JSON.parse(res.data.info)
      // if (seed == -1) {
      //   seed = info.seed
      // }
      let images = response.data.images
      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 subImagePath = []
      let outImagePath = ''
      // 开始写出图片
      for (let i = 0; i < images.length; i++) {
        const element = images[i];
        // 包含info信息的图片地址
        let infoImgPath = path.join(SdOriginalImage, `info_${bookTaskDetail.name}_${new Date().getTime()}_${i}.png`)
        // 不包含info信息的图片地址
        let imgPath = path.join(SdOriginalImage, `${bookTaskDetail.name}_${new Date().getTime()}_${i}.png`)
        await Base64ToFile(element, infoImgPath)
        // 这边去图片信息
        await DeleteFileExifData(path.join(define.package_path, 'exittool/exiftool.exe'), infoImgPath, 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_FORGE,
        imageClick: subImagePath.join(','),
        imageShow: subImagePath.join(','),
        messageId: "",
        action: MJAction.IMAGINE,
        status: "success",
        subImagePath: subImagePath,
        outImagePath: outImagePath,
        message: "FLUX FORGE 生成图片成功"
      }
      await this.bookServiceBasic.UpdateBookTaskDetailMjMessage(task.bookTaskDetailId, resp)
      global.newWindow[0].win.webContents.send(task.messageName, {
        code: 1,
        message: "FLUX FORGE 生成图片成功",
        data: {
          ...resp,
          id: bookTaskDetail.id
        }
      })
    } catch (error) {
      let errorMsg = "FLUX FORGE 生成图片失败，错误信息如下：" + error.toString()

      await this.bookServiceBasic.UpdateBookTaskDetailMjMessage(task.bookTaskDetailId, {
        mjApiUrl: sdSetting ? sdSetting.setting.webui_api_url : "",
        progress: 0,
        category: MJImageType.FLUX_FORGE,
        imageClick: "",
        imageShow: "",
        messageId: "",
        action: MJAction.IMAGINE,
        status: "error",
        message: errorMsg
      })
      await this.bookServiceBasic.UpdateTaskStatus({
        id: task.id,
        status: BookBackTaskStatus.FAIL,
        errorMessage: errorMsg
      });

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

      throw error
    }
  }
  //#endregion


  //#region flux api

  /**
   * 发送 FLUX API 请求
   * @param url 请求的地址
   * @param key  请求的key
   * @param body 请求的请求体
   * @returns 
   */
  async FluxAPIImageRequest(url: string, key: string, body: { model: string; prompt: string; size: string; }): Promise<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
    }
  }

  /**
   * 调用 FLUX API 生成图片
   * @param task 
   */
  async FluxAPIImage(task: TaskModal.Task): Promise<void> {
    let sdSetting = undefined
    try {
      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}`
      let model = sdSetting.flux.model
      // 一次请求生成一张 多个请求

      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;
      for (let i = 0; i < batchSize; i++) {
        const element = batchSize;
        let imageUrl = await this.FluxAPIImageRequest(url, this.gptService.gptApiKey, {
          model: model,
          prompt: prompt,
          size: size
        })
        // 这边开始处理返回的数据
        if (isEmpty(imageUrl)) {
          throw new Error('FLUX 生图返回的图片地址为空')
        }
        // 下载指定的文件
        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: "FLUX API 生成图片成功"
      }
      await this.bookServiceBasic.UpdateBookTaskDetailMjMessage(task.bookTaskDetailId, resp)
      global.newWindow[0].win.webContents.send(task.messageName, {
        code: 1,
        message: "FLUX API 生成图片成功",
        data: {
          ...resp,
          id: bookTaskDetail.id
        }
      })

    } catch (error) {
      let errorMsg = "FLUX API 生成图片失败，错误信息如下：" + 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
    }
  }
}