feat: Compatible with nai-diffusion-3

package.json

@@ -69,6 +69,7 @@
     "@koishijs/plugin-help": "^2.3.3",
     "@koishijs/translator": "^1.1.1",
     "@koishijs/vitepress": "^3.0.1",
+    "@types/adm-zip": "^0.5.5",
     "@types/libsodium-wrappers-sumo": "^0.7.7",
     "@types/node": "^20.8.9",
     "atsc": "^1.2.2",
@@ -78,6 +79,7 @@
     "vitepress": "1.0.0-rc.4"
   },
   "dependencies": {
+    "adm-zip": "^0.5.10",
     "image-size": "^1.0.2",
     "libsodium-wrappers-sumo": "^0.7.13"
   }

src/config.ts

@@ -9,6 +9,7 @@ export const modelMap = {
   safe: 'safe-diffusion',
   nai: 'nai-diffusion',
   furry: 'nai-diffusion-furry',
+  'nai-v3': 'nai-diffusion-3',
 } as const
 
 export const orientMap = {
@@ -20,9 +21,10 @@ export const orientMap = {
 export const hordeModels = require('../data/horde-models.json') as string[]
 
 const ucPreset = [
-  'nsfw, lowres, bad anatomy, bad hands, text, error, missing fingers',
-  'extra digit, fewer digits, cropped, worst quality, low quality',
-  'normal quality, jpeg artifacts, signature, watermark, username, blurry',
+  // Replace with the prompt words that come with novelai
+  'nsfw, lowres, {bad}, error, fewer, extra, missing, worst quality',
+  'jpeg artifacts, bad quality, watermark, unfinished, displeasing',
+  'chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract]'
 ].join(', ')
 
 type Model = keyof typeof modelMap
@@ -405,9 +407,10 @@ export function parseInput(session: Session, input: string, config: Config, over
       .split('\\{').map(s => s.replace(/\{/g, '(')).join('\\{')
       .split('\\}').map(s => s.replace(/\}/g, ')')).join('\\}')
   } else {
-    input = input
-      .split('\\(').map(s => s.replace(/\(/g, '{')).join('\\(')
-      .split('\\)').map(s => s.replace(/\)/g, '}')).join('\\)')
+    // emmmmm doesn't know why substituting curly brackets makes the picture ugly, so it's temporarily commented
+    // input = input
+    //   .split('\\(').map(s => s.replace(/\(/g, '{')).join('\\(')
+    //   .split('\\)').map(s => s.replace(/\)/g, '}')).join('\\)')
   }
 
   input = input

src/index.ts

@@ -4,6 +4,7 @@ import { ImageData, StableDiffusionWebUI } from './types'
 import { closestMultiple, download, forceDataPrefix, getImageSize, login, NetworkError, project, resizeInput, Size } from './utils'
 import {} from '@koishijs/translator'
 import {} from '@koishijs/plugin-help'
+import AdmZip from 'adm-zip'
 
 export * from './config'
 
@@ -302,13 +303,23 @@ export function apply(ctx: Context, config: Config) {
         }
       })()
 
+      const getResponseType = () => {
+        if (config.type === "stable-horde" || config.type === "sd-webui") {
+          return "json"
+        }
+        return "arraybuffer"
+      }
+
       const getPayload = () => {
         switch (config.type) {
           case 'login':
           case 'token':
           case 'naifu': {
             parameters.sampler = sampler.sd2nai(options.sampler)
             parameters.image = image?.base64 // NovelAI / NAIFU accepts bare base64 encoded image
+            // The latest interface changes uc to negative_prompt, so that needs to be changed here as well
+            parameters.negative_prompt = parameters.uc
+            delete parameters.uc
             if (config.type === 'naifu') return parameters
             return { model, input: prompt, parameters: omit(parameters, ['prompt']) }
           }
@@ -379,6 +390,8 @@ export function apply(ctx: Context, config: Config) {
           const res = await ctx.http.axios(trimSlash(config.endpoint) + path, {
             method: 'POST',
             timeout: config.requestTimeout,
+            // Since novelai's latest interface returns an application/x-zip-compressed, a responseType must be passed in
+            responseType: getResponseType(),
             headers: {
               ...config.headers,
               ...getHeaders(),
@@ -412,6 +425,18 @@ export function apply(ctx: Context, config: Config) {
           // event: newImage
           // id: 1
           // data:
+
+          if (res.headers['content-type'] === 'application/x-zip-compressed') {
+            const buffer = Buffer.from(res.data, 'binary');  // Ensure 'binary' encoding
+            const zip = new AdmZip(buffer);
+
+            // Gets all files in the ZIP file
+            const zipEntries = zip.getEntries();
+            const firstImageBuffer = zip.readFile(zipEntries[0]);
+            const b64 = Buffer.from(firstImageBuffer).toString('base64')
+            return forceDataPrefix(b64, "image/png")
+          }
+
           return forceDataPrefix(res.data?.slice(27))
         }