@mymj/midjourney
Version:
Node.js client for the unofficial MidJourney API.
887 lines • 31.3 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.WsMessage = void 0;
const tslib_1 = require("tslib");
const utils_1 = require("./utils");
const verify_human_1 = require("./verify.human");
const node_util_1 = require("node:util");
const zlib = tslib_1.__importStar(require("zlib-sync"));
class WsMessage {
config;
MJApi;
ws;
closed = false;
event = [];
waitMjEvents = new Map();
skipMessageId = [];
reconnectTime = [];
// private heartbeatTask: NodeJS.Timer | null = null
lastSequence = 0;
inflate = null;
textDecoder = new node_util_1.TextDecoder();
UserId = "";
connecting = false;
constructor(config, MJApi) {
this.config = config;
this.MJApi = MJApi;
this.ws = new this.config.WebSocket(this.config.WsBaseUrl);
this.inflate = new zlib.Inflate({
chunkSize: 65535,
to: 'string',
});
this.connecting = true;
this.ws.addEventListener("open", this.open.bind(this));
this.onSystem("messageCreate", this.onMessageCreate.bind(this));
this.onSystem("messageUpdate", this.onMessageUpdate.bind(this));
this.onSystem("messageDelete", this.onMessageDelete.bind(this));
this.onSystem("ready", this.onReady.bind(this));
this.onSystem("interactionSuccess", this.onInteractionSuccess.bind(this));
this.onSystem("modalCreate", this.onModalCreate.bind(this));
}
async heartbeat(num) {
if (this.reconnectTime[num])
return;
//check if ws is closed
if (this.closed)
return;
if (this.ws.readyState !== this.ws.OPEN) {
this.reconnect();
return;
}
this.log("heartbeat", this.lastSequence);
this.ws.send(JSON.stringify({
op: 1,
d: this.lastSequence,
}));
await this.timeout(1000 * 40);
this.heartbeat(num);
}
// private heartbeat(interval: number) {
// const nextInterval = interval * Math.random();
// this.log(`heartbeat`, `send discord heartbeat after ${Math.round(nextInterval / 1000)}s`);
// if (this.closed) return;
// this.heartbeatTask = setTimeout(() => {
// if (this.ws.readyState === WebSocket.OPEN) {
// this.ws.send(
// JSON.stringify({
// op: 1,
// d: this.lastSequence,
// }),
// );
// this.heartbeat(interval);
// } else {
// this.reconnect();
// }
// }, nextInterval);
// }
close() {
this.closed = true;
this.ws.close();
}
async checkWs() {
if (this.closed)
return;
if (this.ws.readyState !== this.ws.OPEN) {
this.reconnect();
await this.onceReady();
}
}
async isReady() {
if (this.closed)
return false;
if (this.connecting) {
await this.onceReady();
}
else {
await this.checkWs();
}
return true;
}
async onceReady() {
return new Promise((resolve) => {
this.once("ready", (user) => {
//print user nickname
console.log(`🎊 ws ready!!! Hi: ${user.global_name}`);
resolve(this);
});
});
}
//try reconnect
reconnect() {
if (this.closed)
return;
this.inflate = new zlib.Inflate({
chunkSize: 65535,
to: 'string',
});
this.ws = new this.config.WebSocket(this.config.WsBaseUrl);
this.connecting = true;
this.lastSequence = 0;
// clear previours heartbeat
// if (this.heartbeatTask && typeof this.heartbeatTask === 'number') {
// clearInterval(this.heartbeatTask)
// this.heartbeatTask = null
// }
this.ws.addEventListener("open", this.open.bind(this));
}
decodeMessage(data) {
if (this.inflate) {
const decompressable = new Uint8Array(data);
const l = decompressable.length;
const flush = l >= 4 &&
decompressable[l - 4] === 0x00 &&
decompressable[l - 3] === 0x00 &&
decompressable[l - 2] === 0xff &&
decompressable[l - 1] === 0xff;
this.inflate.push(Buffer.from(decompressable), flush ? zlib.Z_SYNC_FLUSH : zlib.Z_NO_FLUSH);
if (this.inflate.err) {
this.log(`${this.inflate.err}${this.inflate.msg ? `: ${this.inflate.msg}` : ''}`);
return null;
}
if (!flush) {
return null;
}
const { result } = this.inflate;
if (!result) {
return null;
}
try {
return JSON.parse(typeof result === 'string' ? result : this.textDecoder.decode(result));
}
catch (error) {
return null;
}
}
else {
return data;
}
}
// After opening ws
async open() {
const num = this.reconnectTime.length;
this.log("open.time", num);
this.reconnectTime.push(false);
this.ws.addEventListener("message", (event) => {
const res = this.decodeMessage(event.data);
this.parseMessage(res);
});
this.ws.addEventListener("error", (event) => {
this.reconnectTime[num] = true;
this.reconnect();
});
this.ws.addEventListener("close", (event) => {
this.reconnectTime[num] = true;
this.reconnect();
});
this.connecting = false;
this.auth();
setTimeout(() => {
this.heartbeat(num);
}, 1000 * 10);
}
// auth
auth() {
this.ws.send(JSON.stringify({
op: 2,
d: {
token: this.config.SalaiToken,
capabilities: 8189,
properties: {
os: "Mac OS X",
browser: "Chrome",
device: "",
},
compress: false,
},
}));
}
async timeout(ms) {
return new Promise((resolve) => setTimeout(resolve, ms));
}
async messageCreate(message) {
const { embeds, id, nonce, components, attachments, content } = message;
const hash = (0, utils_1.componentsToHash)(components);
const contentNonce = this.getNonceFromContent(message.content);
let isJobQueued = false;
if (nonce) {
// this.log("waiting start image or info or error");
// there could be components with 'cancel job' but sometime not
this.updateMjEventIdByNonce(nonce, id);
if (hash) {
this.updateHashByNonce(nonce, hash);
}
if (embeds?.[0]) {
const { color, description, title } = embeds[0];
switch (color) {
case 16711680: //error
if (title === "Action needed to continue") {
if (!description.includes('Our AI moderators feel your prompt might be against our community standards.')) {
return this.continue(message);
}
}
else if (title === "Pending mod message") {
return this.continue(message);
}
else if (title === "Tos not accepted") {
return this.continue(message);
}
const error = new Error(description);
this.EventError(id, error);
return;
case 16776960: //warning
console.warn(description);
break;
case 16239475: // rich
if (title === "Job queued") {
isJobQueued = true;
}
break;
default:
if (title?.includes("continue") &&
description?.includes("verify you're human")) {
//verify human
await this.verifyHuman(message);
return;
}
if (title?.includes("Invalid")) {
//error
const error = new Error(description);
this.EventError(id, error);
return;
}
}
}
if (content === 'Failed to process your command :c') {
this.EventError(id, new Error(content));
return;
}
if (content.includes('Are you sure you want to imagine') && components?.length) {
this.EventError(id, new Error('Permutation is not supported yet, please try again one by one.'));
return;
}
}
else {
// some case error message create new message which do not have nonce
// use message_reference to match previours message
if (embeds?.[0]) {
const { description, title } = embeds[0];
const { message_reference } = message;
if (title?.includes("Invalid")) {
const refId = message_reference?.message_id;
if (refId) {
//error
const error = new Error(description);
this.EventError(refId, error);
return;
}
}
}
// vary region create is special without nonce
if (contentNonce) {
const event = this.getEventByNonce(contentNonce);
if (event) {
this.updateMjEventIdByNonce(contentNonce, id);
if (hash) {
this.updateHashByNonce(contentNonce, hash);
}
}
// } else {
// const event = this.getEventById(id);
// if (!event && hash) {
// event = this.getEventByHash(hash);
// }
// // somecase MJ throw finished image directly.
// if (!event) {
// event = this.getEventByContent(message.content);
// }
// if (event && !event.hash) {
// this.updateMjEventIdByNonce(contentNonce, id);
// if (hash) {
// this.updateHashByNonce(contentNonce, hash);
// }
// }
}
}
if (hash && attachments?.length > 0 && components?.length > 0) {
this.done(message);
return;
}
this.log(`isJobQueued: ${isJobQueued}`);
if (isJobQueued) {
this.processingImage(message, isJobQueued);
}
else {
this.messageUpdate(message);
}
}
messageUpdate(message) {
// this.log("messageUpdate", message);
const { content, embeds, interaction = {}, nonce, id, components, attachments, } = message;
if (!nonce) {
const { name } = interaction;
switch (name) {
case "settings":
this.emit("settings", message);
return;
case "describe":
let uri = embeds?.[0]?.image?.url;
if (this.config.ImageProxy !== "") {
uri = uri.replace("https://cdn.discordapp.com/", this.config.ImageProxy);
}
const describe = {
id: id,
flags: message.flags,
descriptions: embeds?.[0]?.description.split("\n\n"),
uri: uri,
proxy_url: embeds?.[0]?.image?.proxy_url,
options: (0, utils_1.formatOptions)(components),
};
this.emitMJ(id, describe);
break;
case "prefer remix":
if (content != "") {
this.emit("prefer-remix", content);
}
break;
case "shorten":
const shorten = {
description: embeds?.[0]?.description,
prompts: (0, utils_1.formatPrompts)(embeds?.[0]?.description),
options: (0, utils_1.formatOptions)(components),
id,
flags: message.flags,
};
this.emitMJ(id, shorten);
break;
case "info":
this.emit("info", embeds?.[0]?.description);
case "subscribe":
this.emit("subscribe", components?.[0]?.components?.[0]);
return;
}
}
if (embeds?.[0]) {
// message is not able to continue
if (!attachments || !attachments.length || !components || !components.length) {
const { description, title, color } = embeds[0];
switch (color) {
case 16711680: //error
const error = new Error(description);
this.EventError(id, error);
return;
default:
break;
}
if (title === "Duplicate images detected") {
const error = new Error(description);
this.EventError(id, error);
return;
}
}
}
if (content) {
this.processingImage(message);
}
}
//interaction success
async onInteractionSuccess({ nonce, id, }) {
// this.log("interactionSuccess", nonce, id);
const event = this.getEventByNonce(nonce);
if (!event) {
return;
}
// event.onmodal && event.onmodal(nonce, id);
}
//modal create
async onModalCreate({ nonce, id, custom_id, components, }) {
this.log("modalCreate", nonce, id);
const event = this.getEventByNonce(nonce);
if (!event) {
return;
}
const prompt_custom_id = components[0]?.components[0]?.custom_id;
event.onmodal && event.onmodal(nonce, id, custom_id, prompt_custom_id);
}
async onReady(user) {
this.UserId = user.id;
}
async onMessageCreate(message) {
const { channel_id, author, interaction } = message;
if (channel_id !== this.config.ChannelId)
return;
if (author?.id !== this.config.BotId)
return;
if (interaction && interaction.user.id !== this.UserId)
return;
// this.log("[messageCreate]", JSON.stringify(message));
this.messageCreate(message);
}
async onMessageUpdate(message) {
const { channel_id, author, interaction } = message;
if (channel_id !== this.config.ChannelId)
return;
if (author?.id !== this.config.BotId)
return;
if (interaction && interaction.user.id !== this.UserId)
return;
// this.log("[messageUpdate]", JSON.stringify(message));
this.messageUpdate(message);
}
async onMessageDelete(message) {
const { channel_id, id } = message;
if (channel_id !== this.config.ChannelId)
return;
// @ts-ignore
for (const [key, value] of this.waitMjEvents.entries()) {
if (value.id === id) {
this.waitMjEvents.set(key, { ...value, del: true });
}
}
}
// parse message from ws
parseMessage(msg) {
const operate = msg.op;
const type = msg.t;
if (!type) {
this.log("event", JSON.stringify(msg));
return;
}
else {
this.log("event", type);
}
if (!isNaN(msg.s)) {
this.lastSequence = msg.s;
}
// if (operate === 10 && msg?.d?.heartbeat_interval) {
// this.heartbeat(msg.d.heartbeat_interval!);
// }
const message = msg.d;
if (message.channel_id === this.config.ChannelId) {
this.log("message", JSON.stringify(msg));
}
else if (!type.includes("MESSAGE_") && type !== 'READY') {
this.log("operations", msg.t);
}
switch (type) {
case "READY":
if (message.session_id) {
this.config.ActiveSessionId = message.session_id;
}
this.emitSystem("ready", message.user);
break;
case "INTERACTION_IFRAME_MODAL_CREATE":
if (message.nonce) {
this.emit(message.nonce, message);
}
break;
case "INTERACTION_MODAL_CREATE":
if (message.nonce) {
this.emitSystem("modalCreate", message);
}
break;
case "MESSAGE_CREATE":
this.emitSystem("messageCreate", message);
break;
case "MESSAGE_UPDATE":
this.emitSystem("messageUpdate", message);
break;
case "MESSAGE_DELETE":
this.emitSystem("messageDelete", message);
break;
case "INTERACTION_SUCCESS":
this.emitSystem("interactionSuccess", message);
break;
case "INTERACTION_CREATE":
if (message.nonce) {
this.emitSystem("interactionCreate", message);
}
break;
default:
break;
}
}
//continue click appeal or Acknowledged
async continue(message) {
const { components, id, flags, nonce } = message;
const appeal = components[0]?.components[0];
this.log("appeal", appeal);
if (appeal) {
var newnonce = (0, utils_1.nextNonce)();
const httpStatus = await this.MJApi.CustomApi({
msgId: id,
customId: appeal.custom_id,
flags,
nonce: newnonce,
});
this.log("appeal.httpStatus", httpStatus);
if (httpStatus == 204) {
//todo
this.log("new nonce", newnonce);
this.on(newnonce, (data) => {
this.log("nonce data", data);
this.emit(nonce, data);
});
}
}
}
async verifyHuman(message) {
const { HuggingFaceToken } = this.config;
if (HuggingFaceToken === "" || !HuggingFaceToken) {
this.log("HuggingFaceToken is empty");
return;
}
const { embeds, components, id, flags, nonce } = message;
const uri = embeds[0].image.url;
const categories = components[0].components;
const classify = categories.map((c) => c.label);
const verifyClient = new verify_human_1.VerifyHuman(this.config);
const category = await verifyClient.verify(uri, classify);
if (category) {
const custom_id = categories.find((c) => c.label === category).custom_id;
var newnonce = (0, utils_1.nextNonce)();
const httpStatus = await this.MJApi.CustomApi({
msgId: id,
customId: custom_id,
flags,
nonce: newnonce,
});
if (httpStatus == 204) {
this.on(newnonce, (data) => {
this.emit(nonce, data);
});
}
this.log("verifyHumanApi", httpStatus, custom_id, message.id);
}
}
EventError(id, error) {
const event = this.getEventById(id);
if (!event) {
return;
}
const eventMsg = {
error,
};
this.log("MJ ERROR:", error.message);
this.emit(event.nonce, eventMsg);
}
done(message) {
const { content, id, attachments, components, flags } = message;
const { url, proxy_url, width, height } = attachments[0];
let uri = url;
if (this.config.ImageProxy !== "") {
uri = uri.replace("https://cdn.discordapp.com/", this.config.ImageProxy);
}
let hash;
if (components && components.length) {
hash = (0, utils_1.componentsToHash)(components);
}
if (!hash) {
hash = (0, utils_1.uriToHash)(url);
}
const MJmsg = {
id,
flags,
content,
hash: (0, utils_1.uriToHash)(url),
progress: "done",
uri,
proxy_url,
options: (0, utils_1.formatOptions)(components),
width,
height,
};
this.filterMessages(MJmsg);
return;
}
processingImage(message, isJobQueued) {
const { content, id, attachments, flags, components } = message;
let event;
let hash;
let uri, proxy_url, width, height;
if (components && components.length) {
hash = (0, utils_1.componentsToHash)(components);
}
if (attachments && attachments.length) {
const { url, filename } = attachments[0];
uri = url;
proxy_url = attachments[0].proxy_url;
width = attachments[0].width;
height = attachments[0].height;
if (this.config.ImageProxy !== "") {
uri = uri.replace("https://cdn.discordapp.com/", this.config.ImageProxy);
}
if (!hash) {
hash = (0, utils_1.filenameToHash)(filename);
}
}
if (hash) {
// found same hash event
event = this.getEventByHash(hash);
this.log('find event by hash', event);
}
// track event by msg id for imagine at the begining.
if (!event) {
event = this.getEventById(id);
this.log('find old event by id', event);
}
if (event && content) {
event.prompt = content;
// update hash for the first time, in imagine case
if (hash && !event.hash) {
this.log('update event by hash', event.id, hash);
this.updateHashByid(id, hash);
event.hash = hash;
}
}
// while MJ queued, this will tell the status
const progress = isJobQueued ? 'Job queued' : (0, utils_1.content2progress)(content);
const MJmsg = {
uri: uri,
proxy_url: proxy_url,
content: content,
flags: flags,
hash,
options: components && components.length && (0, utils_1.formatOptions)(components),
progress,
width,
height,
};
const eventMsg = {
message: MJmsg,
};
if (event) {
this.emitImage(event.nonce, eventMsg);
}
else {
this.emitImage('OTHER_MSG', eventMsg);
}
}
async filterMessages(MJmsg) {
// delay 300ms for discord message delete
await this.timeout(300);
let event;
if (MJmsg.hash) {
event = this.getEventByHash(MJmsg.hash);
}
// somecase MJ throw finished image directly.
if (!event) {
event = this.getEventByContent(MJmsg.content);
}
const eventMsg = {
message: MJmsg,
};
if (!event) {
this.log("FilterMessages not found", MJmsg.hash, this.waitMjEvents);
this.emitImage('OTHER_MSG', eventMsg);
return;
}
this.emitImage(event.nonce, eventMsg);
}
getEventByContent(content) {
const prompt = (0, utils_1.content2prompt)(content);
//fist del message
// @ts-ignore
for (const [key, value] of this.waitMjEvents.entries()) {
const prompCache = (0, utils_1.content2prompt)(value?.prompt);
if (value.del === true &&
prompCache.includes(prompt)) {
return value;
}
}
// @ts-ignore
for (const [key, value] of this.waitMjEvents.entries()) {
const prompCache = (0, utils_1.content2prompt)(value?.prompt);
if (prompCache.includes(prompt)) {
return value;
}
}
}
getEventByHash(hash) {
// @ts-ignore
for (const [key, value] of this.waitMjEvents.entries()) {
if (value.hash === hash) {
return value;
}
}
}
getEventById(id) {
// @ts-ignore
for (const [key, value] of this.waitMjEvents.entries()) {
if (value.id === id) {
return value;
}
}
}
getEventByNonce(nonce) {
// @ts-ignore
for (const [key, value] of this.waitMjEvents.entries()) {
if (value.nonce === nonce) {
return value;
}
}
}
getNonceFromContent(content) {
return content.match(/\[(.*?)\]/)?.[1];
}
updateMjEventIdByNonce(nonce, id) {
if (nonce === "" || id === "")
return;
let event = this.waitMjEvents.get(nonce);
if (!event)
return;
event.id = id;
this.log("updateMjEventIdByNonce success", event);
return event;
}
updateHashByNonce(nonce, hash) {
if (nonce === "" || hash === "")
return;
let event = this.waitMjEvents.get(nonce);
if (!event)
return;
event.hash = hash;
this.log("updateHashByNonce success", event);
return event;
}
updateHashByid(id, hash) {
if (id === "" || hash === "")
return;
let event = this.getEventById(id);
if (!event)
return;
event.hash = hash;
this.log("updateHashByMjEventId success", event);
return event;
}
async log(...args) {
this.config.Debug && console.info(...args, new Date().toISOString());
}
emit(event, message) {
this.event
.filter((e) => e.event === event)
.forEach((e) => e.callback(message));
}
emitImage(type, message) {
this.emit(type, message);
}
//FIXME: emitMJ rename
emitMJ(id, data) {
const event = this.getEventById(id);
if (!event)
return;
this.emit(event.nonce, data);
}
on(event, callback) {
this.event.push({ event, callback });
}
onSystem(event, callback) {
this.on(event, callback);
}
emitSystem(type, message) {
this.emit(type, message);
}
once(event, callback) {
const once = (message) => {
this.remove(event, once);
callback(message);
};
this.event.push({ event, callback: once });
}
remove(event, callback) {
this.event = this.event.filter((e) => e.event !== event && e.callback !== callback);
}
removeEvent(event) {
this.event = this.event.filter((e) => e.event !== event);
}
onceMJ(nonce, callback) {
const once = (message) => {
this.remove(nonce, once);
//FIXME: removeWaitMjEvent
this.removeWaitMjEvent(nonce);
callback(message);
};
//FIXME: addWaitMjEvent
this.waitMjEvents.set(nonce, { nonce });
this.event.push({ event: nonce, callback: once });
}
removeSkipMessageId(messageId) {
const index = this.skipMessageId.findIndex((id) => id !== messageId);
if (index !== -1) {
this.skipMessageId.splice(index, 1);
}
}
removeWaitMjEvent(nonce) {
this.waitMjEvents.delete(nonce);
}
onceImage(nonce, callback) {
const once = (data) => {
const { message, error } = data;
if (error || (message && message.progress === "done")) {
this.remove(nonce, once);
}
callback(data);
};
this.event.push({ event: nonce, callback: once });
}
async waitImageMessage({ nonce, prompt, onmodal, messageId, loading, }) {
if (messageId)
this.skipMessageId.push(messageId);
return new Promise((resolve, reject) => {
const handleImageMessage = ({ message, error }) => {
if (error) {
this.removeWaitMjEvent(nonce);
reject(error);
return;
}
if (message && message.progress === "done") {
this.removeWaitMjEvent(nonce);
messageId && this.removeSkipMessageId(messageId);
resolve(message);
return;
}
message && loading && loading({
uri: message.uri,
progress: message.progress || "",
options: message.options,
hash: message.hash,
});
};
this.waitMjEvents.set(nonce, {
nonce,
prompt,
onmodal: async (oldnonce, id, custom_id, prompt_custom_id) => {
if (onmodal === undefined) {
// reject(new Error("onmodal is not defined"))
return "";
}
var nonce = await onmodal(oldnonce, id, custom_id, prompt_custom_id);
if (nonce === "") {
// reject(new Error("onmodal return empty nonce"))
return "";
}
this.removeWaitMjEvent(oldnonce);
this.waitMjEvents.set(nonce, { nonce });
this.onceImage(nonce, handleImageMessage);
return nonce;
},
});
this.onceImage(nonce, handleImageMessage);
});
}
async waitOnceMJ(nonce) {
return new Promise((resolve) => {
this.onceMJ(nonce, (message) => {
resolve(message);
});
});
}
async waitOnce(nonce) {
return new Promise((resolve) => {
this.once(nonce, (message) => {
resolve(message);
});
});
}
}
exports.WsMessage = WsMessage;
//# sourceMappingURL=discord.ws.js.map