import { Context, Logger, Schema, trimSlash } from 'koishi' import { } from 'koishi-plugin-puppeteer' export const name = 'sd-switch' export const inject = { optional: ['puppeteer'] } export interface Config { endpoint?: string inputTimeout?: number vaeList?: Array sendAsImage?: boolean picWidth?: number } export const Config: Schema = Schema.intersect([ Schema.object({ endpoint: Schema.string().description('SD-WebUI 服务器地址。').default('http://127.0.0.1:7860'), inputTimeout: Schema.number().description('选择模型的等待时间。').default(10000), vaeList: Schema.array( Schema.string() ).description('vae 列表，请输入去掉结尾 .pt 后的 vae 文件名。'), sendAsImage: Schema.boolean().default(false).description('是否以图片形式发送。'), }).description('基础配置'), Schema.union([ Schema.object({ sendAsImage: Schema.const(true).required(), picWidth: Schema.number().default(256).description('图片输出的宽度。'), }), Schema.object({}), ]) ]) const logger = new Logger(name) const MODELS_ENDPOINT = "/sdapi/v1/sd-models" const OPTIONS_ENDPOINT = "/sdapi/v1/options" export function apply(ctx: Context, config: Config) { ctx.i18n.define('zh', require('./locales/zh')) const cmd = ctx.command('sd-switch') .alias('切换模型') .action(async ({ session }) => { const models = [] let vae = '' let model = '' async function render(content: string, picWidth: number) { // https://github.com/ifrvn/koishi-plugin-send-as-image return ctx.puppeteer.render( `

${content}

` ) } async function send(text: string) { if (config.sendAsImage) { if (ctx.puppeteer) { session.send(await render(text, config.picWidth)) } else session.send(session.text('.pptrErr')) } else session.send(text) } async function input(max: number) { const value = +await session.prompt(config.inputTimeout) if (value * 0 === 0 && Math.floor(value) === value && value > 0 && value <= max) { return value } session.send(session.text('.inputErr')) throw new Error('inputErr') } async function getInfo() { session.send(session.text('.inQuery')) try { const res = await ctx.http.axios(trimSlash(config.endpoint) + OPTIONS_ENDPOINT) vae = res.data.sd_vae model = res.data.sd_model_checkpoint } catch (err) { session.send(session.text('.queryErr')) throw err } } async function getModelList() { try { const res = await ctx.http.axios(trimSlash(config.endpoint) + MODELS_ENDPOINT) res.data.forEach(item => models.push(item.title)) } catch (err) { session.send(session.text('.queryErr')) throw err } } async function sendInfo() { await getInfo() send( `当前模型为：${model}\n` + `当前 VAE 为：${vae}\n` + `输入 1 切换模型，输入 2 切换 VAE：` ) } async function sendModels() { await getModelList() send('当前可用模型有：\n' + models.map((model, i) => `${i + 1}.${model}\n`).join('') + '\n回复模型序号切换模型') } async function sendVaes() { await getModelList() send('当前可用 VAE 有：\n' + config.vaeList.map((vae, i) => `${i + 1}.${vae}\n`).join('') + '\n回复模型序号切换 VAE') } async function switchModel(index: number) { try { session.send(session.text('.switching')) await ctx.http.axios(trimSlash(config.endpoint) + OPTIONS_ENDPOINT, { method: 'POST', data: { sd_model_checkpoint: models[index - 1] } }) } catch (err) { session.send(session.text('.switchFailed')) throw err } await getInfo() send(`已切换至模型：${model}`) } async function switchVae(index: number) { try { session.send(session.text('.switching')) await ctx.http.axios(trimSlash(config.endpoint) + OPTIONS_ENDPOINT, { method: 'POST', data: { sd_vae: config.vaeList[index - 1] } }) } catch (err) { session.send(session.text('.switchFailed')) throw err } await getInfo() send(`已切换至 VAE：${vae}`) } await sendInfo() const choice = await input(2) switch (choice) { case 1: await sendModels() const modelIndex = await input(models.length) await switchModel(modelIndex) break case 2: await sendVaes() const vaeIndex = await input(config.vaeList.length) await switchVae(vaeIndex) break } }) }