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

feat: experimentally support stable horde (#145)

Browse Source
This commit is contained in:
Maiko Sinkyaet Tan 2023-01-05 14:10:28 +08:00 committed by GitHub
parent e09bbc26ca
commit a69241c55b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 250 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
build/fetch-horde-models.js Normal file
View File

@ -0,0 +1,21 @@
const fsp = require('fs/promises')
const https = require('https')
const path = require('path')
const MODELS_URL = 'https://stablehorde.net/api/v2/status/models'
const DATA_JSON_PATH = path.join(__dirname, '..', 'data', 'horde-models.json')
;(async () => {
const db = await new Promise((resolve, reject) => {
https.get(MODELS_URL, res => {
let data = ''
res.on('data', chunk => data += chunk)
res.on('end', () => resolve(JSON.parse(data)))
}).on('error', reject)
})
const models = db.map((model) => model.name)
const json = JSON.stringify(models, null, 2)
await fsp.writeFile(DATA_JSON_PATH, json)
})()

85
data/horde-models.json Normal file
View File

@ -0,0 +1,85 @@
[
"colorbook",
"Papercutcraft",
"Xynthii-Diffusion",
"Redshift Diffusion",
"Mega Merge Diffusion",
"Borderlands",
"Voxel Art Diffusion",
"Elden Ring Diffusion",
"Cyberpunk Anime Diffusion",
"Future Diffusion",
"Microworlds",
"Vintedois Diffusion",
"mo-di-diffusion",
"Min Illust Background",
"Robo-Diffusion",
"stable_diffusion_inpainting",
"Dawgsmix",
"Zeipher Female Model",
"AIO Pixel Art",
"Dungeons and Diffusion",
"trinart",
"Yiffy",
"Midjourney PaintArt",
"App Icon Diffusion",
"Smoke Diffusion",
"Inkpunk Diffusion",
"Supermarionation",
"DnD Item",
"Zack3D",
"Fantasy Card Diffusion",
"stable_diffusion",
"Synthwave",
"Furry Epoch",
"Asim Simpsons",
"PortraitPlus",
"Anything Diffusion",
"Tron Legacy Diffusion",
"Clazy",
"Valorant Diffusion",
"Stable Diffusion 2 Depth",
"Papercut Diffusion",
"ChromaV5",
"stable_diffusion_2.1",
"Archer Diffusion",
"Dark Victorian Diffusion",
"Guohua Diffusion",
"Analog Diffusion",
"vectorartz",
"Classic Animation Diffusion",
"Eimis Anime Diffusion",
"Ranma Diffusion",
"Microscopic",
"Dreamlike Diffusion",
"GTA5 Artwork Diffusion",
"Midjourney Diffusion",
"Poison",
"kurzgesagt",
"Samdoesarts Ultmerge",
"Dreamlike Photoreal",
"Trinart Characters",
"Arcane Diffusion",
"JWST Deep Space Diffusion",
"RPG",
"Hassanblend",
"waifu_diffusion",
"ModernArt Diffusion",
"Darkest Diffusion",
"Balloon Art",
"Comic-Diffusion",
"BubblyDubbly",
"Eternos",
"Van Gogh Diffusion",
"Double Exposure Diffusion",
"Squishmallow Diffusion",
"Funko Diffusion",
"Hentai Diffusion",
"Ghibli Diffusion",
"Seek.art MEGA",
"Knollingcase",
"Spider-Verse Diffusion",
"ACertainThing",
"Wavyfusion",
"Nitro Diffusion"
]

3
package.json
View File

@ -6,7 +6,8 @@
"typings": "lib/index.d.ts",
"files": [
"lib",
"dist"
"dist",
"data"
],
"browser": {
"image-size": false,

53
src/config.ts
View File

@ -13,6 +13,8 @@ export const orientMap = {
square: { height: 640, width: 640 },
} as const
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',
@ -54,6 +56,29 @@ export namespace sampler {
'plms': 'PLMS',
}
export const horde = {
k_lms: 'LMS',
k_heun: 'Heun',
k_euler: 'Euler',
k_euler_a: 'Euler a',
k_dpm_2: 'DPM2',
k_dpm_2_a: 'DPM2 a',
k_dpm_fast: 'DPM fast',
k_dpm_adaptive: 'DPM adaptive',
k_dpmpp_2m: 'DPM++ 2M',
k_dpmpp_2s_a: 'DPM++ 2S a',
k_lms_ka: 'LMS Karras',
k_heun_ka: 'Heun Karras',
k_euler_ka: 'Euler Karras',
k_euler_a_ka: 'Euler a Karras',
k_dpm_2_ka: 'DPM2 Karras',
k_dpm_2_a_ka: 'DPM2 a Karras',
k_dpm_fast_ka: 'DPM fast Karras',
k_dpm_adaptive_ka: 'DPM adaptive Karras',
k_dpmpp_2m_ka: 'DPM++ 2M Karras',
k_dpmpp_2s_a_ka: 'DPM++ 2S a Karras',
}
export function createSchema(map: Dict<string>) {
return Schema.union(Object.entries(map).map(([key, value]) => {
return Schema.const(key).description(value)
@ -135,7 +160,7 @@ interface ParamConfig {
}
export interface Config extends PromptConfig, ParamConfig {
type: 'token' | 'login' | 'naifu' | 'sd-webui'
type: 'token' | 'login' | 'naifu' | 'sd-webui' | 'stable-horde'
token?: string
email?: string
password?: string
@ -143,11 +168,14 @@ export interface Config extends PromptConfig, ParamConfig {
allowAnlas?: boolean | number
endpoint?: string
headers?: Dict<string>
nsfw?: 'disallow' | 'censor' | 'allow'
maxIterations?: number
maxRetryCount?: number
requestTimeout?: number
recallTimeout?: number
maxConcurrency?: number
pollInterval?: number
trustedWorkerOnly?: boolean
}
export const Config = Schema.intersect([
@ -157,6 +185,7 @@ export const Config = Schema.intersect([
...process.env.KOISHI_ENV === 'browser' ? [] : [Schema.const('login' as const).description('账号密码')],
Schema.const('naifu' as const).description('naifu'),
Schema.const('sd-webui' as const).description('sd-webui'),
Schema.const('stable-horde' as const).description('Stable Horde'),
] as const).description('登录方式'),
}).description('登录设置'),
@ -192,6 +221,18 @@ export const Config = Schema.intersect([
endpoint: Schema.string().description('API 服务器地址。').required(),
headers: Schema.dict(String).description('要附加的额外请求头。'),
}),
Schema.object({
type: Schema.const('stable-horde'),
endpoint: Schema.string().description('API 服务器地址。').default('https://stablehorde.net/'),
token: Schema.string().description('授权令牌 (API Key)。').role('secret').default('0000000000'),
nsfw: Schema.union([
Schema.const('disallow').description('禁止'),
Schema.const('censor').description('屏蔽'),
Schema.const('allow').description('允许'),
]).description('是否允许 NSFW 内容。').default('allow'),
trustedWorkerOnly: Schema.boolean().description('是否只请求可信任工作节点。').default(false),
pollInterval: Schema.number().role('time').description('轮询进度间隔时长。').default(Time.second),
}),
]),
Schema.union([
@ -201,6 +242,11 @@ export const Config = Schema.intersect([
upscaler: Schema.union(upscalers).description('默认的放大算法。').default('Lanczos'),
hiresFix: Schema.boolean().description('是否启用高分辨率修复。').default(false),
}).description('参数设置'),
Schema.object({
type: Schema.const('stable-horde'),
sampler: sampler.createSchema(sampler.horde),
model: Schema.union(hordeModels),
}),
Schema.object({
type: Schema.const('naifu'),
sampler: sampler.createSchema(sampler.nai),
@ -236,11 +282,6 @@ export const Config = Schema.intersect([
Schema.const('default').description('发送图片和关键信息'),
Schema.const('verbose').description('发送全部信息'),
]).description('输出方式。').default('default'),
allowAnlas: Schema.union([
Schema.const(true).description('允许'),
Schema.const(false).description('禁止'),
Schema.natural().description('权限等级').default(1),
]).default(true).description('是否启用高级功能 (例如图片增强和手动设置某些参数)。'),
maxIterations: Schema.natural().description('允许的最大绘制次数。').default(1),
maxRetryCount: Schema.natural().description('连接失败时最大的重试次数。').default(3),
requestTimeout: Schema.number().role('time').description('当请求超过这个时间时会中止并提示超时。').default(Time.minute),

126
src/index.ts
View File

@ -247,6 +247,8 @@ export function apply(ctx: Context, config: Config) {
switch (config.type) {
case 'sd-webui':
return image ? '/sdapi/v1/img2img' : '/sdapi/v1/txt2img'
case 'stable-horde':
return '/api/v2/generate/async'
case 'naifu':
return '/generate-stream'
default:
@ -255,49 +257,100 @@ export function apply(ctx: Context, config: Config) {
})()
const getPayload = () => {
if (config.type !== 'sd-webui') {
parameters.sampler = sampler.sd2nai(options.sampler)
parameters.image = image?.base64 // NovelAI / NAIFU accepts bare base64 encoded image
if (config.type === 'naifu') return parameters
return { model, input: prompt, parameters: omit(parameters, ['prompt']) }
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
if (config.type === 'naifu') return parameters
return { model, input: prompt, parameters: omit(parameters, ['prompt']) }
}
case 'sd-webui': {
return {
sampler_index: sampler.sd[options.sampler],
init_images: image && [image.dataUrl], // sd-webui accepts data URLs with base64 encoded image
...project(parameters, {
prompt: 'prompt',
batch_size: 'n_samples',
seed: 'seed',
negative_prompt: 'uc',
cfg_scale: 'scale',
steps: 'steps',
width: 'width',
height: 'height',
denoising_strength: 'strength',
}),
}
}
case 'stable-horde': {
return {
prompt: parameters.prompt,
params: {
sampler_name: options.sampler.replace('_ka', ''),
cfg_scale: parameters.scale,
denoising_strength: parameters.strength,
seed: parameters.seed.toString(),
height: parameters.height,
width: parameters.width,
post_processing: [],
karras: options.sampler.includes('_ka'),
steps: parameters.steps,
n: 1,
},
nsfw: config.nsfw !== 'disallow',
trusted_workers: config.trustedWorkerOnly,
censor_nsfw: config.nsfw === 'censor',
models: [options.model],
source_image: image?.base64,
source_processing: image ? 'img2img' : undefined,
}
}
}
}
return {
sampler_index: sampler.sd[options.sampler],
init_images: image && [image.dataUrl], // sd-webui accepts data URLs with base64 encoded image
enable_hr: options.hiresFix ?? config.hiresFix ?? false,
...project(parameters, {
prompt: 'prompt',
batch_size: 'n_samples',
seed: 'seed',
negative_prompt: 'uc',
cfg_scale: 'scale',
steps: 'steps',
width: 'width',
height: 'height',
denoising_strength: 'strength',
}),
const getHeaders = () => {
switch (config.type) {
case 'login':
case 'token':
case 'naifu':
return { Authorization: `Bearer ${token}` }
case 'stable-horde':
return { apikey: token }
}
}
const iterate = async () => {
const request = () => ctx.http.axios(trimSlash(config.endpoint) + path, {
method: 'POST',
timeout: config.requestTimeout,
headers: {
...config.headers,
authorization: 'Bearer ' + token,
},
data: getPayload(),
}).then((res) => {
const request = async () => {
const res = await ctx.http.axios(trimSlash(config.endpoint) + path, {
method: 'POST',
timeout: config.requestTimeout,
headers: {
...config.headers,
...getHeaders(),
},
data: getPayload(),
})
if (config.type === 'sd-webui') {
return stripDataPrefix((res.data as StableDiffusionWebUI.Response).images[0])
}
if (config.type === 'stable-horde') {
const uuid = res.data.id
const check = () => ctx.http.get(trimSlash(config.endpoint) + '/api/v2/generate/check/' + uuid).then((res) => res.done)
const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms))
while(await check() === false) {
await sleep(config.pollInterval)
}
const result = await ctx.http.get(trimSlash(config.endpoint) + '/api/v2/generate/status/' + uuid)
return result.generations[0].img
}
// event: newImage
// id: 1
// data:
return res.data?.slice(27)
})
}
let base64: string, count = 0
while (true) {
@ -373,11 +426,22 @@ export function apply(ctx: Context, config: Config) {
})
ctx.accept(['scale', 'model', 'sampler', 'output'], (config) => {
const getSamplers = () => {
switch (config.type) {
case 'sd-webui':
return sampler.sd
case 'stable-horde':
return sampler.horde
default:
return sampler.nai
}
}
cmd._options.output.fallback = config.output
cmd._options.scale.fallback = config.scale
cmd._options.model.fallback = config.model
cmd._options.sampler.fallback = config.sampler
cmd._options.sampler.type = Object.keys(config.type === 'sd-webui' ? sampler.sd : sampler.nai)
cmd._options.sampler.type = Object.keys(getSamplers())
}, { immediate: true })
const subcmd = ctx.intersect(() => config.type === 'sd-webui')