LaiTool/src/main/Public/SD.js

import axios from "axios";
import path from "path";
import { DEFINE_STRING } from "../../define/define_string";
import { define } from "../../define/define";
import { ImageStyleDefine } from "../../define/iamgeStyleDefine";
import { cloneDeep } from 'lodash';
let fspromises = require("fs").promises;
const sharp = require('sharp');
import { SdSettingDefine } from "../../define/setting/sdSettingDefine";
import { PublicMethod } from "./publicMethod";
import { Tools } from "../tools";
import { errorMessage, successMessage } from "../generalTools";
import { SdApi } from "../../api/sdApi";
const { v4: uuidv4 } = require('uuid');

export class SD {
    constructor(global) {
        this.global = global;
        this.pm = new PublicMethod(global);
        this.tools = new Tools();
        this.sdApi = new SdApi();
    }

    /**
     * 获取当前SD服务器所有的lora信息
     */
    async GetAllLoras(baseURL = null) {
        try {
            let data = await this.sdApi.getAllLoras(baseURL);
            return successMessage(data);
        } catch (error) {
            return errorMessage(error.toString());
        }
    }

    /**
     * 获取所有的checkpoint模型
     * @param {*} baseURL 
     * @returns 
     */
    async GetAllSDModel(baseURL = null) {
        try {
            let data = await this.sdApi.getAllSDModel(baseURL);
            return successMessage(data);
        } catch (error) {
            return errorMessage(error.toString());
        }
    }

    /**
     * 获取所有的采样器
     * @param {*} baseURL 
     * @returns 
     */
    async GetAllSamplers(baseURL = null) {
        try {
            let data = await this.sdApi.getAllSamplers(baseURL);
            return successMessage(data);
        } catch (error) {
            return errorMessage(error.toString());
        }
    }

    /**
     * 加载所有的SD数据
     * @param {*} baseURL 
     * @returns 
     */
    async LoadSDServiceData(baseURL = null) {
        try {
            // 加载大模型
            let sd_model = await this.GetAllSDModel(baseURL);
            // 往sd_model中添加一个默认的选项
            sd_model.data.data.unshift({
                title: "无",
                name: "无",
                description: "无",
            })
            // 加载Lora
            let lora = await this.GetAllLoras(baseURL);
            lora.data.data.unshift({
                Key: "无",
                name: "无",
                description: "无",
            })
            // 加载采样器
            let sampler = await this.GetAllSamplers(baseURL);
            sampler.data.data.unshift({
                name: "无",
                description: "无",
            })

            if (!(sd_model.code & lora.code & sampler.code)) {
                throw new Error("获取SD数据错误，请检查SD WEBUI链接！");
            }

            for (let i = 0; i < lora.data.data.length; i++) {
                delete lora.data.data[i].metadata;
            }
            let data = {
                sd_model: sd_model.data.data,
                lora: lora.data.data,
                sampler: sampler.data.data
            }
            // 处理当前获取的数据，保存到配置文件中
            await SdSettingDefine.SavePropertyValue("sd_model", data.sd_model);
            await SdSettingDefine.SavePropertyValue("lora", data.lora);
            await SdSettingDefine.SavePropertyValue("sampler", data.sampler);

            return successMessage(data);

        } catch (error) {
            return errorMessage(error.toString());
        }
    }

    /**
     * 获取图片风格菜单
     * @returns 返回图片风格菜单
     *  
     * */
    async GetImageStyleMenu() {
        try {
            let style = ImageStyleDefine.getImageStyleMenu();
            return {
                code: 1,
                data: style
            }
        } catch (error) {
            return {
                code: 0,
                message: "不可能出现错误"
            }
        }
    }

    /**
     * 获取指定的ID的风格信息，传入的是一个数组
     * @param {*} value id集合
     */
    async GetImageStyleInfomation(value) {
        try {
            if (value) {
                value = JSON.parse(value);
            } else {
                value = [];
            }
            value = value ? value : [];
            let style = ImageStyleDefine.getAllSubStyle();
            let tmp = [];
            for (let i = 0; i < value.length; i++) {
                const element = value[i];
                for (let j = 0; j < style.length; j++) {
                    const item = style[j];
                    if (item.id == element) {
                        tmp.push(item);
                        break;
                    }
                }
            }
            let newSubStyle = cloneDeep(tmp);
            for (let i = 0; i < newSubStyle.length; i++) {
                const element = newSubStyle[i];
                element.image = path.join(define.image_path, "style/" + element.image);
            }
            return {
                code: 1,
                data: newSubStyle
            }

        } catch (error) {
            return {
                code: 0,
                message: error.toString()
            }
        }
    }

    /**
     * 获取指定ID的分类的子风格信息
     * @param {*} value ID
     * @returns 返回ID对应的子风格的详细信息
     */
    async GetStyleImageSubList(value) {
        try {
            let subStyle = ImageStyleDefine.getImagePathById(value);
            let newSubStyle = cloneDeep(subStyle);
            for (let i = 0; i < newSubStyle.length; i++) {
                const element = newSubStyle[i];
                element.image = path.join(define.image_path, "style/" + element.image);
            }
            return {
                code: 1,
                data: newSubStyle
            }

        } catch (error) {
            return {
                code: 0,
                message: error.toString()
            }
        }
    }

    /**
     * 单张生图
     * @param {*} value 0 生图的参数，1 图片的表示，用于保存 ，2 baseUrl
     * @returns 
     */
    async txt2img(value) {
        try {
            value = JSON.parse(value);
            let data = value[0];
            let res = await this.sdApi.txt2img(data);
            // 将base· 64的图片转换为图片
            // 将当前的图片保存到指定的文件夹中，然后返回文件路径，并且可以复制到指定的文件，删除exif信息
            let image_paths = [];
            for (let i = 0; res.data.images && i < res.data.images.length; i++) {
                const element = res.data.images[i];
                let image_data = {
                    base64: element
                }
                // 将保存图片添加到队列中
                let image_name = `sd_${Date.now()}_${uuidv4()}.png`;
                let image_path = path.join(define.temp_sd_image, image_name);
                image_path = await this.tools.saveBase64ToImage(element, image_path);
                image_data["image_path"] = image_path;
                image_paths.push(image_data);
            }
            return successMessage(image_paths);
        } catch (error) {
            return errorMessage("生图失败，错误信息如下：" + error.toString());
        }
    }

    /**
     * 生成一次图片的方法。可以区分模式
     * @param {图片名称 } image
     * @param {任务队列信息} task_list 301198499
     */
    async OneImageGeneration(image, task_list, seed = -1) {
        let taskPath = path.join(this.global.config.project_path, "scripts/task_list.json")
        try {
            let imageJson = JSON.parse(await fspromises.readFile(image + '.json', 'utf-8'));
            let sd_setting = JSON.parse(await fspromises.readFile(define.sd_setting, 'utf-8'));
            let model = imageJson.model;
            let image_json = JSON.parse(await fspromises.readFile(image + '.json', 'utf-8'));
            let image_path = "";
            let target_image_path = "";
            if (image_json.name) {
                image_path = path.join(this.global.config.project_path, `tmp/${task_list.out_folder}/tmp_${image_json.name}`)
                target_image_path = path.join(this.global.config.project_path, `tmp/${task_list.out_folder}/${image_json.name}`)
            } else {
                image_path = image.replaceAll("input_crop", task_list.out_folder).split(".png")[0] + "_tmp.png";
                target_image_path = image.replaceAll("input_crop", task_list.out_folder);
            }
            let image_styles = await ImageStyleDefine.getImageStyleStringByIds(task_list.image_style_list ? task_list.image_style_list : []);
            let prompt = sd_setting.webui.prompt + image_styles;
            // 拼接提示词
            if (task_list.image_style != null) {
                prompt += `((${task_list.image_style})), `;
            }
            if (task_list.lora != null) {
                prompt += `${task_list.lora}, `;
            }
            prompt += imageJson.webui_config.prompt;
            // 判断当前是不是有开修脸修手
            let ADetailer = {
                args: sd_setting.adetailer
            };
            if (model == "img2img") {
                let web_api = this.global.config.webui_api_url + 'sdapi/v1/img2img'
                let sd_config = imageJson["webui_config"];
                sd_config.prompt = prompt;
                sd_config.seed = seed;
                let im = await fspromises.readFile(image, 'binary');
                sd_config.init_images = [new Buffer.from(im, 'binary').toString('base64')];
                if (imageJson.adetailer) {
                    let ta = {
                        ADetailer: ADetailer
                    }
                    sd_config.alwayson_scripts = ta;
                }
                sd_config.height = sd_setting.webui.height;
                sd_config.width = sd_setting.webui.width;
                const response = await axios.post(web_api, sd_config);
                let info = JSON.parse(response.data.info);
                if (seed == -1) {
                    seed = info.seed;
                }
                // 目前是单图出图
                let images = response.data.images;
                let imageData = Buffer.from(images[0].split(",", 1)[0], 'base64');
                await sharp(imageData)
                    .toFile(image_path)
                    .then(async () => {
                        // console.log("图生图成功" + image_path);
                        await this.tools.deletePngAndDeleteExifData(image_path, target_image_path);
                    })
                    .catch(err => {
                        throw new Error(err);
                    });
                return seed;

            } else if (model == "txt2img") {
                let body = {
                    "prompt": prompt,
                    "negative_prompt": imageJson.webui_config.negative_prompt,
                    "seed": seed,
                    "sampler_name": imageJson.webui_config.sampler_name,
                    // 提示词相关性
                    "cfg_scale": imageJson.webui_config.cfg_scale,
                    "width": sd_setting.webui.width,
                    "height": sd_setting.webui.height,
                    "batch_size": 1,
                    "n_iter": 1,
                    "steps": imageJson.webui_config.steps,
                    "save_images": false,
                }
                let web_api = this.global.config.webui_api_url + 'sdapi/v1/txt2img';

                if (imageJson.adetailer) {
                    let ta = {
                        ADetailer: ADetailer
                    }
                    body.alwayson_scripts = ta;
                }
                const response = await axios.post(web_api, body);
                let info = JSON.parse(response.data.info);
                if (seed == -1) {
                    seed = info.seed;
                }
                // 目前是单图出图
                let images = response.data.images;
                let imageData = Buffer.from(images[0].split(",", 1)[0], 'base64');
                await sharp(imageData)
                    .toFile(image_path)
                    .then(async () => {
                        // console.log("文生图成功" + image_path);
                        await this.tools.deletePngAndDeleteExifData(image_path, target_image_path);
                    })
                    .catch(err => {
                        // console.log(err)
                        throw new Error(err);
                    });
                return seed;
            } else {
                throw new Error("SD 模式错误");
            }

        } catch (error) {
            // 当前队列执行失败移除整个批次的任务
            this.global.requestQuene.removeTask(task_list.out_folder, null)
            this.global.fileQueue.enqueue(async () => {
                // 记录失败状态
                let task_list_json = JSON.parse(await fspromises.readFile(taskPath, 'utf-8'));
                // 修改指定的列表的数据
                task_list_json.task_list.map(a => {
                    if (a.id == task_list.id) {
                        a.status = "error";
                        a.errorMessage = error.toString();
                    }
                })
                // 写入
                await fspromises.writeFile(taskPath, JSON.stringify(task_list_json));
                this.global.newWindow[0].win.webContents.send(DEFINE_STRING.IMAGE_TASK_STATUS_REFRESH, {
                    out_folder: task_list.out_folder,
                    status: "error"
                });
            })
            throw error;
        }
    }
}