From a57d8bc6a79636a435e8ebc492114f46a69ed2a8 Mon Sep 17 00:00:00 2001 From: lvqq Date: Sun, 21 May 2023 19:39:10 +0800 Subject: [PATCH] feat: support upscale --- README.md | 29 ++++++++++++- package.json | 5 ++- pnpm-lock.yaml | 10 +++++ src/config.ts | 8 ++++ src/interface.ts | 31 ++++++++++++++ src/midjourney.ts | 104 +++++++++++++++++++++++++++++++++++++--------- src/utils.ts | 17 ++++++-- tsup.config.ts | 2 +- 8 files changed, 180 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index d511151..4917114 100644 --- a/README.md +++ b/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 diff --git a/package.json b/package.json index d1b8aa7..2fd18f6 100644 --- a/package.json +++ b/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" } } \ No newline at end of file diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 6b8548b..4371bb9 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1,5 +1,10 @@ lockfileVersion: '6.0' +dependencies: + '@sapphire/snowflake': + specifier: ^3.5.1 + version: 3.5.1 + devDependencies: '@typescript-eslint/eslint-plugin': specifier: ^5.59.1 @@ -356,6 +361,11 @@ packages: fastq: 1.15.0 dev: true + /@sapphire/snowflake@3.5.1: + resolution: {integrity: sha512-BxcYGzgEsdlG0dKAyOm0ehLGm2CafIrfQTZGWgkfKYbj+pNNsorZ7EotuZukc2MT70E0UbppVbtpBrqpzVzjNA==} + engines: {node: '>=v14.0.0', npm: '>=7.0.0'} + dev: false + /@types/chai-subset@1.3.3: resolution: {integrity: sha512-frBecisrNGz+F4T6bcc+NLeolfiojh5FxW2klu669+8BARtyQv2C/GkNW6FUodVe4BroGMP/wER/YDGc7rEllw==} dependencies: diff --git a/src/config.ts b/src/config.ts index a7bc335..fcccd1f 100644 --- a/src/config.ts +++ b/src/config.ts @@ -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', +}; diff --git a/src/interface.ts b/src/interface.ts index 1cac5d0..9bdd242 100644 --- a/src/interface.ts +++ b/src/interface.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; + index: number; + } + | { + type?: Extract; + }; + +export interface UpscaleProps { + messageId: string; + index: number; + hash?: string; + customId?: string; } diff --git a/src/midjourney.ts b/src/midjourney.ts index b1b2ba2..fc97574 100644 --- a/src/midjourney.ts +++ b/src/midjourney.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; } } diff --git a/src/utils.ts b/src/utils.ts index 6d0b541..82253bd 100644 --- a/src/utils.ts +++ b/src/utils.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 ); }; diff --git a/tsup.config.ts b/tsup.config.ts index 2a4e890..c987328 100644 --- a/tsup.config.ts +++ b/tsup.config.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, });