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 } from "../generalTools";

/**
 * MJ原创生图的类
 */
export class MJOriginalImageGenerate {
    constructor(global) {
        this.global = global;
        this.pm = new PublicMethod(global);
        this.discordWorker = new DiscordWorker();
        this.tools = new Tools();
    }

    /**
     * 下载指定的图片地址并且分割
     * @param {*} value 
     */
    async DownloadImageUrlAndSplit(value) {
        try {
            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 param = [];
            // 循环数据，直传需要的数据
            for (let i = 0; i < value.length; i++) {
                const element = value[i];
                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);
            // 开始执行脚本
            let result = await discordSimple.ExecuteScript(define.discordScript, `GetGeneratedMJImageAndSplit(${JSON.stringify(param)})`);

            let res = [];
            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\\${element.image_id}.png`);


                let ds = this.DownloadImageUrlAndSplit(JSON.stringify[element, element.result, image_path]);
                if (ds.code == 0) {
                    throw new Error(ds.message);
                }

                res.push(ds.data);
            }

            // 全部分割完毕，返回
            return {
                code: 1,
                data: res
            }

        } catch (error) {
            return {
                code: 0,
                message: "获取已经生图完成的数据，并获取图片错误，错误信息如下" + 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
            }
        }

    }

    /**
     * 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);


            // 检查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;


                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);
                        await discordSimple.WritePromptToInput(prompt);

                        // 发送命令完成（删除当前正在执行。开始下一个任务）

                    } catch (error) {
                        throw error;
                    }


                }, tasK_id, batch);
            }

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

            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, {
                        code: 0,
                        message: message
                    })
                } else {
                    if (show_global_message) {
                        this.global.newWindow[0].win.webContents.send(DEFINE_STRING.SHOW_MESSAGE_DIALOG, {
                            code: 1,
                            message: "所有MJ生图任务完成"
                        })
                    }
                }
            });


            return {
                code: 1,
            }


        } catch (error) {
            return {
                code: 0,
                message: "MJ生图错误，错误信息如下" + error.message
            }
        }
    }

}