feat: support upscale

Author: lvqq

README.md

@@ -4,6 +4,7 @@
 Fetch api for midjourney on discord
 
 ## Usage
+### imagine
 ```typescript
 import { Midjourney } from 'midjourney-fetch'
 
@@ -13,9 +14,33 @@ const midjourney = new Midjourney({
   token: 'your token',
 })
 
-const images = await midjourney.imagine('your prompt')
+const data = await midjourney.imagine('your prompt')
 
-console.log(images[0].url)
+// generated image url
+console.log(data.attachments[0].url)
+```
+
+### upscale
+```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.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'
+})
+
+// generated image url
+console.log(data.attachments[0].url)
 ```
 
 ## How to get Ids and Token

package.json

@@ -1,6 +1,6 @@
 {
   "name": "midjourney-fetch",
-  "version": "0.1.4",
+  "version": "1.0.0-beta.0",
   "description": "",
   "type": "module",
   "main": "./dist/index.js",
@@ -55,5 +55,8 @@
       "eslint --fix --quiet",
       "prettier --write"
     ]
+  },
+  "dependencies": {
+    "@sapphire/snowflake": "^3.5.1"
   }
 }

pnpm-lock.yaml

@@ -2,3 +2,11 @@ export const configs = {
   timeout: 5 * 60 * 1000, // 5 min
   interval: 15 * 1000, // every 15 second
 };
+
+export const defaultSessionId = 'ab318945494d4aa96c97ce6fce934b97';
+
+export const midjourneyBotConfigs = {
+  applicationId: '936929561302675456',
+  version: '1077969938624553050',
+  id: '938956540159881230',
+};

src/config.ts

@@ -20,6 +20,14 @@ export interface MessageAttachment {
   id: string;
 }
 
+export interface MessageComponent {
+  custom_id: string; // starts with MJ::JOB::upsample or MJ::JOB::variation
+  label?: 'U1' | 'U2' | 'U3' | 'U4' | 'V1' | 'V2' | 'V3' | 'V4';
+  emoji?: { name: string };
+  style: number; // 1 - used; 2 - free
+  type: number;
+}
+
 export interface MessageItem {
   application_id: string;
   attachments: MessageAttachment[];
@@ -30,4 +38,27 @@ export interface MessageItem {
   channel_id: string;
   content: string;
   id: string;
+  type: number; // 19 - upscale; 0 - imagine
+  components: Array<{
+    components: MessageComponent[];
+    type: number;
+  }>;
+}
+
+export type MessageType = 'imagine' | 'upscale';
+
+export type MessageTypeProps =
+  | {
+      type: Extract<MessageType, 'upscale'>;
+      index: number;
+    }
+  | {
+      type?: Extract<MessageType, 'imagine'>;
+    };
+
+export interface UpscaleProps {
+  messageId: string;
+  index: number;
+  hash?: string;
+  customId?: string;
 }

src/interface.ts

@@ -1,8 +1,10 @@
-import { configs } from './config';
+import { DiscordSnowflake } from '@sapphire/snowflake';
+import { configs, defaultSessionId, midjourneyBotConfigs } from './config';
 import type {
-  MessageAttachment,
   MessageItem,
+  MessageTypeProps,
   MidjourneyProps,
+  UpscaleProps,
 } from './interface';
 import { findMessageByPrompt, isInProgress } from './utils';
 
@@ -43,16 +45,16 @@ export class Midjourney {
     }
   }
 
-  async interactions(prompt: string) {
+  async createImage(prompt: string) {
     const payload = {
       type: 2,
-      application_id: '936929561302675456',
+      application_id: midjourneyBotConfigs.applicationId,
       guild_id: this.serverId,
       channel_id: this.channelId,
-      session_id: 'ab318945494d4aa96c97ce6fce934b97',
+      session_id: defaultSessionId,
       data: {
-        version: '1077969938624553050',
-        id: '938956540159881230',
+        version: midjourneyBotConfigs.version,
+        id: midjourneyBotConfigs.id,
         name: 'imagine',
         type: 1,
         options: [
@@ -63,9 +65,9 @@ export class Midjourney {
           },
         ],
         application_command: {
-          id: '938956540159881230',
-          application_id: '936929561302675456',
-          version: '1077969938624553050',
+          id: midjourneyBotConfigs.id,
+          application_id: midjourneyBotConfigs.applicationId,
+          version: midjourneyBotConfigs.version,
           default_permission: true,
           default_member_permissions: null,
           type: 1,
@@ -85,6 +87,7 @@ export class Midjourney {
         },
         attachments: [],
       },
+      nonce: DiscordSnowflake.generate().toString(),
     };
 
     const res = await fetch(`https://discord.com/api/v9/interactions`, {
@@ -100,17 +103,55 @@ export class Midjourney {
       try {
         const data = await res.json();
         if (this.debugger) {
-          this.log('Interactions failed', JSON.stringify(data));
+          this.log('Create image failed', JSON.stringify(data));
         }
         message = data?.message;
       } catch (e) {
         // catch JSON error
       }
-      throw new Error(message || `Interactions failed with ${res.status}`);
+      throw new Error(message || `Create image failed with ${res.status}`);
     }
   }
 
-  async getMessage(prompt: string) {
+  async createUpscale({ messageId, index, hash, customId }: UpscaleProps) {
+    const payload = {
+      type: 3,
+      nonce: DiscordSnowflake.generate().toString(),
+      guild_id: this.serverId,
+      channel_id: this.channelId,
+      message_flags: 0,
+      message_id: messageId,
+      application_id: midjourneyBotConfigs.applicationId,
+      session_id: defaultSessionId,
+      data: {
+        component_type: 2,
+        custom_id: customId || `MJ::JOB::upsample::${index}::${hash}`,
+      },
+    };
+    const res = await fetch(`https://discord.com/api/v9/interactions`, {
+      method: 'POST',
+      body: JSON.stringify(payload),
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: this.token,
+      },
+    });
+    if (res.status >= 400) {
+      let message = '';
+      try {
+        const data = await res.json();
+        if (this.debugger) {
+          this.log('Create upscale failed', JSON.stringify(data));
+        }
+        message = data?.message;
+      } catch (e) {
+        // catch JSON error
+      }
+      throw new Error(message || `Create upscale failed with ${res.status}`);
+    }
+  }
+
+  async getMessage(prompt: string, options?: MessageTypeProps) {
     const res = await fetch(
       `https://discord.com/api/v10/channels/${this.channelId}/messages?limit=50`,
       {
@@ -120,7 +161,7 @@ export class Midjourney {
       }
     );
     const data: MessageItem[] = await res.json();
-    const message = findMessageByPrompt(data, prompt);
+    const message = findMessageByPrompt(data, prompt, options);
     this.log(JSON.stringify(message), '\n');
     return message;
   }
@@ -129,24 +170,49 @@ export class Midjourney {
    * Same with /imagine command
    */
   async imagine(prompt: string) {
-    await this.interactions(prompt);
+    await this.createImage(prompt);
     const times = this.timeout / this.interval;
     let count = 0;
-    let image: MessageAttachment | null = null;
+    let result: MessageItem | undefined;
     while (count < times) {
       try {
         count += 1;
         await new Promise((res) => setTimeout(res, this.interval));
-        this.log(count);
+        this.log(count, 'imagine');
         const message = await this.getMessage(prompt);
         if (message && !isInProgress(message)) {
-          [image] = message.attachments;
+          result = message;
+          break;
+        }
+      } catch {
+        continue;
+      }
+    }
+    return result;
+  }
+
+  async upscale({ prompt, ...params }: UpscaleProps & { prompt: string }) {
+    await this.createUpscale(params);
+    const times = this.timeout / this.interval;
+    let count = 0;
+    let result: MessageItem | undefined;
+    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,
+        });
+        if (message && !isInProgress(message)) {
+          result = message;
           break;
         }
       } catch {
         continue;
       }
     }
-    return image ? [image] : [];
+    return result;
   }
 }

src/midjourney.ts

@@ -1,15 +1,26 @@
-import { type MessageItem } from './interface';
+import { midjourneyBotConfigs } from './config';
+import type { MessageTypeProps, MessageItem } from './interface';
 
 export const findMessageByPrompt = (
   messages: MessageItem[],
-  prompt: string
+  prompt: string,
+  options?: MessageTypeProps
 ) => {
   // trim and merge spaces
   const filterPrompt = prompt.split(' ').filter(Boolean).join(' ');
+  if (options?.type === 'upscale') {
+    return messages.find(
+      (msg) =>
+        msg.type === 19 &&
+        msg.content.includes(filterPrompt) &&
+        msg.content.includes(`Image #${options.index}`) &&
+        msg.author.id === midjourneyBotConfigs.applicationId
+    );
+  }
   return messages.find(
     (msg) =>
       msg.content.includes(filterPrompt) &&
-      msg.author.id === '936929561302675456'
+      msg.author.id === midjourneyBotConfigs.applicationId
   );
 };
 

src/utils.ts

@@ -3,7 +3,7 @@ import { defineConfig } from 'tsup';
 export default defineConfig({
   entry: ['src/index.ts'],
   format: ['esm', 'cjs'],
-  target: 'node16',
+  target: 'node18',
   dts: true,
   legacyOutput: true,
 });