fix Bug: novelai队列429错误 #217

koishijs · Feb 8, 2024 · b6498fc · b6498fc
1 parent dcfda09
commit b6498fc
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git 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<Config>
 

diff --git 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)
             }