Merge pull request #18 from erictik:zcpua/issue5

Upscale no msg return

example/upscale.ts

@@ -14,7 +14,7 @@ async function main() {
     <string>process.env.SALAI_TOKEN,
     true
   );
-  const msg = await client.Imagine("a cool cat, blue ears, red hat");
+  const msg = await client.Imagine("a cool cat, blue ears, yellow hat");
   console.log({ msg });
   if (!msg) {
     console.log("no message");

src/midjourney.message.ts

@@ -19,7 +19,8 @@ export class MidjourneyMessage {
   async FilterMessages(
     prompt: string,
     loading?: (uri: string) => void,
-    options?: string
+    options?: string,
+    index?: number
   ) {
     const data = await this.safeRetrieveMessages(this.Limit);
     for (let i = 0; i < data.length; i++) {
@@ -28,8 +29,16 @@ export class MidjourneyMessage {
         item.author.id === "936929561302675456" &&
         item.content.includes(`**${prompt}`)
       ) {
-        // this.log(JSON.stringify(item))
-        if (options && !item.content.includes(options)) {
+        this.log(JSON.stringify(item));
+        // Upscaled or Variation
+        if (
+          options &&
+          !(
+            item.content.includes(options) ||
+            (options === "Upscaled" &&
+              item.content.includes(` - Image #${index}`))
+          )
+        ) {
           this.log("no options");
           continue;
         }
@@ -39,7 +48,7 @@ export class MidjourneyMessage {
         }
         const imageUrl = item.attachments[0].url;
         if (!imageUrl.endsWith(".png")) {
-          loading && loading(imageUrl);
+          loading?.(imageUrl);
           break;
         }
         const content = item.content.split("**")[1];
@@ -82,6 +91,25 @@ export class MidjourneyMessage {
       await sleep(1000 * 2);
     }
   }
+  async WaitUpscaledMessage(
+    content: string,
+    index: number,
+    loading?: (uri: string) => void
+  ) {
+    for (let i = 0; i < this.maxWait; i++) {
+      const msg = await this.FilterMessages(
+        content,
+        loading,
+        "Upscaled",
+        index
+      );
+      if (msg !== null) {
+        return msg;
+      }
+      this.log(i, content, "wait no message found");
+      await sleep(1000 * 2);
+    }
+  }
 
   // limit the number of concurrent interactions
   protected async safeRetrieveMessages(limit = 50) {

src/midjourney.ts

@@ -159,7 +159,7 @@ export class Midjourney extends MidjourneyMessage {
       throw new Error(`VariationApi failed with status ${httpStatus}`);
     }
     this.log(`await generate image`);
-    return await this.WaitOptionMessage(content, `Upscaled`, loading);
+    return await this.WaitUpscaledMessage(content, index, loading);
   }
 
   async UpscaleApi(index: number, messageId: string, messageHash: string) {

test/test.ts

@@ -1,44 +1,45 @@
-type TaskFunction<T> = () => Promise<T>;
-
-interface QueueTask<T> {
-  task: TaskFunction<T>;
-  callback: (result: T) => void;
+import "dotenv/config";
+import { Midjourney } from "../src";
+/**
+ *
+ * a simple example of how to use the imagine command
+ * ```
+ * npx tsx test/test.ts
+ * ```
+ */
+async function main() {
+  const client = new Midjourney(
+    <string>process.env.SERVER_ID,
+    <string>process.env.CHANNEL_ID,
+    <string>process.env.SALAI_TOKEN,
+    true
+  );
+  const msg = await client.RetrieveMessages(10);
+  console.log(JSON.stringify(msg));
 }
-import async from "async";
-
-const queue = async.queue(async ({ task, callback }: QueueTask<any>) => {
-  const result = await task();
-  callback(result);
-}, 10);
-
-const task1: TaskFunction<number> = async () => {
-  // do some async work
-  return 42;
-};
-
-queue.push({
-  task: task1,
-  callback: (result: number) => {
-    console.log(result); // output: 42
-  },
-});
-queue.push({
-  task: task1,
-  callback: (result: number) => {
-    console.log(result); // output: 42
-  },
-});
-
-queue.push({
-  task: task1,
-  callback: (result: number) => {
-    console.log(result); // output: 42
-  },
-});
-
-queue.push({
-  task: task1,
-  callback: (result: number) => {
-    console.log(result); // output: 42
-  },
-});
+// main().catch((err) => {
+//   console.error(err);
+//   process.exit(1);
+// });
+function test2() {
+  const item = {
+    content:
+      "**a cool cat, blue ears, blue hat --seed 9113 --v 5** - Image #2 <@1017020769332637730>",
+  };
+  const options = "Upscaled";
+  const index = 2;
+  // if (options === "Upscaled" && !item.content.includes(` - Image #${index}`)) {
+  //   console.log("123123");
+  // }
+  if (
+    item.content.includes(options) ||
+    (options === "Upscaled" && item.content.includes(` - Image #${index}`))
+  ) {
+    console.log("has options");
+  } else {
+    console.log("no options");
+    return;
+  }
+  console.log("no options");
+}
+test2();