Merge pull request #90 from erictik/mjbot-dm

DM Midjourney Bot

example/imagine-dm.ts

@@ -0,0 +1,34 @@
+import "dotenv/config";
+import { Midjourney } from "../src";
+/**
+ *
+ * a simple example of using the imagine api via DM Midjourney Bot
+ * ```
+ * npx tsx example/imagine-dm.ts
+ * ```
+ */
+async function main() {
+  const client = new Midjourney({
+    SalaiToken: <string>process.env.SALAI_TOKEN,
+    Debug: true,
+    Ws: true,
+  });
+  await client.init();
+  const msg = await client.Imagine(
+    "A little white dog",
+    (uri: string, progress: string) => {
+      console.log("loading", uri, "progress", progress);
+    }
+  );
+  console.log({ msg });
+}
+main()
+  .then(() => {
+    console.log("finished");
+    process.exit(0);
+  })
+  .catch((err) => {
+    console.log("finished");
+    console.error(err);
+    process.exit(1);
+  });

src/interfaces/config.ts

@@ -1,19 +1,19 @@
-export interface MessageConfig {
+export interface MJConfig {
   ChannelId: string;
   SalaiToken: string;
   Debug: boolean;
   Limit: number;
   MaxWait: number;
-  ServerId: string;
   SessionId: string;
+  ServerId?: string;
   Ws?: boolean;
   HuggingFaceToken?: string;
-  DiscordBaseUrl?: string;
-  WsBaseUrl?: string;
+  DiscordBaseUrl: string;
+  WsBaseUrl: string;
 }
-export interface MessageConfigParam {
-  ChannelId: string;
+export interface MJConfigParam {
   SalaiToken: string;
+  ChannelId?: string;
   Debug?: boolean;
   Limit?: number;
   MaxWait?: number;
@@ -24,28 +24,14 @@ export interface MessageConfigParam {
   DiscordBaseUrl?: string;
   WsBaseUrl?: string;
 }
-export interface MidjourneyConfig extends MessageConfig {
-  ServerId: string;
-  SessionId: string;
-}
 
-export interface MidjourneyConfigParam extends MessageConfigParam {
-  ServerId: string;
-  SessionId?: string;
-}
-export const DefaultMessageConfig: MessageConfig = {
-  ChannelId: "",
+export const DefaultMJConfig: MJConfig = {
+  ChannelId: "1077800642086703114",
   SalaiToken: "",
-  ServerId: "",
-  SessionId: "",
+  SessionId: "8bb7f5b79c7a49f7d0824ab4b8773a81",
   Debug: false,
   Limit: 50,
   MaxWait: 100,
   DiscordBaseUrl: "https://discord.com",
-  WsBaseUrl: "wss://gateway.discord.gg"
-};
-export const DefaultMidjourneyConfig: MidjourneyConfig = {
-  ...DefaultMessageConfig,
-  ServerId: "",
-  SessionId: "8bb7f5b79c7a49f7d0824ab4b8773a81",
+  WsBaseUrl: "wss://gateway.discord.gg?v=9&encoding=json&compress=gzip-stream`",
 };

src/midjourne.api.ts

@@ -0,0 +1,157 @@
+import { MJConfig } from "./interfaces";
+import { CreateQueue } from "./queue";
+import { nextNonce, sleep } from "./utls";
+export class MidjourneyApi {
+  private ApiQueue = CreateQueue(1);
+  constructor(public config: MJConfig) {}
+  // limit the number of concurrent interactions
+  protected async safeIteractions(payload: any) {
+    return this.ApiQueue.addTask(
+      () =>
+        new Promise<number>((resolve) => {
+          this.interactions(payload, (res) => {
+            resolve(res);
+          });
+        })
+    );
+  }
+  protected async interactions(
+    payload: any,
+    callback?: (result: number) => void
+  ) {
+    try {
+      const headers = {
+        "Content-Type": "application/json",
+        Authorization: this.config.SalaiToken,
+      };
+      const response = await fetch(
+        `${this.config.DiscordBaseUrl}/api/v9/interactions`,
+        {
+          method: "POST",
+          body: JSON.stringify(payload),
+          headers: headers,
+        }
+      );
+      callback && callback(response.status);
+      //discord api rate limit
+      await sleep(950);
+      if (response.status >= 400) {
+        console.error("api.error.config", { config: this.config });
+      }
+      return response.status;
+    } catch (error) {
+      console.error(error);
+      callback && callback(500);
+    }
+  }
+  async ImagineApi(prompt: string, nonce: string = nextNonce()) {
+    const guild_id = this.config.ServerId;
+    const payload = {
+      type: 2,
+      application_id: "936929561302675456",
+      guild_id,
+      channel_id: this.config.ChannelId,
+      session_id: this.config.SessionId,
+      data: {
+        version: "1077969938624553050",
+        id: "938956540159881230",
+        name: "imagine",
+        type: 1,
+        options: [
+          {
+            type: 3,
+            name: "prompt",
+            value: prompt,
+          },
+        ],
+        application_command: {
+          id: "938956540159881230",
+          application_id: "936929561302675456",
+          version: "1077969938624553050",
+          default_permission: true,
+          default_member_permissions: null,
+          type: 1,
+          nsfw: false,
+          name: "imagine",
+          description: "Create images with Midjourney",
+          dm_permission: true,
+          options: [
+            {
+              type: 3,
+              name: "prompt",
+              description: "The prompt to imagine",
+              required: true,
+            },
+          ],
+        },
+        attachments: [],
+      },
+      nonce,
+    };
+    return this.safeIteractions(payload);
+  }
+  async VariationApi(
+    index: number,
+    messageId: string,
+    messageHash: string,
+    nonce?: string
+  ) {
+    const payload = {
+      type: 3,
+      guild_id: this.config.ServerId,
+      channel_id: this.config.ChannelId,
+      message_flags: 0,
+      message_id: messageId,
+      application_id: "936929561302675456",
+      session_id: "1f3dbdf09efdf93d81a3a6420882c92c",
+      data: {
+        component_type: 2,
+        custom_id: `MJ::JOB::variation::${index}::${messageHash}`,
+      },
+      nonce,
+    };
+    return this.safeIteractions(payload);
+  }
+  async UpscaleApi(
+    index: number,
+    messageId: string,
+    messageHash: string,
+    nonce?: string
+  ) {
+    const guild_id = this.config.ServerId;
+    const payload = {
+      type: 3,
+      guild_id,
+      channel_id: this.config.ChannelId,
+      message_flags: 0,
+      message_id: messageId,
+      application_id: "936929561302675456",
+      session_id: "ec6524c8d2926e285a8232f7ed1ced98",
+      data: {
+        component_type: 2,
+        custom_id: `MJ::JOB::upsample::${index}::${messageHash}`,
+      },
+      nonce,
+    };
+    return this.safeIteractions(payload);
+  }
+
+  async ClickBtnApi(messageId: string, customId: string, nonce?: string) {
+    const guild_id = this.config.ServerId;
+    const payload = {
+      type: 3,
+      nonce,
+      guild_id,
+      channel_id: this.config.ChannelId,
+      message_flags: 0,
+      message_id: messageId,
+      application_id: "936929561302675456",
+      session_id: this.config.SessionId,
+      data: {
+        component_type: 2,
+        custom_id: customId,
+      },
+    };
+    return this.safeIteractions(payload);
+  }
+}

src/midjourney.message.ts

@@ -1,25 +1,24 @@
 import {
-  DefaultMessageConfig,
+  DefaultMJConfig,
   LoadingHandler,
   MJMessage,
-  MessageConfig,
-  MessageConfigParam,
-  MidjourneyConfig,
+  MJConfig,
+  MJConfigParam,
 } from "./interfaces";
 import { CreateQueue } from "./queue";
 import { sleep } from "./utls";
 
 export class MidjourneyMessage {
   private magApiQueue = CreateQueue(1);
-  public config: MessageConfig;
-  constructor(defaults: MessageConfigParam) {
-    const { ChannelId, SalaiToken } = defaults;
-    if (!ChannelId || !SalaiToken) {
-      throw new Error("ChannelId and SalaiToken are required");
+  public config: MJConfig;
+  constructor(defaults: MJConfigParam) {
+    const { SalaiToken } = defaults;
+    if (!SalaiToken) {
+      throw new Error("SalaiToken are required");
     }
 
     this.config = {
-      ...DefaultMessageConfig,
+      ...DefaultMJConfig,
       ...defaults,
     };
   }
@@ -151,7 +150,7 @@ export class MidjourneyMessage {
   async RetrieveMessages(limit = this.config.Limit) {
     const headers = { authorization: this.config.SalaiToken };
     const response = await fetch(
-        `${this.config.DiscordBaseUrl}/api/v10/channels/${this.config.ChannelId}/messages?limit=${limit}`,
+      `${this.config.DiscordBaseUrl}/api/v10/channels/${this.config.ChannelId}/messages?limit=${limit}`,
       {
         headers: headers,
       }