Skip to content

Commit

Permalink
feat: support upscale
Browse files Browse the repository at this point in the history
  • Loading branch information
lvqq committed May 21, 2023
1 parent 6c6a6f9 commit a57d8bc
Show file tree
Hide file tree
Showing 8 changed files with 180 additions and 26 deletions.
29 changes: 27 additions & 2 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
Fetch api for midjourney on discord

## Usage
### imagine
```typescript
import { Midjourney } from 'midjourney-fetch'

Expand All @@ -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
Expand Down
5 changes: 4 additions & 1 deletion package.json
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
{
"name": "midjourney-fetch",
"version": "0.1.4",
"version": "1.0.0-beta.0",
"description": "",
"type": "module",
"main": "./dist/index.js",
Expand Down Expand Up @@ -55,5 +55,8 @@
"eslint --fix --quiet",
"prettier --write"
]
},
"dependencies": {
"@sapphire/snowflake": "^3.5.1"
}
}
10 changes: 10 additions & 0 deletions pnpm-lock.yaml

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

8 changes: 8 additions & 0 deletions src/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
};
31 changes: 31 additions & 0 deletions src/interface.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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[];
Expand All @@ -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;
}
104 changes: 85 additions & 19 deletions src/midjourney.ts
Original file line number Diff line number Diff line change
@@ -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';

Expand Down Expand Up @@ -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: [
Expand All @@ -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,
Expand All @@ -85,6 +87,7 @@ export class Midjourney {
},
attachments: [],
},
nonce: DiscordSnowflake.generate().toString(),
};

const res = await fetch(`https://discord.com/api/v9/interactions`, {
Expand All @@ -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`,
{
Expand All @@ -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;
}
Expand All @@ -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;
}
}
17 changes: 14 additions & 3 deletions src/utils.ts
Original file line number Diff line number Diff line change
@@ -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
);
};

Expand Down
2 changes: 1 addition & 1 deletion tsup.config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
});

0 comments on commit a57d8bc

Please sign in to comment.