@mymj/midjourney/src/midjourney.api.ts

Node.js client for the unofficial MidJourney API.

import {
  CustomZoomModalSubmitID,
  DescribeModalSubmitID,
  DiscordImage,
  MJConfig,
  RemixModalSubmitID,
  ShortenModalSubmitID,
  UploadParam,
  UploadSlot,
  QueueItem,
} from "./interfaces";

import { nextNonce, sleep } from "./utils";
import { Command, CommandName } from "./command";
import async, { QueueObject } from "async";

export class MidjourneyApi extends Command {
  private queue: QueueObject<QueueItem>;
  UpId = Date.now() % 10; // upload id
  constructor(public config: MJConfig) {
    super(config);
    this.queue = async.queue(this.processRequest, 1);
  }
  private safeIteractions = (request: any) => {
    return new Promise<number>((resolve, reject) => {
      this.queue.push(
        {
          request,
          callback: (any: any) => {
            resolve(any);
          },
        },
        (error: any, result: any) => {
          if (error) {
            reject(error);
          } else {
            resolve(result);
          }
        }
      );
    });
  };
  private processRequest = async ({
    request,
    callback,
  }: QueueItem) => {
    console.log('process request', request);
    const httpStatus = request.mask ? await this.inpaint(request) : await this.interactions(request);
    callback(httpStatus);
    await sleep(this.config.ApiInterval);
  };
  private interactions = async (payload: any) => {
    try {
      const headers = {
        "Content-Type": "application/json",
        Authorization: this.config.SalaiToken,
      };
      const response = await this.config.fetch(
        `${this.config.DiscordBaseUrl}/api/v9/interactions`,
        {
          method: "POST",
          body: JSON.stringify(payload),
          headers: headers,
        }
      );
      if (response.status >= 400) {
        console.error("api.error.config", {
          payload: JSON.stringify(payload),
          config: this.config,
        });
      }
      return response.status;
    } catch (error) {
      console.error(error);
      return 500;
    }
  };

  private inpaint = async (payload: any) => {
    try {
      const headers = {
        "Content-Type": "application/json",
      };
      const response = await this.config.fetch(
        `${this.config.DiscordBotUrl}/inpaint/api/submit-job`,
        {
          method: "POST",
          body: JSON.stringify(payload),
          headers: headers,
        }
      );
      if (response.status >= 400) {
        console.error("api.error.config", {
          payload: JSON.stringify(payload),
          config: this.config,
          response,
        });
      }
      return response.status;
    } catch (error) {
      console.error(error);
      return 500;
    }
  };

  async ImagineApi(prompt: string, nonce: string = nextNonce()) {
    const payload = await this.imaginePayload(prompt, nonce);
    return this.safeIteractions(payload);
  }

  async SwitchRemixApi(nonce: string = nextNonce()) {
    const payload = await this.PreferPayload(nonce);
    return this.safeIteractions(payload);
  }

  async ShortenApi(prompt: string, nonce: string = nextNonce()) {
    const payload = await this.shortenPayload(prompt, nonce);
    return this.safeIteractions(payload);
  }

  async VariationApi({
    index,
    msgId,
    hash,
    nonce = nextNonce(),
    flags = 0,
  }: {
    index: 1 | 2 | 3 | 4;
    msgId: string;
    hash: string;
    nonce?: string;
    flags?: number;
  }) {
    return this.CustomApi({
      msgId,
      customId: `MJ::JOB::variation::${index}::${hash}`,
      flags,
      nonce,
    });
  }

  async UpscaleApi({
    index,
    msgId,
    hash,
    nonce = nextNonce(),
    flags,
  }: {
    index: 1 | 2 | 3 | 4;
    msgId: string;
    hash: string;
    nonce?: string;
    flags: number;
  }) {
    return this.CustomApi({
      msgId,
      customId: `MJ::JOB::upsample::${index}::${hash}`,
      flags,
      nonce,
    });
  }

  async RerollApi({
    msgId,
    hash,
    nonce = nextNonce(),
    flags,
  }: {
    msgId: string;
    hash: string;
    nonce?: string;
    flags: number;
  }) {
    return this.CustomApi({
      msgId,
      customId: `MJ::JOB::reroll::0::${hash}::SOLO`,
      flags,
      nonce,
    });
  }

  async CustomApi({
    msgId,
    customId,
    flags,
    nonce = nextNonce(),
  }: {
    msgId: string;
    customId: string;
    flags: number;
    nonce?: string;
  }) {
    if (!msgId) throw new Error("msgId is empty");
    if (flags === undefined) throw new Error("flags is undefined");
    // VaryRegion must use active session id after auth
    // The others use static session id
    const session_id = this.config.ActiveSessionId;
    const payload = {
      type: 3,
      nonce,
      guild_id: this.config.ServerId,
      channel_id: this.config.ChannelId,
      message_flags: flags,
      message_id: msgId,
      application_id: this.config.BotId,
      session_id,
      data: {
        component_type: 2,
        custom_id: customId,
      },
    };
    return this.safeIteractions(payload);
  }

  async InpaintApi({
    customId,
    prompt,
    mask,
  }: {
    customId: string;
    prompt: string,
    mask: string,
  }) {
    if (!customId) throw new Error("customId is empty");
    if (mask === undefined) throw new Error("mask is undefined");
    const payload = {
      customId,
      prompt,
      mask: mask.replace(/^data:.+?;base64,/, ''),
      userId: '0',
      username: '0',
      full_prompt: null,
    };
    return this.safeIteractions(payload);
  }

  //FIXME: get SubmitCustomId from discord api
  async ModalSubmitApi({
    nonce,
    msgId,
    customId,
    prompt,
    submitCustomId,
  }: {
    nonce: string;
    msgId: string;
    customId: string;
    prompt: string;
    submitCustomId: string;
  }) {
    var payload = {
      type: 5,
      application_id: this.config.BotId,
      channel_id: this.config.ChannelId,
      guild_id: this.config.ServerId,
      data: {
        id: msgId,
        custom_id: customId,
        components: [
          {
            type: 1,
            components: [
              {
                type: 4,
                custom_id: submitCustomId,
                value: prompt,
              },
            ],
          },
        ],
      },
      session_id: this.config.SessionId,
      nonce,
    };
    console.log("submitCustomId", JSON.stringify(payload));
    return this.safeIteractions(payload);
  }
  async RemixApi({
    nonce,
    msgId,
    customId,
    prompt,
    submitCustomId,
  }: {
    nonce: string;
    msgId: string;
    customId: string;
    prompt: string;
    submitCustomId?: string;
  }) {
    return this.ModalSubmitApi({
      nonce,
      msgId,
      customId,
      prompt,
      submitCustomId: submitCustomId || RemixModalSubmitID,
    });
  }

  async ShortenImagineApi({
    nonce,
    msgId,
    customId,
    prompt,
  }: {
    nonce: string;
    msgId: string;
    customId: string;
    prompt: string;
  }) {
    return this.ModalSubmitApi({
      nonce,
      msgId,
      customId,
      prompt,
      submitCustomId: ShortenModalSubmitID,
    });
  }

  async DescribeImagineApi({
    nonce,
    msgId,
    customId,
    prompt,
  }: {
    nonce: string;
    msgId: string;
    customId: string;
    prompt: string;
  }) {
    return this.ModalSubmitApi({
      nonce,
      msgId,
      customId,
      prompt,
      submitCustomId: DescribeModalSubmitID,
    });
  }

  async CustomZoomImagineApi({
    nonce,
    msgId,
    customId,
    prompt,
  }: {
    nonce: string;
    msgId: string;
    customId: string;
    prompt: string;
  }) {
    // customId = customId.replace(
    //   "MJ::CustomZoom",
    //   "MJ::OutpaintCustomZoomModal"
    // );
    return this.ModalSubmitApi({
      nonce,
      msgId,
      customId,
      prompt,
      submitCustomId: CustomZoomModalSubmitID,
    });
  }

  private simpleCommand = async (name: CommandName, nonce?: string) => {
    const payload = await this.commandPayload(name, nonce);
    return this.safeIteractions(payload);
  }

  async InfoApi(nonce?: string) {
    return await this.simpleCommand('info', nonce);
  }

  async SubscribeApi(nonce?: string) {
    return await this.simpleCommand('subscribe', nonce);
  }

  async SettingsApi(nonce?: string) {
    return await this.simpleCommand('settings', nonce);
  }

  async TurboApi(nonce?: string) {
    return await this.simpleCommand('turbo', nonce);
  }

  async FastApi(nonce?: string) {
    return await this.simpleCommand('fast', nonce);
  }

  async RelaxApi(nonce?: string) {
    return await this.simpleCommand('relax', nonce);
  }

  async StealthApi(nonce?: string) {
    return await this.simpleCommand('stealth', nonce);
  }

  /**
   *
   * @param fileUrl http file path
   * @returns
   */
  async UploadImageByUri(fileUrl: string) {
    const response = await this.config.fetch(fileUrl);
    const fileData = await response.arrayBuffer();
    const mimeType = response.headers.get("content-type");
    const filename = fileUrl.split("/").pop() || "image.png";
    const file_size = fileData.byteLength;
    if (!mimeType) {
      throw new Error("Unknown mime type");
    }
    const { attachments } = await this.attachments({
      filename,
      file_size,
      id: this.UpId++,
    });
    const UploadSlot = attachments[0];
    await this.uploadImage(UploadSlot, fileData, mimeType);
    const resp: DiscordImage = {
      id: UploadSlot.id,
      filename: UploadSlot.upload_filename.split("/").pop() || "image.png",
      upload_filename: UploadSlot.upload_filename,
    };
    return resp;
  }

  async UploadImageByBole(blob: Blob, filename = nextNonce() + ".png") {
    const fileData = await blob.arrayBuffer();
    const mimeType = blob.type;
    const file_size = fileData.byteLength;
    if (!mimeType) {
      throw new Error("Unknown mime type");
    }
    const { attachments } = await this.attachments({
      filename,
      file_size,
      id: this.UpId++,
    });
    const UploadSlot = attachments[0];
    await this.uploadImage(UploadSlot, fileData, mimeType);
    const resp: DiscordImage = {
      id: UploadSlot.id,
      filename: UploadSlot.upload_filename.split("/").pop() || "image.png",
      upload_filename: UploadSlot.upload_filename,
    };
    return resp;
  }

  /**
   * prepare an attachement to upload an image.
   */
  private async attachments(
    ...files: UploadParam[]
  ): Promise<{ attachments: UploadSlot[] }> {
    const { SalaiToken, DiscordBaseUrl, ChannelId, fetch } = this.config;
    const headers = {
      Authorization: SalaiToken,
      "content-type": "application/json",
    };
    const url = new URL(
      `${DiscordBaseUrl}/api/v9/channels/${ChannelId}/attachments`
    );
    const body = { files };
    const response = await this.config.fetch(url, {
      headers,
      method: "POST",
      body: JSON.stringify(body),
    });
    if (response.status === 200) {
      return (await response.json()) as { attachments: UploadSlot[] };
    }
    const error = `Attachments return ${response.status} ${
      response.statusText
    } ${await response.text()}`;
    throw new Error(error);
  }

  private async uploadImage(
    slot: UploadSlot,
    data: ArrayBuffer,
    contentType: string
  ): Promise<void> {
    const body = new Uint8Array(data);
    const headers = { "content-type": contentType };
    const response = await this.config.fetch(slot.upload_url, {
      method: "PUT",
      headers,
      body,
    });
    if (!response.ok) {
      throw new Error(
        `uploadImage return ${response.status} ${
          response.statusText
        } ${await response.text()}`
      );
    }
  }

  async DescribeApi(image: DiscordImage, nonce?: string) {
    const payload = await this.describePayload(image, nonce);
    return this.safeIteractions(payload);
  }
  async upImageApi(image: DiscordImage, nonce?: string) {
    const { SalaiToken, DiscordBaseUrl, ChannelId, fetch } = this.config;
    const payload = {
      content: "",
      nonce,
      channel_id: ChannelId,
      type: 0,
      sticker_ids: [],
      attachments: [image],
    };

    const url = new URL(
      `${DiscordBaseUrl}/api/v9/channels/${ChannelId}/messages`
    );
    const headers = {
      Authorization: SalaiToken,
      "content-type": "application/json",
    };
    const response = await fetch(url, {
      headers,
      method: "POST",
      body: JSON.stringify(payload),
    });

    return response.status;
  }
}