WindyMadman/novelai-bot
1
0
Fork 0
mirror of https://github.com/koishijs/novelai-bot synced 2025-01-07 07:06:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: Compatible with nai-diffusion-3 (#213)

Browse Source 
Co-authored-by: Shigma <shigma10826@gmail.com>
Co-authored-by: Maiko Sinkyaet Tan <maiko.tan.coding@gmail.com>
This commit is contained in:
风宝宝 2023-11-25 23:04:25 +08:00 committed by GitHub
parent 052e57ac3d
commit 97a98eb52f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 25 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
package.json
View File

@ -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"
}

8
src/config.ts
View File

@ -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

18
src/index.ts
View File

@ -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))
}