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 '../Public/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 + '?t=' + new Date().getTime() ) } 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 } } }