add puppeteer render

MirrorCY · Nov 1, 2023 · 6aa1afd · 6aa1afd
1 parent b7fafc9
commit 6aa1afd
Show file tree

Hide file tree

Showing 4 changed files with 90 additions and 13 deletions.
diff --git a/node_modules/.bin/koishi b/node_modules/.bin/koishi
diff --git a/node_modules/.bin/koishi.cmd b/node_modules/.bin/koishi.cmd
diff --git a/src/index.ts b/src/index.ts
@@ -1,20 +1,34 @@
 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<string>
+  sendAsImage?: boolean
+  picWidth?: number
 }
 
-export const Config: Schema<Config> = 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 文件名。')
-})
+export const Config: Schema<Config> = 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)
 
@@ -30,6 +44,46 @@ export function apply(ctx: Context, config: Config) {
       let vae = ''
       let model = ''
 
+      async function render(content: string, picWidth: number) {
+        // https://github.com/ifrvn/koishi-plugin-send-as-image
+        return ctx.puppeteer.render(
+          `<html>
+        <head>
+          <style>
+            @font-face {
+              font-family: AlibabaPuHuiTi-2-55-Regular;
+              src:url(https://puhuiti.oss-cn-hangzhou.aliyuncs.com/AlibabaPuHuiTi-2/AlibabaPuHuiTi-2-55-Regular/AlibabaPuHuiTi-2-55-Regular.woff2) format('woff2');
+            }
+            html {
+              font-family: 'AlibabaPuHuiTi-2-55-Regular', 'Microsoft YaHei', 'Noto Sans SC', sans-serif;
+              width: ${picWidth}px;
+              height: 0;
+              background: white;
+            }
+            p {
+              padding: 10px;
+              word-wrap: break-word;
+              white-space: pre-wrap;
+            }
+          </style>
+        </head>
+        <body>
+          <p>${content}</p>
+        </body>
+        </html>`
+        )
+      }
+
+      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) {
@@ -63,7 +117,7 @@ export function apply(ctx: Context, config: Config) {
 
       async function sendInfo() {
         await getInfo()
-        session.send(
+        send(
           `当前模型为：${model}\n` +
           `当前 VAE 为：${vae}\n` +
           `输入 1 切换模型，输入 2 切换 VAE：`
@@ -72,14 +126,14 @@ export function apply(ctx: Context, config: Config) {
 
       async function sendModels() {
         await getModelList()
-        session.send('当前可用模型有：\n' +
+        send('当前可用模型有：\n' +
           models.map((model, i) => `${i + 1}.${model}\n`).join('') +
           '\n回复模型序号切换模型')
       }
 
       async function sendVaes() {
         await getModelList()
-        session.send('当前可用 VAE 有：\n' +
+        send('当前可用 VAE 有：\n' +
           config.vaeList.map((vae, i) => `${i + 1}.${vae}\n`).join('') +
           '\n回复模型序号切换 VAE')
       }
@@ -96,7 +150,7 @@ export function apply(ctx: Context, config: Config) {
           throw err
         }
         await getInfo()
-        session.send(`已切换至模型：${model}`)
+        send(`已切换至模型：${model}`)
       }
 
       async function switchVae(index: number) {
@@ -111,7 +165,7 @@ export function apply(ctx: Context, config: Config) {
           throw err
         }
         await getInfo()
-        session.send(`已切换至 VAE：${vae}`)
+        send(`已切换至 VAE：${vae}`)
       }
 
       await sendInfo()

diff --git a/src/locales/zh.yml b/src/locales/zh.yml
@@ -7,7 +7,8 @@ commands:
 
     messages:
       switching: 模型切换中
-      inputErr: 你输入了啥？
+      inputErr: 输入错误或超时
+      pptrErr: 未启用 puppeteer 插件，无法使用图像渲染
       inQuery: 查询模型中
       queryErr: 查询失败
       modelSucceed: 切换成功{0}