From b6498fce380a997533a8cbdbe1239a161b44e993 Mon Sep 17 00:00:00 2001 From: 42 <37006258+MirrorCY@users.noreply.github.com> Date: Thu, 8 Feb 2024 17:54:43 +0000 Subject: [PATCH 1/4] =?UTF-8?q?fix=20Bug:=20novelai=E9=98=9F=E5=88=97429?= =?UTF-8?q?=E9=94=99=E8=AF=AF=20#217?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/config.ts | 2 ++ src/index.ts | 9 +++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/src/config.ts b/src/config.ts index 5f87954..17cf1b4 100644 --- a/src/config.ts +++ b/src/config.ts @@ -237,6 +237,7 @@ export interface Config extends PromptConfig, ParamConfig { requestTimeout?: number recallTimeout?: number maxConcurrency?: number + globalConcurrency?: number pollInterval?: number trustedWorkers?: boolean } @@ -389,6 +390,7 @@ export const Config = Schema.intersect([ requestTimeout: Schema.number().role('time').description('当请求超过这个时间时会中止并提示超时。').default(Time.minute), recallTimeout: Schema.number().role('time').description('图片发送后自动撤回的时间 (设置为 0 以禁用此功能)。').default(0), maxConcurrency: Schema.number().description('单个频道下的最大并发数量 (设置为 0 以禁用此功能)。').default(0), + globalConcurrency: Schema.number().min(0).description('全局的最大并发数量 (设置为 0 以禁用此功能)。').default(0), }).description('高级设置'), ]) as Schema diff --git a/src/index.ts b/src/index.ts index 74ee3a3..01d4c51 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,4 +1,4 @@ -import { Computed, Context, Dict, h, Logger, omit, Quester, Session, SessionError, trimSlash } from 'koishi' +import { Computed, Context, Dict, h, Logger, omit, Quester, Session, SessionError, sleep, trimSlash } from 'koishi' import { Config, modelMap, models, orientMap, parseInput, sampler, upscalers, scheduler } from './config' import { ImageData, StableDiffusionWebUI } from './types' import { closestMultiple, download, forceDataPrefix, getImageSize, login, NetworkError, project, resizeInput, Size } from './utils' @@ -285,6 +285,12 @@ export function apply(ctx: Context, config: Config) { session.send(globalTasks.size ? session.text('.pending', [globalTasks.size]) : session.text('.waiting')) + + if (config.globalConcurrency) { + while (globalTasks.size >= config.globalConcurrency) { + await sleep(100) + } + } container.forEach((id) => globalTasks.add(id)) const cleanUp = (id: string) => { @@ -428,7 +434,6 @@ export function apply(ctx: Context, config: Config) { 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) } From 572a2fa8b1406d00c2ce08ed78330d000cd98bd8 Mon Sep 17 00:00:00 2001 From: Maiko Tan Date: Thu, 15 Aug 2024 10:31:30 +0800 Subject: [PATCH 2/4] refa: use event for triggering next iteration --- src/index.ts | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/src/index.ts b/src/index.ts index 97cb9e1..1446696 100644 --- a/src/index.ts +++ b/src/index.ts @@ -10,6 +10,12 @@ import { readFile } from 'fs/promises' export * from './config' +declare module 'koishi' { + interface Events { + 'novelai/finish'(id: string): void + } +} + export const reactive = true export const name = 'novelai' @@ -45,6 +51,7 @@ export function apply(ctx: Context, config: Config) { const tasks: Dict> = Object.create(null) const globalTasks = new Set() + const globalPending = new Set() let tokenTask: Promise = null const getToken = () => tokenTask ||= login(ctx) @@ -297,10 +304,15 @@ export function apply(ctx: Context, config: Config) { session.send(globalTasks.size ? session.text('.pending', [globalTasks.size]) : session.text('.waiting')) - + if (config.globalConcurrency) { while (globalTasks.size >= config.globalConcurrency) { - await sleep(100) + globalPending.add(container.pop()) + await new Promise((resolve) => ctx.once('novelai/finish', (id) => { + if (id === id) { + resolve() + } + })) } } @@ -308,6 +320,11 @@ export function apply(ctx: Context, config: Config) { const cleanUp = (id: string) => { tasks[session.cid]?.delete(id) globalTasks.delete(id) + if (globalPending.size) { + const id = globalPending.values().next().value + globalPending.delete(id) + ctx.parallel('novelai/finish', id) + } } const path = (() => { @@ -424,8 +441,8 @@ export function apply(ctx: Context, config: Config) { const body = new FormData() const capture = /^data:([\w/.+-]+);base64,(.*)$/.exec(image.dataUrl) const [, mime,] = capture - - let name = Date.now().toString() + + let name = Date.now().toString() const ext = mime === 'image/jpeg' ? 'jpg' : mime === 'image/png' ? 'png' : '' if (ext) name += `.${ext}` const imageFile = new Blob([image.buffer], {type:mime}) @@ -441,7 +458,7 @@ export function apply(ctx: Context, config: Config) { const data = res.data let imagePath = data.name if (data.subfolder) imagePath = data.subfolder + '/' + imagePath - + for (const nodeId in prompt) { if (prompt[nodeId].class_type === 'LoadImage') { prompt[nodeId].inputs.image = imagePath @@ -466,7 +483,7 @@ export function apply(ctx: Context, config: Config) { const negativeeNodeId = prompt[nodeId].inputs.negative[0] const latentImageNodeId = prompt[nodeId].inputs.latent_image[0] prompt[positiveNodeId].inputs.text = parameters.prompt - prompt[negativeeNodeId].inputs.text = parameters.uc + prompt[negativeeNodeId].inputs.text = parameters.uc prompt[latentImageNodeId].inputs.width = parameters.width prompt[latentImageNodeId].inputs.height = parameters.height prompt[latentImageNodeId].inputs.batch_size = parameters.n_samples From a6bd059670069229960a858fcd6f89dd4289ce6b Mon Sep 17 00:00:00 2001 From: Maiko Tan Date: Thu, 15 Aug 2024 10:39:30 +0800 Subject: [PATCH 3/4] fix: pending id error --- src/index.ts | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/index.ts b/src/index.ts index 1446696..a798b67 100644 --- a/src/index.ts +++ b/src/index.ts @@ -306,10 +306,11 @@ export function apply(ctx: Context, config: Config) { : session.text('.waiting')) if (config.globalConcurrency) { - while (globalTasks.size >= config.globalConcurrency) { - globalPending.add(container.pop()) + if (globalTasks.size >= config.globalConcurrency) { + const pendingId = container.pop() + globalPending.add(pendingId) await new Promise((resolve) => ctx.once('novelai/finish', (id) => { - if (id === id) { + if (id === pendingId) { resolve() } })) From 7b2aa62fbe4976be00e9237cc82e8df542890f44 Mon Sep 17 00:00:00 2001 From: Maiko Sinkyaet Tan Date: Fri, 16 Aug 2024 00:43:09 +0800 Subject: [PATCH 4/4] Update src/index.ts Co-authored-by: Shigma --- src/index.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/index.ts b/src/index.ts index a798b67..74cbacf 100644 --- a/src/index.ts +++ b/src/index.ts @@ -309,9 +309,11 @@ export function apply(ctx: Context, config: Config) { if (globalTasks.size >= config.globalConcurrency) { const pendingId = container.pop() globalPending.add(pendingId) - await new Promise((resolve) => ctx.once('novelai/finish', (id) => { - if (id === pendingId) { + await new Promise((resolve) => { + const dispose = ctx.on('novelai/finish', (id) => { + if (id !== pendingId) return resolve() + dispose() } })) }