koishi-plugin-hhs-huatu 0.3.2 → 0.3.3

package/lib/index.js

@@ -471,9 +471,6 @@ async function login(ctx, email, password) {
       let startIndex = ctx.runtime.currentTokenIndex;
       let strategy = ctx.config.tokenStrategy || "round-robin";
       logger.debug(`使用token策略: ${strategy}`);
-      if (strategy === "parallel" && !ctx.runtime.tokenUsage) {
-        ctx.runtime.tokenUsage = {};
-      }
       if (strategy === "random") {
         startIndex = Math.floor(Math.random() * ctx.config.token.length);
         ctx.runtime.currentTokenIndex = startIndex;
@@ -481,9 +478,9 @@ async function login(ctx, email, password) {
         let foundUnusedToken = false;
         for (let i = 0; i < ctx.config.token.length; i++) {
           const checkIndex = (ctx.runtime.currentTokenIndex + i) % ctx.config.token.length;
-          if (!ctx.runtime.tokenUsage[checkIndex]) {
+          const tokenUsage = ctx.runtime.tokenUsage || {};
+          if (!tokenUsage[checkIndex]) {
             startIndex = checkIndex;
-            ctx.runtime.tokenUsage[checkIndex] = true;
             ctx.runtime.currentTokenIndex = (checkIndex + 1) % ctx.config.token.length;
             foundUnusedToken = true;
             logger.debug(`并行策略选择token索引 ${checkIndex}，下一次从索引 ${ctx.runtime.currentTokenIndex} 开始检查`);
@@ -1523,12 +1520,19 @@ __name(auditImage, "auditImage");
 
 // src/queueSystem.ts
 var QueueSystem = class {
-  constructor(ctx, config, generateImageFn, membershipSystem) {
+  constructor(ctx, config, generateImageFn, membershipSystem, initialTokenUsage) {
     this.ctx = ctx;
     this.config = config;
     this.generateImageFn = generateImageFn;
     this.membershipSystem = membershipSystem || null;
     this.maxConcurrentTasks = config.tokenStrategy === "parallel" && Array.isArray(config.token) ? Math.min(config.maxConcurrentRequests || 1, config.token.length) : config.maxConcurrentRequests || 1;
+    if (initialTokenUsage) {
+      this.tokenUsage = { ...initialTokenUsage };
+    } else if (config.tokenStrategy === "parallel" && Array.isArray(config.token)) {
+      for (let i = 0; i < config.token.length; i++) {
+        this.tokenUsage[i] = false;
+      }
+    }
   }
   static {
     __name(this, "QueueSystem");
@@ -1552,6 +1556,9 @@ var QueueSystem = class {
   maxConcurrentTasks;
   // 会员系统引用
   membershipSystem = null;
+  // Token管理相关
+  tokenLock = Promise.resolve();
+  tokenUsage = /* @__PURE__ */ Object.create(null);
   // 获取锁
   async acquireRedrawLock() {
     if (!this.redrawLock) {
@@ -1589,6 +1596,20 @@ var QueueSystem = class {
     this.usedTokenIndices.add(newIndex);
     return newIndex;
   }
+  // 获取token锁，确保token选择的原子性
+  async acquireTokenLock() {
+    const release = this.tokenLock;
+    let releaseFn = null;
+    this.tokenLock = new Promise((resolve3) => {
+      releaseFn = resolve3;
+    });
+    await release;
+    return () => {
+      if (releaseFn) {
+        releaseFn();
+      }
+    };
+  }
   // 处理队列
   async processQueue() {
     if (this.taskQueue.length === 0) return;
@@ -1596,10 +1617,8 @@ var QueueSystem = class {
     const tasksToProcess = Math.min(this.maxConcurrentTasks - this.processingTasks, this.taskQueue.length);
     if (this.config.tokenStrategy === "parallel" && Array.isArray(this.ctx.config.token)) {
       this.ctx.logger.debug(`处理队列中的任务，当前处理中: ${this.processingTasks}，即将处理: ${tasksToProcess}，队列剩余: ${this.taskQueue.length}，最大并发: ${this.maxConcurrentTasks}`);
-      if (this.ctx.runtime?.tokenUsage) {
-        const usedTokens = Object.entries(this.ctx.runtime.tokenUsage).filter(([_, used]) => used).map(([index]) => index);
-        this.ctx.logger.debug(`当前token使用情况: 已用 ${usedTokens.length}/${this.ctx.config.token.length} (索引: ${usedTokens.join(", ")})`);
-      }
+      const usedTokens = Object.entries(this.tokenUsage).filter(([_, used]) => used).map(([index]) => index);
+      this.ctx.logger.debug(`当前token使用情况: 已用 ${usedTokens.length}/${this.ctx.config.token.length} (索引: ${usedTokens.join(", ")})`);
     }
     for (let i = 0; i < tasksToProcess; i++) {
       this.processingTasks++;
@@ -1625,6 +1644,7 @@ var QueueSystem = class {
         }
       }
       Promise.resolve().then(async () => {
+        const releaseTokenLock = await this.acquireTokenLock();
         try {
           const result = await this.generateImageFn(task.session, task.options, task.input);
           task.resolve(result);
@@ -1632,9 +1652,10 @@ var QueueSystem = class {
           task.reject(err);
         } finally {
           this.processingTasks--;
+          releaseTokenLock();
           if (Array.isArray(this.ctx.config.token)) {
-            if (this.config.tokenStrategy === "parallel" && this.ctx.runtime?.tokenUsage) {
-              const usedTokens = Object.entries(this.ctx.runtime.tokenUsage).filter(([_, used]) => used).map(([index]) => index);
+            if (this.config.tokenStrategy === "parallel") {
+              const usedTokens = Object.entries(this.tokenUsage).filter(([_, used]) => used).map(([index]) => index);
               this.ctx.logger.debug(`任务完成后token使用情况(parallel): 已用 ${usedTokens.length}/${this.ctx.config.token.length} (索引: ${usedTokens.join(", ")})`);
             } else if (this.config.tokenStrategy === "round-robin") {
               this.ctx.logger.debug(`任务完成后token索引(round-robin): ${this.ctx.runtime?.currentTokenIndex}/${this.ctx.config.token.length - 1}`);
@@ -2168,6 +2189,16 @@ function apply(ctx, config) {
   const tasks = /* @__PURE__ */ Object.create(null);
   const globalTasks = /* @__PURE__ */ new Set();
   let queueSystem;
+  const syncTokenUsage = /* @__PURE__ */ __name(() => {
+    if (ctx.runtime && Array.isArray(ctx.config.token)) {
+      if (!ctx.runtime.tokenUsage) {
+        ctx.runtime.tokenUsage = {};
+        for (let i = 0; i < ctx.config.token.length; i++) {
+          ctx.runtime.tokenUsage[i] = false;
+        }
+      }
+    }
+  }, "syncTokenUsage");
   let tokenTask = null;
   const getToken = /* @__PURE__ */ __name(async (session) => {
     try {
@@ -2284,6 +2315,7 @@ function apply(ctx, config) {
     }
     return { width, height, custom: true };
   }, "resolution");
+  syncTokenUsage();
   async function generateImage(session, options2, input) {
     ctx.logger.debug(`generateImage开始处理，sessionId=${session.id}，userId=${session.userId}`);
     const sessionObj = session;
@@ -2940,7 +2972,7 @@ function apply(ctx, config) {
     }
   }
   __name(generateImage, "generateImage");
-  queueSystem = new QueueSystem(ctx, config, generateImage, membershipSystem);
+  queueSystem = new QueueSystem(ctx, config, generateImage, membershipSystem, ctx.runtime?.tokenUsage);
   const cmd = ctx.command("novelai [prompts...]").alias("nai").alias("imagine").alias("nai4", { options: { model: "nai-v4-full", sampler: "k_euler_a", iterations: 1, batch: 1 } }).alias("nai4c", { options: { model: "nai-v4-curated-preview", sampler: "k_euler_a", iterations: 1, batch: 1 } }).alias("nai4-5c", { options: { model: "nai-v4-5-curated", sampler: "k_euler_a", iterations: 1, batch: 1 } }).alias("nai4-5", { options: { model: "nai-v4-5-full", sampler: "k_euler_a", iterations: 1, batch: 1 } }).userFields(["authority"]).shortcut("imagine", { i18n: true, fuzzy: true }).shortcut("enhance", { i18n: true, fuzzy: true, options: { enhance: true } }).option("enhance", "-e", { hidden: some(restricted, thirdParty, noImage) }).option("model", "-m <model>", { type: models, hidden: thirdParty }).option("resolution", "-r <resolution>", { type: resolution2 }).option("output", "-o", { type: ["minimal", "default", "verbose"] }).option("override", "-O", { hidden: restricted }).option("sampler", "-s <sampler>").option("seed", "-x <seed:number>").option("steps", "-t <step>", { type: step, hidden: restricted }).option("scale", "-c <scale:number>").option("rescale", "-R <rescale:number>").option("noise", "-n <noise:number>", { hidden: some(restricted, thirdParty) }).option("strength", "-N <strength:number>", { hidden: restricted }).option("hiresFix", "-H", { hidden: /* @__PURE__ */ __name(() => config.type !== "sd-webui", "hidden") }).option("hiresFixSteps", "<step>", { type: step, hidden: /* @__PURE__ */ __name(() => config.type !== "sd-webui", "hidden") }).option("smea", "-S", { hidden: /* @__PURE__ */ __name(() => config.model !== "nai-v3", "hidden") }).option("smeaDyn", "-d", { hidden: /* @__PURE__ */ __name(() => config.model !== "nai-v3", "hidden") }).option("scheduler", "-C <scheduler:string>", {
     hidden: /* @__PURE__ */ __name(() => config.type === "naifu", "hidden"),
     type: ["token", "login"].includes(config.type) ? scheduler.nai : config.type === "sd-webui" ? scheduler.sd : config.type === "stable-horde" ? scheduler.horde : []
@@ -3247,8 +3279,9 @@ function apply(ctx, config) {
               queueSystem.userTasks[currentUserId]--;
               targetBot.sendMessage(targetChannelId, value);
               if (Array.isArray(ctx.config.token)) {
-                if (config.tokenStrategy === "parallel" && ctx.runtime?.tokenUsage) {
-                  const usedTokens = Object.entries(ctx.runtime.tokenUsage).filter(([_, used]) => used).map(([index2]) => index2);
+                if (config.tokenStrategy === "parallel") {
+                  const queueTokenUsage = queueSystem.tokenUsage || {};
+                  const usedTokens = Object.entries(queueTokenUsage).filter(([_, used]) => used).map(([index2]) => index2);
                   ctx.logger.debug(`重画任务 ${taskUniqueId} 完成，当前token使用情况: ${usedTokens.join(", ")}`);
                 } else if (config.tokenStrategy === "round-robin") {
                   ctx.logger.debug(`重画任务 ${taskUniqueId} 完成，当前token索引: ${ctx.runtime?.currentTokenIndex}`);
@@ -3266,8 +3299,9 @@ function apply(ctx, config) {
                 handleError(ctx, session, err)
               );
               if (Array.isArray(ctx.config.token)) {
-                if (config.tokenStrategy === "parallel" && ctx.runtime?.tokenUsage) {
-                  const usedTokens = Object.entries(ctx.runtime.tokenUsage).filter(([_, used]) => used).map(([index2]) => index2);
+                if (config.tokenStrategy === "parallel") {
+                  const queueTokenUsage = queueSystem.tokenUsage || {};
+                  const usedTokens = Object.entries(queueTokenUsage).filter(([_, used]) => used).map(([index2]) => index2);
                   ctx.logger.debug(`重画任务 ${taskUniqueId} 失败，当前token使用情况: ${usedTokens.join(", ")}`);
                 } else if (config.tokenStrategy === "round-robin") {
                   ctx.logger.debug(`重画任务 ${taskUniqueId} 失败，当前token索引: ${ctx.runtime?.currentTokenIndex}`);

package/lib/queueSystem.d.ts

@@ -28,10 +28,13 @@ export declare class QueueSystem {
     private redrawWaitQueue;
     maxConcurrentTasks: number;
     private membershipSystem;
-    constructor(ctx: Context, config: Config, generateImageFn: (session: Session<'authority'>, options: any, input: string) => Promise<any>, membershipSystem?: MembershipSystem);
+    private tokenLock;
+    private tokenUsage;
+    constructor(ctx: Context, config: Config, generateImageFn: (session: Session<'authority'>, options: any, input: string) => Promise<any>, membershipSystem?: MembershipSystem, initialTokenUsage?: Dict<boolean>);
     acquireRedrawLock(): Promise<void>;
     releaseRedrawLock(): void;
     getUniqueTokenIndex(currentIndex: number, tokenCount: number): number;
+    private acquireTokenLock;
     processQueue(): Promise<void>;
     addTask(task: TaskQueueItem): Promise<any>;
     getQueueStatus(userId: string): {

package/package.json

@@ -1,7 +1,7 @@
 {
   "name": "koishi-plugin-hhs-huatu",
   "description": "Generate images by diffusion models",
-  "version": "0.3.2",
+  "version": "0.3.3",
   "main": "lib/index.js",
   "typings": "lib/index.d.ts",
   "files": [