feat: support variation

Author: lvqq

README.md

@@ -34,7 +34,6 @@ const image = await midjourney.imagine('your prompt')
 
 const data = await midjourney.upscale('your prompt', {
   messageId: image.id,
-  index: 1,
   // custom_id could be found at image.component, for example: MJ::JOB::upsample::1::0c266431-26c6-47fa-bfee-2e1e11c7a66f
   customId: 'component custom_id'
 })
@@ -43,6 +42,28 @@ const data = await midjourney.upscale('your prompt', {
 console.log(data.attachments[0].url)
 ```
 
+### variation
+```typescript
+import { Midjourney } from 'midjourney-fetch'
+
+const midjourney = new Midjourney({
+  channelId: 'your channelId',
+  serverId: 'your serverId',
+  token: 'your token',
+})
+
+const image = await midjourney.imagine('your prompt')
+
+const data = await midjourney.variation('your prompt', {
+  messageId: image.id,
+  // custom_id could be found at image.component, for example: MJ::JOB::variation::1::0c266431-26c6-47fa-bfee-2e1e11c7a66f
+  customId: 'component custom_id'
+})
+
+// generated image url
+console.log(data.attachments[0].url)
+```
+
 ## How to get Ids and Token
 - [How to find ids](https://docs.statbot.net/docs/faq/general/how-find-id/)
 - [Get discord token](https://www.androidauthority.com/get-discord-token-3149920/)

package.json

@@ -1,6 +1,6 @@
 {
   "name": "midjourney-fetch",
-  "version": "1.0.0-beta.0",
+  "version": "1.0.0",
   "description": "",
   "type": "module",
   "main": "./dist/index.js",
@@ -57,6 +57,7 @@
     ]
   },
   "dependencies": {
-    "@sapphire/snowflake": "^3.5.1"
+    "@sapphire/snowflake": "^3.5.1",
+    "dayjs": "^1.11.7"
   }
 }

pnpm-lock.yaml

@@ -43,22 +43,33 @@ export interface MessageItem {
     components: MessageComponent[];
     type: number;
   }>;
+  message_reference?: {
+    channel_id: string;
+    guild_id: string;
+    message_id: string;
+  };
+  timestamp?: string;
 }
 
-export type MessageType = 'imagine' | 'upscale';
+export type MessageType = 'imagine' | 'upscale' | 'variation';
 
-export type MessageTypeProps =
+export type MessageTypeProps = {
+  timestamp: string;
+} & (
+  | {
+      type: Extract<MessageType, 'variation'>;
+      index: number;
+    }
   | {
       type: Extract<MessageType, 'upscale'>;
       index: number;
     }
   | {
       type?: Extract<MessageType, 'imagine'>;
-    };
+    }
+);
 
 export interface UpscaleProps {
   messageId: string;
-  index: number;
-  hash?: string;
-  customId?: string;
+  customId: string;
 }

src/interface.ts

@@ -2,11 +2,16 @@ import { DiscordSnowflake } from '@sapphire/snowflake';
 import { configs, defaultSessionId, midjourneyBotConfigs } from './config';
 import type {
   MessageItem,
+  MessageType,
   MessageTypeProps,
   MidjourneyProps,
   UpscaleProps,
 } from './interface';
-import { findMessageByPrompt, isInProgress } from './utils';
+import {
+  findMessageByPrompt,
+  getHashFromCustomId,
+  isInProgress,
+} from './utils';
 
 export class Midjourney {
   protected readonly channelId: string;
@@ -45,6 +50,17 @@ export class Midjourney {
     }
   }
 
+  async interactions(payload: any) {
+    return fetch(`https://discord.com/api/v9/interactions`, {
+      method: 'POST',
+      body: JSON.stringify(payload),
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: this.token,
+      },
+    });
+  }
+
   async createImage(prompt: string) {
     const payload = {
       type: 2,
@@ -90,14 +106,7 @@ export class Midjourney {
       nonce: DiscordSnowflake.generate().toString(),
     };
 
-    const res = await fetch(`https://discord.com/api/v9/interactions`, {
-      method: 'POST',
-      body: JSON.stringify(payload),
-      headers: {
-        'Content-Type': 'application/json',
-        Authorization: this.token,
-      },
-    });
+    const res = await this.interactions(payload);
     if (res.status >= 400) {
       let message = '';
       try {
@@ -113,7 +122,10 @@ export class Midjourney {
     }
   }
 
-  async createUpscale({ messageId, index, hash, customId }: UpscaleProps) {
+  async createUpscaleOrVariation(
+    type: Exclude<MessageType, 'imagine'>,
+    { messageId, customId }: UpscaleProps
+  ) {
     const payload = {
       type: 3,
       nonce: DiscordSnowflake.generate().toString(),
@@ -125,33 +137,26 @@ export class Midjourney {
       session_id: defaultSessionId,
       data: {
         component_type: 2,
-        custom_id: customId || `MJ::JOB::upsample::${index}::${hash}`,
+        custom_id: customId,
       },
     };
-    const res = await fetch(`https://discord.com/api/v9/interactions`, {
-      method: 'POST',
-      body: JSON.stringify(payload),
-      headers: {
-        'Content-Type': 'application/json',
-        Authorization: this.token,
-      },
-    });
+    const res = await this.interactions(payload);
     if (res.status >= 400) {
       let message = '';
       try {
         const data = await res.json();
         if (this.debugger) {
-          this.log('Create upscale failed', JSON.stringify(data));
+          this.log(`Create ${type} failed`, JSON.stringify(data));
         }
         message = data?.message;
       } catch (e) {
         // catch JSON error
       }
-      throw new Error(message || `Create upscale failed with ${res.status}`);
+      throw new Error(message || `Create ${type} failed with ${res.status}`);
     }
   }
 
-  async getMessage(prompt: string, options?: MessageTypeProps) {
+  async getMessage(prompt: string, options: MessageTypeProps) {
     const res = await fetch(
       `https://discord.com/api/v10/channels/${this.channelId}/messages?limit=50`,
       {
@@ -170,7 +175,10 @@ export class Midjourney {
    * Same with /imagine command
    */
   async imagine(prompt: string) {
+    const timestamp = new Date().toISOString();
+
     await this.createImage(prompt);
+
     const times = this.timeout / this.interval;
     let count = 0;
     let result: MessageItem | undefined;
@@ -179,7 +187,7 @@ export class Midjourney {
         count += 1;
         await new Promise((res) => setTimeout(res, this.interval));
         this.log(count, 'imagine');
-        const message = await this.getMessage(prompt);
+        const message = await this.getMessage(prompt, { timestamp });
         if (message && !isInProgress(message)) {
           result = message;
           break;
@@ -192,18 +200,63 @@ export class Midjourney {
   }
 
   async upscale({ prompt, ...params }: UpscaleProps & { prompt: string }) {
-    await this.createUpscale(params);
+    const { index } = getHashFromCustomId('upscale', params.customId);
     const times = this.timeout / this.interval;
     let count = 0;
     let result: MessageItem | undefined;
+
+    if (!index) {
+      throw new Error('Create upscale failed with 400, unknown customId');
+    }
+
+    const timestamp = new Date().toISOString();
+
+    await this.createUpscaleOrVariation('upscale', params);
+
     while (count < times) {
       try {
         count += 1;
         await new Promise((res) => setTimeout(res, this.interval));
         this.log(count, 'upscale');
         const message = await this.getMessage(prompt, {
           type: 'upscale',
-          index: params.index,
+          index,
+          timestamp,
+        });
+        if (message && !isInProgress(message)) {
+          result = message;
+          break;
+        }
+      } catch {
+        continue;
+      }
+    }
+    return result;
+  }
+
+  async variation({ prompt, ...params }: UpscaleProps & { prompt: string }) {
+    const { index } = getHashFromCustomId('variation', params.customId);
+    const times = this.timeout / this.interval;
+    let count = 0;
+    let result: MessageItem | undefined;
+
+    if (!index) {
+      throw new Error('Create variation failed with 400, unknown customId');
+    }
+
+    const timestamp = new Date().toISOString();
+
+    await this.createUpscaleOrVariation('variation', params);
+
+    while (count < times) {
+      try {
+        count += 1;
+        await new Promise((res) => setTimeout(res, this.interval));
+        this.log(count, 'variation');
+        const message = await this.getMessage(prompt, {
+          type: 'variation',
+          index,
+          timestamp,
         });
         if (message && !isInProgress(message)) {
           result = message;

src/midjourney.ts

@@ -1,24 +1,41 @@
+import dayjs from 'dayjs';
 import { midjourneyBotConfigs } from './config';
-import type { MessageTypeProps, MessageItem } from './interface';
+import type { MessageTypeProps, MessageItem, MessageType } from './interface';
 
 export const findMessageByPrompt = (
   messages: MessageItem[],
   prompt: string,
-  options?: MessageTypeProps
+  options: MessageTypeProps
 ) => {
   // trim and merge spaces
   const filterPrompt = prompt.split(' ').filter(Boolean).join(' ');
   if (options?.type === 'upscale') {
     return messages.find(
       (msg) =>
+        msg.timestamp &&
+        dayjs(msg.timestamp).isAfter(options.timestamp) &&
         msg.type === 19 &&
         msg.content.includes(filterPrompt) &&
         msg.content.includes(`Image #${options.index}`) &&
         msg.author.id === midjourneyBotConfigs.applicationId
     );
   }
+  if (options?.type === 'variation') {
+    return messages.find(
+      (msg) =>
+        msg.timestamp &&
+        dayjs(msg.timestamp).isAfter(options.timestamp) &&
+        msg.type === 19 &&
+        msg.content.includes(filterPrompt) &&
+        // 0 means reroll
+        (!options?.index || msg.content.includes('Variations')) &&
+        msg.author.id === midjourneyBotConfigs.applicationId
+    );
+  }
   return messages.find(
     (msg) =>
+      msg.timestamp &&
+      dayjs(msg.timestamp).isAfter(options.timestamp) &&
       msg.content.includes(filterPrompt) &&
       msg.author.id === midjourneyBotConfigs.applicationId
   );
@@ -28,3 +45,18 @@ export const isInProgress = (message: MessageItem) =>
   message.attachments.length === 0 ||
   (message.attachments[0]?.filename?.startsWith('grid') &&
     message.attachments[0]?.filename?.endsWith('.webp'));
+
+export const getHashFromCustomId = (type: MessageType, id: string) => {
+  let regex: RegExp | null = null;
+  if (type === 'upscale') {
+    regex = /(upsample)::(\d+)::(.+)/;
+  } else if (type === 'variation') {
+    regex = /(variation|reroll)::(\d+)::(.+)/;
+  }
+  if (!regex) return { index: null, hash: null };
+  const match = id.match(regex);
+  const model = match?.[1]; // upsample|variation|reroll
+  const index = match?.[2] ? Number(match[2]) : null;
+  const hash = match?.[3] || null;
+  return { model, index, hash };
+};