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

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'
 
@@ -310,6 +311,9 @@ export function apply(ctx: Context, config: Config) {
             parameters.sampler = sampler.sd2nai(options.sampler)
             parameters.image = image?.base64 // NovelAI / NAIFU accepts bare base64 encoded image
             if (config.type === 'naifu') return parameters
+            // 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
             return { model, input: prompt, parameters: omit(parameters, ['prompt']) }
           }
           case 'sd-webui': {
@@ -379,6 +383,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: ['login', 'token'].includes(config.type) ? 'arraybuffer' : 'json',
             headers: {
               ...config.headers,
               ...getHeaders(),
@@ -412,6 +418,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))
         }