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 { AwesomeRegx } 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";
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);
    }
    /**
     * 返回指定的人物到前端
     * @param {*} data 
     */
    sendChangeMessage(data) {
        this.global.newWindow[0].win.webContents.send(DEFINE_STRING.DISCORD.MAIN_DISCORD_MESSAGE_CHANGE, data);
    }

    /**
     * 初始化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.request_model;
            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.mj_api_url);
                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.api_key);
                    if (task_res.code == 0) {
                        task_res["id"] = element.id;
                        task_res["mj_api_url"] = mjSetting.mj_api_url;
                        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) {
        try {
            // 获取当前的API url
            let apiUrl = await this.InitMJAPIUrl(mjSetting.mj_api_url);
            let imagine_url = apiUrl.mj_url.imagine;
            let once_get_task = apiUrl.mj_url.once_get_task;
            let task_count = mjSetting.task_count ? mjSetting.task_count : 3;
            let request_model = mjSetting.request_model ? mjSetting.request_model : "relaxed";
            let res;
            // 判断当前的API是哪个
            if (imagine_url.includes("mjapi.deepwl.net")) {
                // DrawAPI(MJ)
                let data = {
                    prompt: prompt,
                    mode: request_model == "fast" ? "FAST" : "RELAX",
                }
                let headers = {
                    "Authorization": mjSetting.api_key
                }
                res = await this.discordAPI.mjApiImagine(imagine_url, data, headers);
            } else if (imagine_url.includes("api.ephone.ai")) {
                // ePhoneAPI
                let headers = {
                    "Authorization": mjSetting.api_key
                }
                let data = {
                    prompt: prompt,
                    botType: "MID_JOURNEY",
                    accountFilter: {
                        modes: [
                            request_model == "fast" ? "FAST" : "RELAX"
                        ]
                    }
                }
                res = await this.discordAPI.mjApiImagine(imagine_url, data, headers);
            }
            // 创建成功，开始下一个
            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.mj_api_url
            });
            this.global.mjGenerateQuene.setCurrentCreateItem(null);
            // 开始监听当前ID是不是的生图任务完成
            // 这边设置一个循环监听，每隔一段时间去请求一次
            let timeoutId;
            let startInterval = () => {
                timeoutId = setTimeout(async () => {
                    // 执行你的操作
                    let task_res = await this.discordAPI.GetMJAPITaskByID(res.result, once_get_task, mjSetting.api_key)
                    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.mj_api_url;
                    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 mjSetting = await this.InitMjSetting();

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

            let style_ids = await this.pm.GetConfigJson(JSON.stringify(["image_style", []]), false);
            let image_styles = await ImageStyleDefine.getImageStyleStringByIds(style_ids.data);

            // 替换风格的逻辑
            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 = "  " + image_styles + old_prompt + (mjSetting.image_suffix ? mjSetting.image_suffix : "");

                // 判断当前生图模式
                let request_model = mjSetting.request_model
                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)
                        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.task_count ? mjSetting.task_count : 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)
        }
    }

}