├── .env.template
├── .github
├── ISSUE_TEMPLATE
│ ├── banned_prompt_report.yml
│ └── new_issue.yml
└── workflows
│ └── docker-build.yml
├── .gitignore
├── Dockerfile
├── README-EN.md
├── README.md
├── __init__.py
├── app
├── __init__.py
├── handler.py
├── routers.py
├── schema.py
└── server.py
├── banned_words.txt
├── build.sh
├── entrypoint.sh
├── exceptions.py
├── lib
├── __init__.py
├── api
│ ├── __init__.py
│ ├── callback.py
│ └── discord.py
└── prompt.py
├── requirements.txt
├── server.py
├── start.sh
├── task
├── __init__.py
└── bot
│ ├── __init__.py
│ ├── _typing.py
│ ├── handler.py
│ └── listener.py
├── task_bot.py
└── util
├── _queue.py
└── fetch.py
/.env.template:
--------------------------------------------------------------------------------
1 | LOG_LEVEL=DEBUG
2 | USER_TOKEN=xxx
3 | BOT_TOKEN=xxx
4 | GUILD_ID=xxx
5 | CHANNEL_ID=xxx
6 | CALLBACK_URL=xxx
7 | # midjourney 并发数
8 | CONCUR_SIZE=3
9 | # 等待队列数
10 | WAIT_SIZE=10
11 | # 监听 midjourney bot 处理完任务后，回调 API 服务清除队列
12 | QUEUE_RELEASE_API=http://127.0.0.1:8062/v1/api/trigger/queue/release
13 |
14 | # discord api version
15 | DRAW_VERSION=1237876415471554623
16 | # proxy, default None
17 | PROXY_URL=
18 |
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/banned_prompt_report.yml:
--------------------------------------------------------------------------------
1 | name: "🔥 禁用词添加"
2 | description: 禁用词添加
3 | title: "Banned prompt: "
4 | labels: ["banned prompt"]
5 | body:
6 | - type: markdown
7 | attributes:
8 | value: |
9 | 非常感谢你的 issue report (๑><๑）
10 | - type: input
11 | id: prompt
12 | attributes:
13 | label: 禁用词
14 | description: 输入您需要添加的 Midjourney 禁用词～
15 | placeholder: |
16 | Prompt
17 | validations:
18 | required: true
19 | - type: checkboxes
20 | id: checkboxes
21 | attributes:
22 | label: 一点点的自我检查
23 | description: 在你提交 issue 之前，麻烦确认自己是否已经完成了以下检查：
24 | options:
25 | - label: 如果是网络问题，已经检查网络连接、设置是否正常，并经过了充分测试
26 | required: true
27 | - label: 本 prompt 在 [banned_words](https://raw.githubusercontent.com/yokonsan/midjourney-api/master/banned_words.txt) 中不存在
28 | required: true
29 |
--------------------------------------------------------------------------------
/.github/ISSUE_TEMPLATE/new_issue.yml:
--------------------------------------------------------------------------------
1 | name: "🐛 新建 issue"
2 | description: 上报一个新的 issue
3 | title: "🐛 "
4 | labels: []
5 | body:
6 | - type: markdown
7 | attributes:
8 | value: |
9 | 非常感谢你的 issue report (๑><๑），为了使我们能够更快地定位问题来源，请尽可能完整地填写本 Issue 表格
10 | - type: textarea
11 | id: description
12 | attributes:
13 | label: 问题简述
14 | description: 详述你所遇到的问题（如有报错也请粘贴在这里）～
15 | placeholder: |
16 | 如果方便，请提供更加详细的日志信息
17 | validations:
18 | required: true
19 | - type: textarea
20 | id: reproduction
21 | attributes:
22 | label: 复现方式
23 | description: |
24 | 请在这里提供你复现改问题的方法。
25 | 为了节省彼此交流的时间，麻烦在提交 issue 前多次测试该问题是能够反复复现的（非网络问题）。
26 | placeholder: "注意在粘贴的命令中隐去所有隐私信息哦(*/ω＼*)"
27 | validations:
28 | required: true
29 | - type: textarea
30 | id: environment-info
31 | attributes:
32 | label: 环境信息
33 | description: 请尽可能详细地供以下信息～
34 | placeholder: 你的环境信息～
35 | value: |
36 | - OS: 操作系统类型及其版本号
37 | - Python: Python 版本号 (`python --version`)
38 | - Others: 其它信息
39 | validations:
40 | required: true
41 | - type: textarea
42 | id: additional-context
43 | attributes:
44 | label: 额外信息
45 | description: 请尽可能提供一些你认为可能产生该问题的一些原因
46 | placeholder: 如有额外的信息，请填写在这里～
47 | validations:
48 | required: false
49 | - type: checkboxes
50 | id: checkboxes
51 | attributes:
52 | label: 一点点的自我检查
53 | description: 在你提交 issue 之前，麻烦确认自己是否已经完成了以下检查：
54 | options:
55 | - label: 充分阅读 [README.md](https://github.com/yokonsan/midjourney-api)，特别是与本 issue 相关的部分
56 | required: true
57 | - label: 如果是网络问题，已经检查网络连接、设置是否正常，并经过充分测试认为这是项目本身的问题
58 | required: true
59 | - label: 本 issue 在 [issues](https://github.com/yokonsan/midjourney-api/issues) 中并没有重复问题
60 | required: true
61 |
--------------------------------------------------------------------------------
/.github/workflows/docker-build.yml:
--------------------------------------------------------------------------------
1 | name: Docker Build CI
2 |
3 | on:
4 | workflow_dispatch:
5 | release:
6 | types: [published]
7 |
8 | jobs:
9 |
10 | build:
11 |
12 | runs-on: ubuntu-latest
13 |
14 | steps:
15 | -
16 | name: Check out the repo
17 | uses: actions/checkout@v3
18 | -
19 | name: Set up QEMU
20 | uses: docker/setup-qemu-action@v2
21 | -
22 | name: Set up Docker Buildx
23 | uses: docker/setup-buildx-action@v2
24 | -
25 | name: Login to Docker Hub
26 | uses: docker/login-action@v2
27 | with:
28 | username: ${{ secrets.DOCKERHUB_USERNAME }}
29 | password: ${{ secrets.DOCKERHUB_TOKEN }}
30 | -
31 | name: Build and push
32 | uses: docker/build-push-action@v4
33 | with:
34 | push: true
35 | tags: kunyu/midjourney-api:1.0
36 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | .cache
2 | .idea
3 | .env*
4 | log
5 | test/log
6 | test/**/log
7 | .pytest_cache
8 | .vscode
9 | *__pycache__*
10 | data
11 | test/data
12 | test/**/data
13 |
14 | # pipenv
15 | Pipfile
16 | Pipfile.lock
17 |
--------------------------------------------------------------------------------
/Dockerfile:
--------------------------------------------------------------------------------
1 | FROM python:3.10.6
2 | LABEL creator="yokon" email="944682328@qq.com"
3 |
4 | WORKDIR /app
5 |
6 | COPY . .
7 | RUN pip install --upgrade pip \
8 | && pip install -i https://pypi.douban.com/simple/ -r requirements.txt \
9 | && chmod +x entrypoint.sh
10 |
11 | ENTRYPOINT ["bash", "entrypoint.sh"]
12 | EXPOSE 8062
13 | CMD ["http"]
--------------------------------------------------------------------------------
/README-EN.md:
--------------------------------------------------------------------------------
1 | # midjourney-api
2 |
3 | English / [简体中文](./README.md)
4 |
5 | Based on the Discord Midjourney API.
6 |
7 | **Add Midjourney Banned Words Entry [issue](https://github.com/yokonsan/midjourney-api/issues/new?assignees=&labels=banned+prompt&projects=&template=banned_prompt_report.yml&title=Banned+prompt%3A+)**
8 |
9 | Refer to the project integration demo: [issue31](https://github.com/yokonsan/midjourney-api/issues/31)
10 |
11 |
12 |
13 | ## Midjourney API Recommend
14 | If the above process & code is difficult for you to operate, or maintaining a pool of accounts consumes a lot of your energy and cost. Here is a highly integrated and cost effective [Midjouney API](https://ttapi.io/midjourney) platform recommendation **[TTAPI](https://ttapi.io)** for price and stability.
15 |
16 | **The TTAPI supports functions related to Midjourney:**
17 | - imagine
18 | - U V pan zoom
19 | - describe
20 | - blend
21 | - vary_region and all other advanced features included in Midjourney.
22 | - In addition to this there is support for the **Luma API**, the **Face Swap API**, the relatively inexpensive **ChatGPT**, and the **Claude API**.
23 |
24 | **How to start with TTAPI**
25 |
26 | - Just [Sign In With Github](https://ttapi.io/login), and than you will get 30 quota of free credits (can be used to 10 times imagine endpoint in fast mode for testing).
27 | - [For complete docs of TTAPI](https://ttapi.io/docs)
28 |
29 |
30 |
31 |
32 | ## UML
33 |
34 | ```mermaid
35 | sequenceDiagram
36 | participant ThirdServer
37 | participant APIServer
38 | participant DiscordAPI
39 |
40 | ThirdServer->>APIServer: Request interface to trigger task
41 | APIServer->>APIServer: Add to task queue
42 | APIServer->>DiscordAPI: Call interface to trigger drawing task
43 | APIServer-->>ThirdServer: Return whether the trigger was successful
44 |
45 | DiscordAPI->>DiscordAPI: Trigger Midjourney bot drawing task
46 | DiscordAPI->>DiscordAPI: Monitor MidJourney bot messages
47 | DiscordAPI-->>ThirdServer: Return real-time message monitoring
48 | DiscordAPI-->>APIServer: Clear queue task
49 | ```
50 |
51 | ## Usage Requirements
52 |
53 | 1. Ensure the program's runtime environment can access Discord
54 | 2. Have Midjourney and Discord accounts
55 | 3. Create a Discord channel and add the bot, refer to the tutorial [Midjourney｜How to Integrate into Your Own Platform](https://mp.weixin.qq.com/s?__biz=Mzg4MjkzMzc1Mg==&mid=2247484029&idx=1&sn=d3c458bba9459f19f05d13ab23f5f67e&chksm=cf4e68eaf839e1fc2db025bd9940d0f5e57862f1788c88215b4a66cb23f553a30c5f37ac3ae8&token=79614426&lang=zh_CN#rd)
56 |
57 |
58 | ## Installation and Startup
59 |
60 | ```bash
61 | git clone
62 | pip install -r requirements.txt
63 | ```
64 |
65 | Rename the file `.env.template` to `.env` and fill in the parameter values:
66 |
67 | ```
68 | USER_TOKEN=User token
69 | BOT_TOKEN=Bot token
70 | GUILD_ID=Server ID
71 | CHANNEL_ID=Channel ID
72 | CALLBACK_URL=Callback URL, default is HTTP POST request, used for receiving midjourney drawing progress and results
73 | ```
74 |
75 | ### Direct Startup
76 |
77 | ```bash
78 | # Start the bot listener
79 | python task_bot.py
80 | # Start the HTTP service
81 | python server.py
82 | ```
83 |
84 | #### Update
85 |
86 | ```bash
87 | git pull
88 |
89 | # Start the bot listener
90 | python task_bot.py
91 | # Start the HTTP service
92 | python server.py
93 | ```
94 |
95 | ### Docker Startup
96 |
97 | Fill in the environment variables after `-e` in [start.sh](./start.sh) and start directly:
98 |
99 | ```bash
100 | sh start.sh
101 | ```
102 |
103 | Or build the image locally:
104 |
105 | ```bash
106 | # Build the image
107 | sh build.sh
108 | # Start the container
109 | sh start.sh
110 | ```
111 |
112 | #### Update
113 |
114 | ```bash
115 | docker rmi kunyu/midjourney-api:1.0
116 | sh start.sh
117 | ```
118 |
119 | API `swagger` documentation: [http://127.0.0.1:8062/docs](http://127.0.0.1:8062/docs)
120 |
121 | `midjourney-api` provides interfaces:
122 |
123 | - [x] `/v1/api/trigger/imagine`: Trigger drawing task (image-to-image, add image link before the prompt)
124 | - [x] `/v1/api/trigger/upscale`: U
125 | - [x] `/v1/api/trigger/variation`: V
126 | - [x] `/v1/api/trigger/solo_variation`: Make Variations
127 | - [x] `/v1/api/trigger/solo_low_variation`: Vary(Subtle)
128 | - [x] `/v1/api/trigger/solo_high_variation`: Vary(Strong)
129 | - [x] `/v1/api/trigger/zoomout`: Zoom Out 2x/1.5x
130 | - [x] `/v1/api/trigger/expand`: ⬅️ ➡️ ⬆️ ⬇️
131 | - [x] `/v1/api/trigger/reset`: Redraw
132 | - [x] `/v1/api/trigger/upload`: Upload image
133 | - [x] `/v1/api/trigger/describe`: Generate prompt by uploading image name
134 | - [x] `/v1/api/trigger/message`: Send image message, return image link for image-to-image functionality
135 |
136 |
137 | ## Usage
138 |
139 | ### Imagine
140 |
141 | Text-to-image
142 |
143 | ```bash
144 | curl -X 'POST' \
145 | 'http://127.0.0.1:8062/v1/api/trigger/imagine' \
146 | -H 'accept: application/json' \
147 | -H 'Content-Type: application/json' \
148 | -d '{
149 | "prompt": "a cute cat"
150 | }'
151 | ```
152 |
153 | Image-to-image, must include the image URL
154 |
155 | ```bash
156 | curl -X 'POST' \
157 | 'http://127.0.0.1:8062/v1/api/trigger/imagine' \
158 | -H 'accept: application/json' \
159 | -H 'Content-Type: application/json' \
160 | -d '{
161 | "prompt": "a cute cat",
162 | "picurl": "https://xxxxxx/xxxxxxxxxxxx.jpg"
163 | }'
164 | ```
165 |
166 | ### Upscale
167 |
168 | ```bash
169 | curl -X 'POST' \
170 | 'http://127.0.0.1:8062/v1/api/trigger/upscale' \
171 | -H 'accept: application/json' \
172 | -H 'Content-Type: application/json' \
173 | -d '{
174 | "index": 1,
175 | "msg_id": "xxxxxxxxxx",
176 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
177 | "trigger_id": "xxxxxxxxxx"
178 | }'
179 | ```
180 |
181 | - `index`: Image index, values: 1, 2, 3, 4
182 | - `msg_id`: `imagine` callback message `id` field after drawing completion
183 | - `msg_hash`: `imagine` callback message `attachments[0].filename.split("_")[-1].split(".")[0]` after drawing completion
184 | - `trigger_id`: `imagine` callback message `trigger_id` field after drawing completion
185 |
186 | ### Variation
187 |
188 | ```bash
189 | curl -X 'POST' \
190 | 'http://127.0.0.1:8062/v1/api/trigger/variation' \
191 | -H 'accept: application/json' \
192 | -H 'Content-Type: application/json' \
193 | -d '{
194 | "index": 2,
195 | "msg_id": "xxxxxxxxxx",
196 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
197 | "trigger_id": "xxxxxxxxxx"
198 | }'
199 | ```
200 |
201 | ### Solo Variation
202 |
203 | Perform "Make Variations" on a single image from `upscale`
204 |
205 | ```bash
206 | curl -X 'POST' \
207 | 'http://127.0.0.1:8062/v1/api/trigger/solo_variation' \
208 | -H 'accept: application/json' \
209 | -H 'Content-Type: application/json' \
210 | -d '{
211 | "index": 1,
212 | "msg_id": "xxxxxxxxxx",
213 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
214 | "trigger_id": "xxxxxxxxxx"
215 | }'
216 | ```
217 |
218 | - `index`: Image index, not used here, value: 1
219 | - `msg_id`: `upscale` callback message `id` field after drawing completion
220 | - `msg_hash`: `upscale` callback message `attachments[0].filename.split("_")[-1].split(".")[0]` after drawing completion
221 | - `trigger_id`: `upscale` callback message `trigger_id` field after drawing completion
222 |
223 | ### Solo Low Variation
224 |
225 | Perform "Vary(Subtle)" on a single image from `upscale`
226 |
227 | ```bash
228 | curl -X 'POST' \
229 | 'http://127.0.0.1:8062/v1/api/trigger/solo_low_variation' \
230 | -H 'accept: application/json' \
231 | -H 'Content-Type: application/json' \
232 | -d '{
233 | "index": 1,
234 | "msg_id": "xxxxxxxxxx",
235 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
236 | "trigger_id": "xxxxxxxxxx"
237 | }'
238 | ```
239 |
240 | - `index`: Image index, not used here, value: 1
241 | - `msg_id`: `upscale` callback message `id` field after drawing completion
242 | - `msg_hash`: `upscale` callback message `attachments[0].filename.split("_")[-1].split(".")[0]` after drawing completion
243 | - `trigger_id`: `upscale` callback message `trigger_id` field after drawing completion
244 |
245 | ### Solo High Variation
246 |
247 | Perform "Vary(Strong)" on a single image from `upscale`
248 |
249 | ```bash
250 | curl -X 'POST' \
251 | 'http://127.0.0.1:8062/v1/api/trigger/solo_high_variation' \
252 | -H 'accept: application/json' \
253 | -H 'Content-Type: application/json' \
254 | -d '{
255 | "index": 1,
256 | "msg_id": "xxxxxxxxxx",
257 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
258 | "trigger_id": "xxxxxxxxxx"
259 | }'
260 | ```
261 |
262 | - `index`: Image index, not used here, value: 1
263 | - `msg_id`: `upscale` callback message `id` field after drawing completion
264 | - `msg_hash`: `upscale` callback message `attachments[0].filename.split("_")[-1].split(".")[0]` after drawing completion
265 | - `trigger_id`: `upscale` callback message `trigger_id` field after drawing completion
266 |
267 |
268 | ### Zoom Out
269 |
270 | Perform Zoom Out 2x/1.5x on a single image from `upscale`
271 |
272 | ```bash
273 | curl -X 'POST' \
274 | 'http://127.0.0.1:8062/v1/api/trigger/zoomout' \
275 | -H 'accept: application/json' \
276 | -H 'Content-Type: application/json' \
277 | -d '{
278 | "msg_id": "xxxxxxxxxx",
279 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
280 | "zoomout": 50
281 | "trigger_id": "xxxxxxxxxx"
282 | }'
283 | ```
284 |
285 | - `zoomout`: Image enlargement (Outpaint) factor, 2x -> 50, 1.5x -> 75
286 |
287 |
288 | ### Expand
289 |
290 | Perform expansion in a specific direction on a single image from `upscale`
291 |
292 | ```bash
293 | curl -X 'POST' \
294 | 'http://127.0.0.1:8062/v1/api/trigger/expand' \
295 | -H 'accept: application/json' \
296 | -H 'Content-Type: application/json' \
297 | -d '{
298 | "msg_id": "xxxxxxxxxx",
299 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
300 | "direction": "up"
301 | "trigger_id": "xxxxxxxxxx"
302 | }'
303 | ```
304 |
305 | - `direction`: Image expansion direction, values: left/right/up/down
306 |
307 |
308 | ### Reset
309 |
310 | ```bash
311 | curl -X 'POST' \
312 | 'http://127.0.0.1:8062/v1/api/trigger/reset' \
313 | -H 'accept: application/json' \
314 | -H 'Content-Type: application/json' \
315 | -d '{
316 | "msg_id": "xxxxxxxxxx",
317 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
318 | "trigger_id": "xxxxxxxxxx"
319 | }'
320 | ```
321 |
322 | ### Describe
323 |
324 | 1. First, upload the image
325 |
326 | ```bash
327 | curl -X 'POST' \
328 | 'http://127.0.0.1:8062/v1/api/trigger/upload' \
329 | -H 'accept: application/json' \
330 | -H 'Content-Type: multipart/form-data' \
331 | -F 'file=@cH16Ifh.jpg;type=image/jpeg'
332 | ```
333 |
334 | 2. Based on the returned image filename, call describe
335 |
336 | ```bash
337 | curl -X 'POST' \
338 | 'http://127.0.0.1:8062/v1/api/trigger/describe' \
339 | -H 'accept: application/json' \
340 | -H 'Content-Type: application/json' \
341 | -d '{
342 | "upload_filename": "b56ca21a-5fbe-40b4-89ab-6e0aa732f561/9231228408.jpg",
343 | "trigger_id": "9231228408"
344 | }'
345 | ```
346 |
347 | - `trigger_id`: Use the trigger_id returned by upload first
348 | - `upload_filename`: Filename returned by upload
349 |
350 | ### Message
351 |
352 | Same as `describe`, first `/v1/api/trigger/upload` to upload the image, then send the message based on the returned filename:
353 |
354 | ```bash
355 | curl -X 'POST' \
356 | 'http://127.0.0.1:8062/v1/api/trigger/message' \
357 | -H 'accept: application/json' \
358 | -H 'Content-Type: application/json' \
359 | -d '{
360 | "upload_filename": "560a1e26-36a2-4d5f-a48d-9dd877642b51/7185811546.jpg"
361 | }'
362 | ```
363 |
364 | After sending the image, an image link will be returned.
365 | This link is used in image-to-image, concatenating the prompt in the form `ImageURL Prompt`, calling `/v1/api/trigger/imagine`.
366 |
367 |
368 | ## Features
369 |
370 | - [x] imagine
371 | - [x] upscale
372 | - [x] variation
373 | - [x] solo_variation
374 | - [x] solo_low_variation
375 | - [x] solo_high_variation
376 | - [x] zoomout
377 | - [x] expand
378 | - [x] reset
379 | - [x] describe
380 | - [x] Image-to-image (obtain the link of the uploaded image)
381 | - [x] Sensitive word filtering and reporting
382 | - [x] Task queue (in-memory storage, external storage is not preferred, but can add exception handling for persistence)
383 | - [ ] tests
384 |
385 |
386 | ## enjoy it
387 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # midjourney-api
2 |
3 | 简体中文 / [English](./README-EN.md)
4 |
5 | 基于 Discord 的 Midjourney API。
6 |
7 | **添加 Midjourney 违禁词入口 [issue](https://github.com/yokonsan/midjourney-api/issues/new?assignees=&labels=banned+prompt&projects=&template=banned_prompt_report.yml&title=Banned+prompt%3A+)**
8 |
9 | 项目集成 Demo 参考：[issue31](https://github.com/yokonsan/midjourney-api/issues/31)
10 |
11 |
12 |
13 |
14 | ## Midjourney API 平台推荐
15 | 如果以上流程&代码对你来说操作有困难，或者维护账户池消耗了你们大量的精力与成本。这里有一个高度集成且价格与稳定性具有超高性价比的 [Midjouney API](https://ttapi.io/midjourney) 平台推荐 **[TTAPI](https://ttapi.io)**
16 |
17 | **TTAPI 平台对于Midjourney相关功能支持：**
18 | - imagine
19 | - U V pan zoom
20 | - describe
21 | - blend
22 | - vary_region等等所有Midjourney官方包含的高级功能
23 | - 除此以外还支持**Luma API**、**换脸API**、相对价格比较优惠的**ChatGPT以及Claude API**
24 |
25 | **TTAPI 接入**
26 |
27 | - 进入[平台github一键登录](https://ttapi.io/login)即可获取30配额的免费额度（可用于fast模式下执行10次imagine接口用作测试）
28 | - [平台接口文档地址](https://docs-zh.mjapiapp.com/)
29 |
30 |
31 |
32 |
33 | ## UML
34 |
35 | ```mermaid
36 | sequenceDiagram
37 | participant ThirdServer
38 | participant APIServer
39 | participant DiscordAPI
40 |
41 | ThirdServer->>APIServer: 请求接口触发任务
42 | APIServer->>APIServer: 放入任务队列
43 | APIServer->>DiscordAPI: 调接口触发绘画任务
44 | APIServer-->>ThirdServer: 返回是否触发成功
45 |
46 | DiscordAPI->>DiscordAPI: 触发Midjourney bot绘画任务
47 | DiscordAPI->>DiscordAPI: 监听MidJourney bot消息
48 | DiscordAPI-->>ThirdServer: 返回监听实时消息
49 | DiscordAPI-->>APIServer: 清除队列任务
50 | ```
51 |
52 |
53 | ## 使用条件
54 |
55 | 1. 确保程序启动环境能访问 Discord
56 | 2. 已有 Midjourney、Discord 账户
57 | 3. 创建 Discord 频道并添加机器人，参考教程 [Midjourney｜如何集成到自己的平台](https://mp.weixin.qq.com/s?__biz=Mzg4MjkzMzc1Mg==&mid=2247484029&idx=1&sn=d3c458bba9459f19f05d13ab23f5f67e&chksm=cf4e68eaf839e1fc2db025bd9940d0f5e57862f1788c88215b4a66cb23f553a30c5f37ac3ae8&token=79614426&lang=zh_CN#rd)
58 |
59 |
60 | ## 安装启动
61 |
62 | ```bash
63 | git clone
64 | pip install -r requirements.txt
65 | ```
66 |
67 | 将文件`.env.template`重命名为`.env`，并填入参数值：
68 |
69 | ```
70 | USER_TOKEN=用户token
71 | BOT_TOKEN=机器人token
72 | GUILD_ID=服务器ID
73 | CHANNEL_ID=频道ID
74 | CALLBACK_URL=回调地址，默认http post请求，用于接收 midjourney 作图进度和结果的服务
75 | ```
76 |
77 | ### 直接启动
78 |
79 | ```bash
80 | # 启动监听机器人
81 | python task_bot.py
82 | # 启动http服务
83 | python server.py
84 | ```
85 |
86 | #### 更新
87 |
88 | ```bash
89 | git pull
90 |
91 | # 启动监听机器人
92 | python task_bot.py
93 | # 启动http服务
94 | python server.py
95 | ```
96 |
97 | ### docker 启动
98 |
99 | 填写 [start.sh](./start.sh) 中 `-e` 后的环境变量，直接启动：
100 |
101 | ```bash
102 | sh start.sh
103 | ```
104 |
105 | 或者本地构建镜像：
106 |
107 | ```bash
108 | # 构建镜像
109 | sh build.sh
110 | # 启动容器
111 | sh start.sh
112 | ```
113 |
114 | #### 更新
115 |
116 | ```bash
117 | docker rmi kunyu/midjourney-api:1.0
118 | sh start.sh
119 | ```
120 |
121 | 接口`swagger`文档：[http://127.0.0.1:8062/docs](http://127.0.0.1:8062/docs)
122 |
123 | `midjourney-api` 提供接口：
124 |
125 | - [x] `/v1/api/trigger/imagine`：触发绘画任务（图生图，Prompt 前加上图片链接即可）
126 | - [x] `/v1/api/trigger/upscale`：U
127 | - [x] `/v1/api/trigger/variation`：V
128 | - [x] `/v1/api/trigger/solo_variation`：Make Variations
129 | - [x] `/v1/api/trigger/solo_low_variation`：Vary(Subtle)
130 | - [x] `/v1/api/trigger/solo_high_variation`：Vary(Strong)
131 | - [x] `/v1/api/trigger/zoomout`：Zoom Out 2x/1.5x
132 | - [x] `/v1/api/trigger/expand`：⬅️ ➡️ ⬆️ ⬇️
133 | - [x] `/v1/api/trigger/reset`：重绘
134 | - [x] `/v1/api/trigger/upload`：上传图片
135 | - [x] `/v1/api/trigger/describe`：通过上传图片名，生成 Prompt
136 | - [x] `/v1/api/trigger/message`：发送图片消息，返回图片链接，用于图生图功能
137 |
138 |
139 | ## 使用
140 |
141 | ### imagine
142 |
143 | 文生图
144 |
145 | ```bash
146 | curl -X 'POST' \
147 | 'http://127.0.0.1:8062/v1/api/trigger/imagine' \
148 | -H 'accept: application/json' \
149 | -H 'Content-Type: application/json' \
150 | -d '{
151 | "prompt": "a cute cat"
152 | }'
153 | ```
154 |
155 | 图生图，需带上图片 URL
156 |
157 | ```bash
158 | curl -X 'POST' \
159 | 'http://127.0.0.1:8062/v1/api/trigger/imagine' \
160 | -H 'accept: application/json' \
161 | -H 'Content-Type: application/json' \
162 | -d '{
163 | "prompt": "a cute cat",
164 | "picurl": "https://xxxxxx/xxxxxxxxxxxx.jpg"
165 | }'
166 | ```
167 |
168 | ### upscale
169 |
170 | ```bash
171 | curl -X 'POST' \
172 | 'http://127.0.0.1:8062/v1/api/trigger/upscale' \
173 | -H 'accept: application/json' \
174 | -H 'Content-Type: application/json' \
175 | -d '{
176 | "index": 1,
177 | "msg_id": "xxxxxxxxxx",
178 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
179 | "trigger_id": "xxxxxxxxxx"
180 | }'
181 | ```
182 |
183 | - `index`: 图片索引，取值：1、2、3、4
184 | - `msg_id`: `imagine` 绘画完成后回调报文 `id` 字段
185 | - `msg_hash`: `imagine` 绘画完成后回调报文 `attachments[0].filename.split("_")[-1].split(".").[0]`
186 | - `trigger_id`: `imagine` 绘画完成后回调报文 `trigger_id` 字段
187 |
188 | ### variation
189 |
190 | ```bash
191 | curl -X 'POST' \
192 | 'http://127.0.0.1:8062/v1/api/trigger/variation' \
193 | -H 'accept: application/json' \
194 | -H 'Content-Type: application/json' \
195 | -d '{
196 | "index": 2,
197 | "msg_id": "xxxxxxxxxx",
198 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
199 | "trigger_id": "xxxxxxxxxx"
200 | }'
201 | ```
202 |
203 | ### solo_variation
204 |
205 | 对 `upscale` 的单张图片进行 "Make Variations" 操作
206 |
207 | ```bash
208 | curl -X 'POST' \
209 | 'http://127.0.0.1:8062/v1/api/trigger/solo_variation' \
210 | -H 'accept: application/json' \
211 | -H 'Content-Type: application/json' \
212 | -d '{
213 | "index": 1,
214 | "msg_id": "xxxxxxxxxx",
215 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
216 | "trigger_id": "xxxxxxxxxx"
217 | }'
218 | ```
219 |
220 | - `index`: 图片索引，此处无用，取值：1
221 | - `msg_id`: `upscale` 绘画完成后回调报文 `id` 字段
222 | - `msg_hash`: `upscale` 绘画完成后回调报文 `attachments[0].filename.split("_")[-1].split(".").[0]`
223 | - `trigger_id`: `upscale` 绘画完成后回调报文 `trigger_id` 字段
224 |
225 | ### solo_low_variation
226 |
227 | 对 `upscale` 的单张图片进行 "Vary(Subtle)" 操作
228 |
229 | ```bash
230 | curl -X 'POST' \
231 | 'http://127.0.0.1:8062/v1/api/trigger/solo_low_variation' \
232 | -H 'accept: application/json' \
233 | -H 'Content-Type: application/json' \
234 | -d '{
235 | "index": 1,
236 | "msg_id": "xxxxxxxxxx",
237 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
238 | "trigger_id": "xxxxxxxxxx"
239 | }'
240 | ```
241 |
242 | - `index`: 图片索引，此处无用，取值：1
243 | - `msg_id`: `upscale` 绘画完成后回调报文 `id` 字段
244 | - `msg_hash`: `upscale` 绘画完成后回调报文 `attachments[0].filename.split("_")[-1].split(".").[0]`
245 | - `trigger_id`: `upscale` 绘画完成后回调报文 `trigger_id` 字段
246 |
247 | ### solo_high_variation
248 |
249 | 对 `upscale` 的单张图片进行 "Vary(Strong)" 操作
250 |
251 | ```bash
252 | curl -X 'POST' \
253 | 'http://127.0.0.1:8062/v1/api/trigger/solo_high_variation' \
254 | -H 'accept: application/json' \
255 | -H 'Content-Type: application/json' \
256 | -d '{
257 | "index": 1,
258 | "msg_id": "xxxxxxxxxx",
259 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
260 | "trigger_id": "xxxxxxxxxx"
261 | }'
262 | ```
263 |
264 | - `index`: 图片索引，此处无用，取值：1
265 | - `msg_id`: `upscale` 绘画完成后回调报文 `id` 字段
266 | - `msg_hash`: `upscale` 绘画完成后回调报文 `attachments[0].filename.split("_")[-1].split(".").[0]`
267 | - `trigger_id`: `upscale` 绘画完成后回调报文 `trigger_id` 字段
268 |
269 |
270 | ### zoomout
271 |
272 | 对 `upscale` 的单张图片进行 Zoom Out 2x/1.5x 操作
273 |
274 | ```bash
275 | curl -X 'POST' \
276 | 'http://127.0.0.1:8062/v1/api/trigger/zoomout' \
277 | -H 'accept: application/json' \
278 | -H 'Content-Type: application/json' \
279 | -d '{
280 | "msg_id": "xxxxxxxxxx",
281 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
282 | "zoomout": 50
283 | "trigger_id": "xxxxxxxxxx"
284 | }'
285 | ```
286 |
287 | - `zoomout`: 图片扩大（Outpaint）系数，2x -> 50、1.5x -> 75
288 |
289 |
290 | ### expand
291 |
292 | 对 `upscale` 的单张图片进行某方向的扩展操作
293 |
294 | ```bash
295 | curl -X 'POST' \
296 | 'http://127.0.0.1:8062/v1/api/trigger/expand' \
297 | -H 'accept: application/json' \
298 | -H 'Content-Type: application/json' \
299 | -d '{
300 | "msg_id": "xxxxxxxxxx",
301 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
302 | "direction": "up"
303 | "trigger_id": "xxxxxxxxxx"
304 | }'
305 | ```
306 |
307 | - `direction`: 图片扩大方向，取值：left/right/up/down
308 |
309 |
310 | ### reset
311 |
312 | ```bash
313 | curl -X 'POST' \
314 | 'http://127.0.0.1:8062/v1/api/trigger/reset' \
315 | -H 'accept: application/json' \
316 | -H 'Content-Type: application/json' \
317 | -d '{
318 | "msg_id": "xxxxxxxxxx",
319 | "msg_hash": "xxxxx-xxx-xxxx-xxxx-xxxxxx",
320 | "trigger_id": "xxxxxxxxxx"
321 | }'
322 | ```
323 |
324 | ### describe
325 |
326 | 1. 先上传图片
327 |
328 | ```bash
329 | curl -X 'POST' \
330 | 'http://127.0.0.1:8062/v1/api/trigger/upload' \
331 | -H 'accept: application/json' \
332 | -H 'Content-Type: multipart/form-data' \
333 | -F 'file=@cH16Ifh.jpg;type=image/jpeg'
334 | ```
335 |
336 | 2. 根据返回的图片文件名，调用 describe
337 |
338 | ```bash
339 | curl -X 'POST' \
340 | 'http://127.0.0.1:8062/v1/api/trigger/describe' \
341 | -H 'accept: application/json' \
342 | -H 'Content-Type: application/json' \
343 | -d '{
344 | "upload_filename": "b56ca21a-5fbe-40b4-89ab-6e0aa732f561/9231228408.jpg",
345 | "trigger_id": "9231228408"
346 | }'
347 | ```
348 |
349 | - `trigger_id` 先用 upload 返回的 trigger_id
350 | - `upload_filename` upload 返回的文件名
351 |
352 | ### message
353 |
354 | 和 `describe` 一样，先 `/v1/api/trigger/upload` 上传图片，然后根据返回文件名，发送消息：
355 |
356 | ```bash
357 | curl -X 'POST' \
358 | 'http://127.0.0.1:8062/v1/api/trigger/message' \
359 | -H 'accept: application/json' \
360 | -H 'Content-Type: application/json' \
361 | -d '{
362 | "upload_filename": "560a1e26-36a2-4d5f-a48d-9dd877642b51/7185811546.jpg"
363 | }'
364 | ```
365 |
366 | 发送图片后，会返回图片链接。
367 | 该链接用于以图生图中，拼接 Prompt 形如 `图片URL Prompt`，调用 `/v1/api/trigger/imagine`。
368 |
369 |
370 | ## 功能
371 |
372 | - [x] imagine
373 | - [x] upscale
374 | - [x] variation
375 | - [x] solo_variation
376 | - [x] solo_low_variation
377 | - [x] solo_high_variation
378 | - [x] zoomout
379 | - [x] expand
380 | - [x] reset
381 | - [x] describe
382 | - [x] 图生图（获取到上传图片的链接）
383 | - [x] 敏感词过滤上报
384 | - [x] 任务队列（内存存储，不希望引入外建，可加入异常落盘）
385 | - [ ] tests
386 |
387 | ## enjoy it
388 |
--------------------------------------------------------------------------------
/__init__.py:
--------------------------------------------------------------------------------
1 | import sys
2 | from os import getenv
3 |
4 | from dotenv import load_dotenv
5 | from loguru import logger
6 |
7 | # env config load
8 | load_dotenv()
9 |
10 | # logger config
11 | _lvl = getenv("LOG_LEVEL", default="ERROR")
12 | _format = '{time:%Y-%m-%d %H:%M:%S} | ' + \
13 | '{level} | ' + \
14 | '{process.id}-{thread.id} | ' + \
15 | '"{file.path}:{line}":{function} ' + \
16 | '- {message}'
17 | logger.remove()
18 | logger.add(
19 | sys.stdout, level=_lvl, format=_format, colorize=True,
20 | )
21 |
22 | logger.add(
23 | f"log/mj-api.log",
24 | level=_lvl,
25 | format=_format,
26 | rotation="00:00",
27 | retention="3 days",
28 | backtrace=True,
29 | diagnose=True,
30 | enqueue=True
31 | )
32 |
--------------------------------------------------------------------------------
/app/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/yokonsan/midjourney-api/a5781c4e5123624e55623bb468f0f59046e643f0/app/__init__.py
--------------------------------------------------------------------------------
/app/handler.py:
--------------------------------------------------------------------------------
1 | import hashlib
2 | import time
3 | from functools import wraps
4 | from typing import Union
5 |
6 | from fastapi import status
7 | from fastapi.responses import JSONResponse
8 |
9 | from exceptions import BannedPromptError
10 | from lib.prompt import BANNED_PROMPT
11 |
12 | PROMPT_PREFIX = "<#"
13 | PROMPT_SUFFIX = "#>"
14 |
15 |
16 | def check_banned(prompt: str):
17 | words = set(w.lower() for w in prompt.split())
18 | if len(words & BANNED_PROMPT) != 0:
19 | raise BannedPromptError(f"banned prompt: {prompt}")
20 |
21 |
22 | def unique_id():
23 | """生成唯一的 10 位数字，作为任务 ID"""
24 | return int(hashlib.sha256(str(time.time()).encode("utf-8")).hexdigest(), 16) % 10**10
25 |
26 |
27 | def prompt_handler(prompt: str, picurl: Union[str, None] = None):
28 | """
29 | 拼接 Prompt 形如: <#1234567890#>a cute cat
30 | """
31 | check_banned(prompt)
32 |
33 | trigger_id = str(unique_id())
34 |
35 | if not picurl and prompt.startswith(("http://", "https://")):
36 | picurl, _, prompt = prompt.partition(" ")
37 |
38 | return trigger_id, f"{picurl+' ' if picurl else ''}{PROMPT_PREFIX}{trigger_id}{PROMPT_SUFFIX}{prompt}"
39 |
40 |
41 | def http_response(func):
42 | @wraps(func)
43 | async def router(*args, **kwargs):
44 | trigger_id, resp = await func(*args, **kwargs)
45 | if resp is not None:
46 | code, trigger_result = status.HTTP_200_OK, "success"
47 | else:
48 | code, trigger_result = status.HTTP_400_BAD_REQUEST, "fail"
49 |
50 | return JSONResponse(
51 | status_code=code,
52 | content={"trigger_id": trigger_id, "trigger_result": trigger_result}
53 | )
54 |
55 | return router
56 |
--------------------------------------------------------------------------------
/app/routers.py:
--------------------------------------------------------------------------------
1 | from fastapi import APIRouter, UploadFile
2 |
3 | from lib.api import discord
4 | from lib.api.discord import TriggerType
5 | from util._queue import taskqueue
6 | from .handler import prompt_handler, unique_id
7 | from .schema import (
8 | TriggerExpandIn,
9 | TriggerImagineIn,
10 | TriggerUVIn,
11 | TriggerResetIn,
12 | QueueReleaseIn,
13 | TriggerResponse,
14 | TriggerZoomOutIn,
15 | UploadResponse,
16 | TriggerDescribeIn,
17 | SendMessageResponse,
18 | SendMessageIn,
19 | )
20 |
21 | router = APIRouter()
22 |
23 |
24 | @router.post("/imagine", response_model=TriggerResponse)
25 | async def imagine(body: TriggerImagineIn):
26 | trigger_id, prompt = prompt_handler(body.prompt, body.picurl)
27 | trigger_type = TriggerType.generate.value
28 |
29 | taskqueue.put(trigger_id, discord.generate, prompt)
30 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
31 |
32 |
33 | @router.post("/upscale", response_model=TriggerResponse)
34 | async def upscale(body: TriggerUVIn):
35 | trigger_id = body.trigger_id
36 | trigger_type = TriggerType.upscale.value
37 |
38 | taskqueue.put(trigger_id, discord.upscale, **body.dict())
39 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
40 |
41 |
42 | @router.post("/variation", response_model=TriggerResponse)
43 | async def variation(body: TriggerUVIn):
44 | trigger_id = body.trigger_id
45 | trigger_type = TriggerType.variation.value
46 |
47 | taskqueue.put(trigger_id, discord.variation, **body.dict())
48 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
49 |
50 |
51 | @router.post("/reset", response_model=TriggerResponse)
52 | async def reset(body: TriggerResetIn):
53 | trigger_id = body.trigger_id
54 | trigger_type = TriggerType.reset.value
55 |
56 | taskqueue.put(trigger_id, discord.reset, **body.dict())
57 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
58 |
59 |
60 | @router.post("/describe", response_model=TriggerResponse)
61 | async def describe(body: TriggerDescribeIn):
62 | trigger_id = body.trigger_id
63 | trigger_type = TriggerType.describe.value
64 |
65 | taskqueue.put(trigger_id, discord.describe, **body.dict())
66 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
67 |
68 |
69 | @router.post("/upload", response_model=UploadResponse)
70 | async def upload_attachment(file: UploadFile):
71 | if not file.content_type.startswith("image/"):
72 | return {"message": "must image"}
73 |
74 | trigger_id = str(unique_id())
75 | filename = f"{trigger_id}.jpg"
76 | file_size = file.size
77 | attachment = await discord.upload_attachment(filename, file_size, await file.read())
78 | if not (attachment and attachment.get("upload_url")):
79 | return {"message": "Failed to upload image"}
80 |
81 | return {
82 | "upload_filename": attachment.get("upload_filename"),
83 | "upload_url": attachment.get("upload_url"),
84 | "trigger_id": trigger_id,
85 | }
86 |
87 |
88 | @router.post("/message", response_model=SendMessageResponse)
89 | async def send_message(body: SendMessageIn):
90 | picurl = await discord.send_attachment_message(body.upload_filename)
91 | if not picurl:
92 | return {"message": "Failed to send message"}
93 |
94 | return {"picurl": picurl}
95 |
96 |
97 | @router.post("/queue/release", response_model=TriggerResponse)
98 | async def queue_release(body: QueueReleaseIn):
99 | """bot 清除队列任务"""
100 | taskqueue.pop(body.trigger_id)
101 |
102 | return body
103 |
104 |
105 | @router.post("/solo_variation", response_model=TriggerResponse)
106 | async def solo_variation(body: TriggerUVIn):
107 | trigger_id = body.trigger_id
108 | trigger_type = TriggerType.solo_variation.value
109 | taskqueue.put(trigger_id, discord.solo_variation, **body.dict())
110 |
111 | # 返回结果
112 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
113 |
114 | @router.post("/solo_low_variation", response_model=TriggerResponse)
115 | async def solo_low_variation(body: TriggerUVIn):
116 | trigger_id = body.trigger_id
117 | trigger_type = TriggerType.solo_low_variation.value
118 | taskqueue.put(trigger_id, discord.solo_low_variation, **body.dict())
119 |
120 | # 返回结果
121 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
122 |
123 | @router.post("/solo_high_variation", response_model=TriggerResponse)
124 | async def solo_high_variation(body: TriggerUVIn):
125 | trigger_id = body.trigger_id
126 | trigger_type = TriggerType.solo_high_variation.value
127 | taskqueue.put(trigger_id, discord.solo_high_variation, **body.dict())
128 |
129 | # 返回结果
130 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
131 |
132 | @router.post("/expand", response_model=TriggerResponse)
133 | async def expand(body: TriggerExpandIn):
134 | trigger_id = body.trigger_id
135 | trigger_type = TriggerType.expand.value
136 | taskqueue.put(trigger_id, discord.expand, **body.dict())
137 |
138 | # 返回结果
139 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
140 |
141 |
142 | @router.post("/zoomout", response_model=TriggerResponse)
143 | async def zoomout(body: TriggerZoomOutIn):
144 | trigger_id = body.trigger_id
145 | trigger_type = TriggerType.zoomout.value
146 | taskqueue.put(trigger_id, discord.zoomout, **body.dict())
147 |
148 | # 返回结果
149 | return {"trigger_id": trigger_id, "trigger_type": trigger_type}
150 |
151 |
--------------------------------------------------------------------------------
/app/schema.py:
--------------------------------------------------------------------------------
1 | from typing import Optional
2 |
3 | from pydantic import BaseModel
4 |
5 |
6 | class TriggerImagineIn(BaseModel):
7 | prompt: str
8 | picurl: Optional[str]
9 |
10 |
11 | class TriggerUVIn(BaseModel):
12 | index: int
13 | msg_id: str
14 | msg_hash: str
15 |
16 | trigger_id: str # 供业务定位触发ID，/trigger/imagine 接口返回的 trigger_id
17 |
18 |
19 | class TriggerResetIn(BaseModel):
20 | msg_id: str
21 | msg_hash: str
22 |
23 | trigger_id: str # 供业务定位触发ID，/trigger/imagine 接口返回的 trigger_id
24 |
25 |
26 | class TriggerExpandIn(BaseModel):
27 | msg_id: str
28 | msg_hash: str
29 | direction: str # right/left/up/down
30 |
31 | trigger_id: str # 供业务定位触发ID，/trigger/imagine 接口返回的 trigger_id
32 |
33 | class TriggerZoomOutIn(BaseModel):
34 | msg_id: str
35 | msg_hash: str
36 | zoomout: int # 2x: 50; 1.5x: 75
37 |
38 | trigger_id: str # 供业务定位触发ID，/trigger/imagine 接口返回的 trigger_id
39 |
40 |
41 | class TriggerDescribeIn(BaseModel):
42 | upload_filename: str
43 | trigger_id: str
44 |
45 |
46 | class QueueReleaseIn(BaseModel):
47 | trigger_id: str
48 |
49 |
50 | class TriggerResponse(BaseModel):
51 | message: str = "success"
52 | trigger_id: str
53 | trigger_type: str = ""
54 |
55 |
56 | class UploadResponse(BaseModel):
57 | message: str = "success"
58 | upload_filename: str = ""
59 | upload_url: str = ""
60 | trigger_id: str
61 |
62 |
63 | class SendMessageIn(BaseModel):
64 | upload_filename: str
65 |
66 |
67 | class SendMessageResponse(BaseModel):
68 | message: str = "success"
69 | picurl: str
70 |
--------------------------------------------------------------------------------
/app/server.py:
--------------------------------------------------------------------------------
1 | import uvicorn
2 | from fastapi import FastAPI, status
3 | from fastapi.exceptions import RequestValidationError
4 | from fastapi.responses import JSONResponse
5 |
6 | from exceptions import APPBaseException, ErrorCode
7 |
8 |
9 | def init_app():
10 | _app = FastAPI(title="Midjourney API")
11 |
12 | register_blueprints(_app)
13 | exc_handler(_app)
14 |
15 | return _app
16 |
17 |
18 | def exc_handler(_app):
19 | @_app.exception_handler(RequestValidationError)
20 | def validation_exception_handler(_, exc: RequestValidationError):
21 | return JSONResponse(
22 | status_code=status.HTTP_400_BAD_REQUEST,
23 | content={
24 | "code": ErrorCode.REQUEST_PARAMS_ERROR.value,
25 | "message": f"request params error: {exc.body}"
26 | },
27 | )
28 |
29 | @_app.exception_handler(APPBaseException)
30 | def validation_exception_handler(_, exc: APPBaseException):
31 | return JSONResponse(
32 | status_code=status.HTTP_200_OK,
33 | content={
34 | "code": exc.code.value,
35 | "message": exc.message
36 | },
37 | )
38 |
39 |
40 | def register_blueprints(_app):
41 | from app import routers
42 | _app.include_router(routers.router, prefix="/v1/api/trigger")
43 |
44 |
45 | def run(host, port):
46 | _app = init_app()
47 | uvicorn.run(_app, port=port, host=host)
48 |
--------------------------------------------------------------------------------
/banned_words.txt:
--------------------------------------------------------------------------------
1 | bloodbath
2 | belly button
3 | ovaries
4 | -backed
5 | patriotic
6 | bruises
7 | thick
8 | blood
9 | sensual
10 | deepfake
11 | titty
12 | tryphophobia
13 | rule34
14 | sexi
15 | mammaries
16 | -edge
17 | jinping
18 | without clothes on
19 | booba
20 | president xi
21 | veiny
22 | orgy
23 | seducing
24 | playboy
25 | legs spread
26 | sexualiz
27 | pleasure
28 | mp5
29 | honkey
30 | crucified
31 | hentai
32 | stripped
33 | wound
34 | shag
35 | vivisection
36 | behind bars
37 | tight white
38 | bloodshot
39 | skimpy
40 | bloody
41 | nude
42 | mommy milker
43 | making love
44 | tied up
45 | brown pudding
46 | visceral
47 | bathing
48 | labia
49 | risque
50 | feces
51 | big ass
52 | ahegao
53 | fascist
54 | jerk off king at pic
55 | mccurry
56 | sperm
57 | sensored
58 | disturbing
59 | busty
60 | sexy
61 | reticulum
62 | rear end
63 | cocaine
64 | pleasures
65 | crack
66 | kill
67 | warts
68 | slaughter
69 | lolita
70 | guts
71 | honkers
72 | dick
73 | dominatrix
74 | sexual
75 | zero clothes
76 | boobs
77 | vagina
78 | boudoir
79 | booty
80 | seductive
81 | porn
82 | cannibal
83 | seduction
84 | loli
85 | crucifixion
86 | bimbo
87 | negligee
88 | hemoglobin
89 | shit
90 | bunghole
91 | bosom
92 | vein
93 | coon
94 | infested
95 | penis
96 | wincest
97 | jail
98 | poop
99 | farts
100 | cleavage
101 | smut
102 | ballgag
103 | cannibalism
104 | torture
105 | waifu
106 | surgery
107 | inappropriate
108 | naughty
109 | intimate
110 | sultry
111 | belle delphine
112 | khorne
113 | explicit
114 | fuck
115 | see through
116 | breasts
117 | anus
118 | horny
119 | cutting
120 | gory
121 | camisole
122 | revealing clothing
123 | heroin
124 | censored
125 | barely dressed
126 | hitler
127 | badonkers
128 | girth
129 | nipple
130 | scantily clad
131 | cronenberg
132 | hardcore
133 | meth
134 | unclothed
135 | pus
136 | wearing nothing
137 | succubus
138 | erotic
139 | massacre
140 | flesh
141 | pinup
142 | infected
143 | vomit
144 | minge
145 | sakimichan
146 | wang
147 | decapitate
148 | phallus
149 | fart
150 | dong
151 | voluptuous
152 | jav
153 | see-through
154 | clunge
155 | large bust
156 | bra
157 | arrest
158 | sensuality
159 | shirtless
160 | no clothes
161 | knob
162 | xxx
163 | prophet mohammed
164 | suicide
165 | silenced
166 | crotch
167 | bodily fluids
168 | seductively
169 | incest
170 | twerk
171 | sexy female
172 | reproduce
173 | killing
174 | car crash
175 | kinbaku
176 | invisible clothes
177 | hooters
178 | sadist
179 | erected
180 | slave
181 | provocative
182 | teratoma
183 | naked
184 | full frontal
185 | lingerie
186 | oppai
187 | petite
188 | bleed
189 | no shirt
190 | brothel
191 | au naturel
192 | organs
193 | bare
194 | nazi
195 | with no shirt
196 | excrement
197 | thot
198 | zedong
199 | corpse
200 | shibari
201 | ass
202 | gruesome
203 | transparent
204 | arse
205 |
--------------------------------------------------------------------------------
/build.sh:
--------------------------------------------------------------------------------
1 | git pull
2 | docker build -t kunyu/midjourney-api:1.0 .
3 |
--------------------------------------------------------------------------------
/entrypoint.sh:
--------------------------------------------------------------------------------
1 | #!/bin/bash
2 | set -e
3 |
4 | if [ "$1" = 'http' ]; then
5 | set -- python server.py
6 | fi
7 |
8 | if [ "$1" = 'bot' ]; then
9 | set -- python task_bot.py "$@"
10 | fi
11 |
12 | exec "$@"
13 |
--------------------------------------------------------------------------------
/exceptions.py:
--------------------------------------------------------------------------------
1 | from enum import Enum
2 | from typing import Union
3 |
4 |
5 | class ErrorCode(Enum):
6 | MISS_REQUIRED_VARIABLE_ERROR = 11
7 | MAX_RETRY_ERROR = 12
8 | REQUEST_PARAMS_ERROR = 13
9 | BANNED_PROMPT_ERROR = 14
10 | QUEUE_FULL_ERROR = 15
11 |
12 |
13 | class SuccessCode(Enum):
14 | SUCCESS = 0
15 |
16 |
17 | ReturnCode = Union[ErrorCode, SuccessCode]
18 |
19 |
20 | class APPBaseException(Exception):
21 | code: ErrorCode
22 | message: str
23 |
24 | def __init__(self, message: str):
25 | super().__init__(message)
26 | self.message = message
27 |
28 |
29 | class MissRequiredVariableError(APPBaseException):
30 | """缺少必需变量"""
31 | code = ErrorCode.MISS_REQUIRED_VARIABLE_ERROR
32 |
33 |
34 | class MaxRetryError(APPBaseException):
35 | """请求最大重试错误"""
36 | code = ErrorCode.MAX_RETRY_ERROR
37 |
38 |
39 | class RequestParamsError(APPBaseException):
40 | """请求参数异常"""
41 | code = ErrorCode.REQUEST_PARAMS_ERROR
42 |
43 |
44 | class BannedPromptError(APPBaseException):
45 | """提示词被禁用"""
46 | code = ErrorCode.BANNED_PROMPT_ERROR
47 |
48 |
49 | class QueueFullError(APPBaseException):
50 | """队列已满"""
51 | code = ErrorCode.QUEUE_FULL_ERROR
52 |
--------------------------------------------------------------------------------
/lib/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/yokonsan/midjourney-api/a5781c4e5123624e55623bb468f0f59046e643f0/lib/__init__.py
--------------------------------------------------------------------------------
/lib/api/__init__.py:
--------------------------------------------------------------------------------
1 | from os import getenv
2 |
3 | from exceptions import MissRequiredVariableError
4 |
5 | GUILD_ID = getenv("GUILD_ID")
6 | CHANNEL_ID = getenv("CHANNEL_ID")
7 | USER_TOKEN = getenv("USER_TOKEN")
8 | CALLBACK_URL = getenv("CALLBACK_URL")
9 | PROXY_URL = getenv("PROXY_URL")
10 |
11 | DRAW_VERSION = getenv("DRAW_VERSION")
12 |
13 | if not all([GUILD_ID, CHANNEL_ID, USER_TOKEN, DRAW_VERSION]):
14 | raise MissRequiredVariableError(
15 | "Missing required environment variable: [GUILD_ID, CHANNEL_ID, USER_TOKEN, DRAW_VERSION]")
16 |
--------------------------------------------------------------------------------
/lib/api/callback.py:
--------------------------------------------------------------------------------
1 | from os import getenv
2 |
3 | import aiohttp
4 | from loguru import logger
5 |
6 | from lib.api import CALLBACK_URL
7 | from util.fetch import fetch
8 |
9 |
10 | async def callback(data):
11 | logger.debug(f"callback data: {data}")
12 | if not CALLBACK_URL:
13 | return
14 |
15 | headers = {"Content-Type": "application/json"}
16 | async with aiohttp.ClientSession(
17 | timeout=aiohttp.ClientTimeout(total=30),
18 | headers=headers
19 | ) as session:
20 | await fetch(session, CALLBACK_URL, json=data)
21 |
22 |
23 | QUEUE_RELEASE_API = getenv("QUEUE_RELEASE_API") \
24 | or "http://127.0.0.1:8062/v1/api/trigger/queue/release"
25 |
26 |
27 | async def queue_release(trigger_id: str):
28 | logger.debug(f"queue_release: {trigger_id}")
29 |
30 | headers = {"Content-Type": "application/json"}
31 | data = {"trigger_id": trigger_id}
32 | async with aiohttp.ClientSession(
33 | timeout=aiohttp.ClientTimeout(total=30),
34 | headers=headers
35 | ) as session:
36 | await fetch(session, QUEUE_RELEASE_API, json=data)
37 |
--------------------------------------------------------------------------------
/lib/api/discord.py:
--------------------------------------------------------------------------------
1 | import json
2 | from enum import Enum
3 | from typing import Dict, Any, Union
4 |
5 | import aiohttp
6 |
7 | from lib.api import CHANNEL_ID, USER_TOKEN, GUILD_ID, DRAW_VERSION, PROXY_URL
8 | from util.fetch import fetch, fetch_json, FetchMethod
9 |
10 | TRIGGER_URL = "https://discord.com/api/v9/interactions"
11 | UPLOAD_ATTACHMENT_URL = f"https://discord.com/api/v9/channels/{CHANNEL_ID}/attachments"
12 | SEND_MESSAGE_URL = f"https://discord.com/api/v9/channels/{CHANNEL_ID}/messages"
13 | HEADERS = {
14 | "Content-Type": "application/json",
15 | "Authorization": USER_TOKEN
16 | }
17 |
18 |
19 | class TriggerType(str, Enum):
20 | generate = "generate"
21 | upscale = "upscale"
22 | variation = "variation"
23 | solo_variation = "solo_variation"
24 | solo_low_variation = "solo_low_variation"
25 | solo_high_variation = "solo_high_variation"
26 | max_upscale = "max_upscale"
27 | reset = "reset"
28 | describe = "describe"
29 | expand = "expand"
30 | zoomout = "zoomout"
31 |
32 |
33 | async def trigger(payload: Dict[str, Any]):
34 | async with aiohttp.ClientSession(
35 | timeout=aiohttp.ClientTimeout(total=30),
36 | headers=HEADERS
37 | ) as session:
38 | return await fetch(session, TRIGGER_URL, data=json.dumps(payload), proxy=PROXY_URL)
39 |
40 |
41 | async def upload_attachment(
42 | filename: str, file_size: int, image: bytes
43 | ) -> Union[Dict[str, Union[str, int]], None]:
44 | payload = {
45 | "files": [{
46 | "filename": filename,
47 | "file_size": file_size,
48 | "id": "0"
49 | }]
50 | }
51 | async with aiohttp.ClientSession(
52 | timeout=aiohttp.ClientTimeout(total=30),
53 | headers=HEADERS
54 | ) as session:
55 | response = await fetch_json(session, UPLOAD_ATTACHMENT_URL, data=json.dumps(payload))
56 | if not response or not response.get("attachments"):
57 | return None
58 |
59 | attachment = response["attachments"][0]
60 |
61 | response = await put_attachment(attachment.get("upload_url"), image)
62 | return attachment if response is not None else None
63 |
64 |
65 | async def put_attachment(url: str, image: bytes):
66 | headers = {"Content-Type": "image/png"}
67 | async with aiohttp.ClientSession(
68 | timeout=aiohttp.ClientTimeout(total=30),
69 | headers=headers
70 | ) as session:
71 | return await fetch(session, url, data=image, method=FetchMethod.put)
72 |
73 |
74 | async def send_attachment_message(upload_filename: str) -> Union[str, None]:
75 | payload = {
76 | "content": "",
77 | "nonce": "",
78 | "channel_id": "1105829904790065223",
79 | "type": 0,
80 | "sticker_ids": [],
81 | "attachments": [{
82 | "id": "0",
83 | "filename": upload_filename.split("/")[-1],
84 | "uploaded_filename": upload_filename
85 | }]
86 | }
87 | async with aiohttp.ClientSession(
88 | timeout=aiohttp.ClientTimeout(total=30),
89 | headers=HEADERS
90 | ) as session:
91 | response = await fetch_json(session, SEND_MESSAGE_URL, data=json.dumps(payload))
92 | if not response or not response.get("attachments"):
93 | return None
94 |
95 | attachment = response["attachments"][0]
96 | return attachment.get("url")
97 |
98 |
99 | def _trigger_payload(type_: int, data: Dict[str, Any], **kwargs) -> Dict[str, Any]:
100 | payload = {
101 | "type": type_,
102 | "application_id": "936929561302675456",
103 | "guild_id": GUILD_ID,
104 | "channel_id": CHANNEL_ID,
105 | "session_id": "cb06f61453064c0983f2adae2a88c223",
106 | "data": data
107 | }
108 | payload.update(kwargs)
109 | return payload
110 |
111 |
112 | async def generate(prompt: str, **kwargs):
113 | payload = _trigger_payload(2, {
114 | "version": DRAW_VERSION,
115 | "id": "938956540159881230",
116 | "name": "imagine",
117 | "type": 1,
118 | "options": [{
119 | "type": 3,
120 | "name": "prompt",
121 | "value": prompt
122 | }],
123 | "attachments": []
124 | })
125 | return await trigger(payload)
126 |
127 |
128 | async def upscale(index: int, msg_id: str, msg_hash: str, **kwargs):
129 | kwargs = {
130 | "message_flags": 0,
131 | "message_id": msg_id,
132 | }
133 | payload = _trigger_payload(3, {
134 | "component_type": 2,
135 | "custom_id": f"MJ::JOB::upsample::{index}::{msg_hash}"
136 | }, **kwargs)
137 | return await trigger(payload)
138 |
139 |
140 | async def variation(index: int, msg_id: str, msg_hash: str, **kwargs):
141 | kwargs = {
142 | "message_flags": 0,
143 | "message_id": msg_id,
144 | }
145 | payload = _trigger_payload(3, {
146 | "component_type": 2,
147 | "custom_id": f"MJ::JOB::variation::{index}::{msg_hash}"
148 | }, **kwargs)
149 | return await trigger(payload)
150 |
151 |
152 | async def solo_variation(msg_id: str, msg_hash: str, **kwargs):
153 | kwargs = {
154 | "message_flags": 0,
155 | "message_id": msg_id,
156 | }
157 | payload = _trigger_payload(3, {
158 | "component_type": 2,
159 | "custom_id": f"MJ::JOB::variation::1::{msg_hash}::SOLO"
160 | }, **kwargs)
161 | return await trigger(payload)
162 |
163 | async def solo_low_variation(msg_id: str, msg_hash: str, **kwargs):
164 | kwargs = {
165 | "message_flags": 0,
166 | "message_id": msg_id,
167 | }
168 | payload = _trigger_payload(3, {
169 | "component_type": 2,
170 | "custom_id": f"MJ::JOB::low_variation::1::{msg_hash}::SOLO"
171 | }, **kwargs)
172 | return await trigger(payload)
173 |
174 | async def solo_high_variation(msg_id: str, msg_hash: str, **kwargs):
175 | kwargs = {
176 | "message_flags": 0,
177 | "message_id": msg_id,
178 | }
179 | payload = _trigger_payload(3, {
180 | "component_type": 2,
181 | "custom_id": f"MJ::JOB::high_variation::1::{msg_hash}::SOLO"
182 | }, **kwargs)
183 | return await trigger(payload)
184 |
185 |
186 | async def expand(msg_id: str, msg_hash: str, direction: str, **kwargs):
187 | kwargs = {
188 | "message_flags": 0,
189 | "message_id": msg_id,
190 | }
191 | payload = _trigger_payload(3, {
192 | "component_type": 2,
193 | "custom_id": f"MJ::JOB::pan_{direction}::1::{msg_hash}::SOLO"
194 | }, **kwargs)
195 | return await trigger(payload)
196 |
197 |
198 | async def zoomout(msg_id: str, msg_hash: str, zoomout: int, **kwargs):
199 | kwargs = {
200 | "message_flags": 0,
201 | "message_id": msg_id,
202 | }
203 | payload = _trigger_payload(3, {
204 | "component_type": 2,
205 | "custom_id": f"MJ::Outpaint::{zoomout}::1::{msg_hash}::SOLO"
206 | }, **kwargs)
207 | return await trigger(payload)
208 |
209 |
210 | async def max_upscale(msg_id: str, msg_hash: str, **kwargs):
211 | kwargs = {
212 | "message_flags": 0,
213 | "message_id": msg_id,
214 | }
215 | payload = _trigger_payload(3, {
216 | "component_type": 2,
217 | "custom_id": f"MJ::JOB::upsample_max::1::{msg_hash}::SOLO"
218 | }, **kwargs)
219 | return await trigger(payload)
220 |
221 |
222 | async def reset(msg_id: str, msg_hash: str, **kwargs):
223 | kwargs = {
224 | "message_flags": 0,
225 | "message_id": msg_id,
226 | }
227 | payload = _trigger_payload(3, {
228 | "component_type": 2,
229 | "custom_id": f"MJ::JOB::reroll::0::{msg_hash}::SOLO"
230 | }, **kwargs)
231 | return await trigger(payload)
232 |
233 |
234 | async def describe(upload_filename: str, **kwargs):
235 | payload = _trigger_payload(2, {
236 | "version": DRAW_VERSION,
237 | "id": "1092492867185950852",
238 | "name": "describe",
239 | "type": 1,
240 | "options": [
241 | {
242 | "type": 11,
243 | "name": "image",
244 | "value": 0
245 | }
246 | ],
247 | "attachments": [{
248 | "id": "0",
249 | "filename": upload_filename.split("/")[-1],
250 | "uploaded_filename": upload_filename,
251 | }]
252 | })
253 | return await trigger(payload)
254 |
--------------------------------------------------------------------------------
/lib/prompt.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 |
3 | import aiofiles
4 |
5 |
6 | async def loads_banned_words():
7 | prompt = set()
8 | filename = "banned_words.txt"
9 | async with aiofiles.open(filename, "r") as r:
10 | for line in await r.readlines():
11 | prompt.add(line.strip())
12 |
13 | return prompt
14 |
15 | BANNED_PROMPT = asyncio.run(loads_banned_words())
16 |
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | discord.py==2.2.3
2 | fastapi==0.95
3 | python-dotenv==1.0.0
4 | uvicorn==0.22.0
5 | pydantic~=1.10.7
6 | aiohttp~=3.8.4
7 | loguru==0.7.0
8 | aiofiles==23.1.0
9 | python-multipart==0.0.6
--------------------------------------------------------------------------------
/server.py:
--------------------------------------------------------------------------------
1 | import __init__ # noqa
2 | from app import server
3 |
4 | api_app = server.init_app()
5 |
6 |
7 | if __name__ == '__main__':
8 | server.run("0.0.0.0", 8062)
9 |
--------------------------------------------------------------------------------
/start.sh:
--------------------------------------------------------------------------------
1 | docker rm -f mj-server mj-bot
2 | #docker rmi kunyu/midjourney-api:1.0
3 |
4 | docker network create mjapi
5 | docker run -d --net mjapi --name mj-server -p 8062:8062 \
6 | -e TZ=Asia/Shanghai \
7 | -e LOG_LEVEL=DEBUG \
8 | -e USER_TOKEN="" \
9 | -e GUILD_ID="" \
10 | -e CHANNEL_ID="" \
11 | -e CONCUR_SIZE=3 \
12 | -e WAIT_SIZE=10 \
13 | kunyu/midjourney-api:1.0
14 |
15 | docker run -d --net mjapi --name mj-bot \
16 | -e TZ=Asia/Shanghai \
17 | -e LOG_LEVEL=DEBUG \
18 | -e USER_TOKEN="" \
19 | -e BOT_TOKEN="" \
20 | -e GUILD_ID="" \
21 | -e CHANNEL_ID="" \
22 | -e CALLBACK_URL="" \
23 | -e QUEUE_RELEASE_API="http://mj-server:8062/v1/api/trigger/queue/release" \
24 | kunyu/midjourney-api:1.0 bot
25 |
--------------------------------------------------------------------------------
/task/__init__.py:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/yokonsan/midjourney-api/a5781c4e5123624e55623bb468f0f59046e643f0/task/__init__.py
--------------------------------------------------------------------------------
/task/bot/__init__.py:
--------------------------------------------------------------------------------
1 | from enum import Enum
2 |
3 |
4 | class TriggerStatus(Enum):
5 | start = "start" # 首次触发 MessageType.chat_input_command
6 | generating = "generating" # 生成中
7 | end = "end" # 生成结束 MessageType.default
8 | error = "error" # 生成错误
9 | banned = "banned" # 提示词被禁
10 |
11 | text = "text" # 文本内容：describe
12 |
13 | verify = "verify" # 需人工验证
14 |
--------------------------------------------------------------------------------
/task/bot/_typing.py:
--------------------------------------------------------------------------------
1 | from typing import TypedDict, List, Union
2 |
3 |
4 | class Attachment(TypedDict):
5 | id: int
6 | url: str
7 | proxy_url: str
8 | filename: str
9 | content_type: str
10 | width: int
11 | height: int
12 | size: int
13 | ephemeral: bool
14 |
15 |
16 | class EmbedsImage(TypedDict):
17 | url: str
18 | proxy_url: str
19 |
20 |
21 | class Embed(TypedDict):
22 | type: str
23 | description: str
24 | image: EmbedsImage
25 |
26 |
27 | class CallbackData(TypedDict):
28 | type: str
29 | id: int
30 | content: str
31 | attachments: List[Attachment]
32 | embeds: List[Embed]
33 |
34 | trigger_id: str
35 |
--------------------------------------------------------------------------------
/task/bot/handler.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 | import re
3 | from typing import Dict, Union, Any
4 |
5 | from discord import Message
6 |
7 | from app.handler import PROMPT_PREFIX, PROMPT_SUFFIX
8 | from lib.api.callback import queue_release, callback
9 | from task.bot._typing import CallbackData, Attachment, Embed
10 |
11 | TRIGGER_ID_PATTERN = f"{PROMPT_PREFIX}(\w+?){PROMPT_SUFFIX}" # 消息 ID 正则
12 |
13 | TEMP_MAP: Dict[str, bool] = {} # 临时存储消息流转信息
14 |
15 |
16 | def get_temp(trigger_id: str):
17 | return TEMP_MAP.get(trigger_id)
18 |
19 |
20 | def set_temp(trigger_id: str):
21 | TEMP_MAP[trigger_id] = True
22 |
23 |
24 | def pop_temp(trigger_id: str):
25 | asyncio.get_event_loop().create_task(queue_release(trigger_id))
26 | try:
27 | TEMP_MAP.pop(trigger_id)
28 | except KeyError:
29 | pass
30 |
31 |
32 | def match_trigger_id(content: str) -> Union[str, None]:
33 | match = re.findall(TRIGGER_ID_PATTERN, content)
34 | return match[0] if match else None
35 |
36 |
37 | async def callback_trigger(trigger_id: str, trigger_status: str, message: Message):
38 | await callback(CallbackData(
39 | type=trigger_status,
40 | id=message.id,
41 | content=message.content,
42 | attachments=[
43 | Attachment(**attachment.to_dict())
44 | for attachment in message.attachments
45 | ],
46 | embeds=[],
47 | trigger_id=trigger_id,
48 | ))
49 |
50 |
51 | async def callback_describe(trigger_status: str, message: Message, embed: Dict[str, Any]):
52 | url = embed.get("image", {}).get("url")
53 | trigger_id = url.split("/")[-1].split(".")[0]
54 |
55 | await callback(CallbackData(
56 | type=trigger_status,
57 | id=message.id,
58 | content=message.content,
59 | attachments=[],
60 | embeds=[
61 | Embed(**embed)
62 | ],
63 | trigger_id=trigger_id,
64 | ))
65 | return trigger_id
66 |
--------------------------------------------------------------------------------
/task/bot/listener.py:
--------------------------------------------------------------------------------
1 | from discord import Intents, Message
2 | from discord.ext import commands
3 | from loguru import logger
4 | from lib.api import PROXY_URL
5 | from task.bot import TriggerStatus
6 | from task.bot.handler import (
7 | match_trigger_id,
8 | set_temp,
9 | pop_temp,
10 | get_temp,
11 | callback_trigger,
12 | callback_describe
13 | )
14 |
15 | intents = Intents.default()
16 | intents.message_content = True
17 | bot = commands.Bot(command_prefix="", intents=intents, proxy=PROXY_URL)
18 |
19 |
20 | @bot.event
21 | async def on_ready():
22 | logger.success(f"Logged in as {bot.user} (ID: {bot.user.id})")
23 |
24 |
25 | @bot.event
26 | async def on_message(message: Message):
27 | if message.author.id != 936929561302675456:
28 | return
29 |
30 | logger.debug(f"on_message: {message.content}")
31 | logger.debug(f"on_message embeds: {message.embeds[0].to_dict() if message.embeds else message.embeds}")
32 | content = message.content
33 | trigger_id = match_trigger_id(content)
34 | if not trigger_id:
35 | return
36 |
37 | if content.find("Waiting to start") != -1:
38 | trigger_status = TriggerStatus.start.value
39 | set_temp(trigger_id)
40 | elif content.find("(Stopped)") != -1:
41 | trigger_status = TriggerStatus.error.value
42 | pop_temp(trigger_id)
43 | else:
44 | trigger_status = TriggerStatus.end.value
45 | pop_temp(trigger_id)
46 |
47 | await callback_trigger(trigger_id, trigger_status, message)
48 |
49 |
50 | @bot.event
51 | async def on_message_edit(_: Message, after: Message):
52 | if after.author.id != 936929561302675456:
53 | return
54 |
55 | logger.debug(f"on_message_edit: {after.content}")
56 | if after.embeds:
57 | embed = after.embeds[0]
58 | if not (embed.image.width and embed.image.height):
59 | return
60 |
61 | embed = embed.to_dict()
62 | logger.debug(f"on_message_edit embeds: {embed}")
63 | trigger_status = TriggerStatus.text.value
64 | trigger_id = await callback_describe(trigger_status, after, embed)
65 | pop_temp(trigger_id)
66 | return
67 |
68 | trigger_id = match_trigger_id(after.content)
69 | if not trigger_id:
70 | return
71 |
72 | if after.webhook_id != "":
73 | await callback_trigger(trigger_id, TriggerStatus.generating.value, after)
74 |
75 |
76 | @bot.event
77 | async def on_message_delete(message: Message):
78 | if message.author.id != 936929561302675456:
79 | return
80 |
81 | trigger_id = match_trigger_id(message.content)
82 | if not trigger_id:
83 | return
84 |
85 | if get_temp(trigger_id) is None:
86 | return
87 |
88 | logger.debug(f"on_message_delete: {message.content}")
89 | logger.warning(f"sensitive content: {message.content}")
90 | trigger_status = TriggerStatus.banned.value
91 | pop_temp(trigger_id)
92 | await callback_trigger(trigger_id, trigger_status, message)
93 |
--------------------------------------------------------------------------------
/task_bot.py:
--------------------------------------------------------------------------------
1 | from os import getenv
2 |
3 | import __init__ # noqa
4 | from exceptions import MissRequiredVariableError
5 | from task.bot.listener import bot
6 |
7 | BOT_TOKEN = getenv("BOT_TOKEN")
8 | if not BOT_TOKEN:
9 | raise MissRequiredVariableError("Missing required environment variable: [BOT_TOKEN]")
10 |
11 |
12 | if __name__ == '__main__':
13 | bot.run(BOT_TOKEN)
14 |
--------------------------------------------------------------------------------
/util/_queue.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 | from collections import deque
3 | from os import getenv
4 | from typing import ParamSpec, Callable, Any, Dict, List, Deque
5 |
6 | from loguru import logger
7 |
8 | from exceptions import QueueFullError
9 |
10 | P = ParamSpec("P")
11 |
12 |
13 | class Task:
14 | def __init__(
15 | self, func: Callable[P, Any], *args: P.args, **kwargs: P.kwargs
16 | ) -> None:
17 | self.func = func
18 | self.args = args
19 | self.kwargs = kwargs
20 |
21 | async def __call__(self) -> None:
22 | await self.func(*self.args, **self.kwargs)
23 |
24 | def __repr__(self) -> str:
25 | return f"{self.func.__name__}({self.args}, {self.kwargs})"
26 |
27 |
28 | class TaskQueue:
29 | def __init__(self, concur_size: int, wait_size: int) -> None:
30 | self._concur_size = concur_size
31 | self._wait_size = wait_size
32 | self._wait_queue: Deque[Dict[str, Task]] = deque()
33 | self._concur_queue: List[str] = []
34 |
35 | def put(
36 | self,
37 | _trigger_id: str,
38 | func: Callable[P, Any],
39 | *args: P.args,
40 | **kwargs: P.kwargs
41 | ) -> None:
42 | if len(self._wait_queue) >= self._wait_size:
43 | raise QueueFullError(f"Task queue is full: {self._wait_size}")
44 |
45 | self._wait_queue.append({
46 | _trigger_id: Task(func, *args, **kwargs)
47 | })
48 | while self._wait_queue and len(self._concur_queue) < self._concur_size:
49 | self._exec()
50 |
51 | def pop(self, _trigger_id: str) -> None:
52 | self._concur_queue.remove(_trigger_id)
53 | if self._wait_queue:
54 | self._exec()
55 |
56 | def _exec(self):
57 | key, task = self._wait_queue.popleft().popitem()
58 | self._concur_queue.append(key)
59 |
60 | logger.debug(f"Task[{key}] start execution: {task}")
61 | loop = asyncio.get_running_loop()
62 | tsk = loop.create_task(task())
63 | # tsk.add_done_callback(
64 | # lambda t: print(t.result())
65 | # ) # todo
66 |
67 | def concur_size(self):
68 | return self._concur_size
69 |
70 | def wait_size(self):
71 | return self._wait_size
72 |
73 | def clear_wait(self):
74 | self._wait_queue.clear()
75 |
76 | def clear_concur(self):
77 | self._concur_queue.clear()
78 |
79 |
80 | taskqueue = TaskQueue(
81 | int(getenv("CONCUR_SIZE") or 9999),
82 | int(getenv("WAIT_SIZE") or 9999),
83 | )
84 |
--------------------------------------------------------------------------------
/util/fetch.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 | from typing import Callable, Coroutine, Any, TypeVar, Union, Dict
3 |
4 | from aiohttp import ClientError, ClientSession, hdrs
5 | from loguru import logger
6 |
7 | from exceptions import MaxRetryError
8 |
9 | T = TypeVar("T")
10 |
11 |
12 | class MaxRetry:
13 | """重试装饰器"""
14 | def __init__(self, max_retry: int = 0):
15 | self.max_retry = max_retry
16 |
17 | def __call__(self, connect_once: Callable[..., Coroutine[Any, Any, T]]) -> Callable[..., Coroutine[Any, Any, T]]:
18 | async def connect_n_times(*args: Any, **kwargs: Any) -> T:
19 | retry = self.max_retry + 1
20 | while retry:
21 | try:
22 | return await connect_once(*args, **kwargs)
23 | except ClientError as e:
24 | await asyncio.sleep(1)
25 | logger.warning(f"请求失败（{e.__class__.__name__}），正在重试，剩余 {retry - 1} 次")
26 | except asyncio.TimeoutError:
27 | logger.warning(f"请求超时，正在重试，剩余 {retry - 1} 次")
28 | finally:
29 | retry -= 1
30 | raise MaxRetryError("超出最大重试次数")
31 |
32 | return connect_n_times
33 |
34 |
35 | class FetchMethod:
36 | get = hdrs.METH_GET
37 | post = hdrs.METH_POST
38 | put = hdrs.METH_PUT
39 |
40 |
41 | @MaxRetry(2)
42 | async def fetch(
43 | session: ClientSession,
44 | url: str,
45 | method: str = FetchMethod.post, **kwargs
46 | ) -> Union[bool, None]:
47 | logger.debug(f"Fetch: {url}, {kwargs}")
48 | async with session.request(method, url, **kwargs) as resp:
49 | if not resp.ok:
50 | return None
51 | return True
52 |
53 |
54 | @MaxRetry(2)
55 | async def fetch_json(
56 | session: ClientSession,
57 | url: str,
58 | method: str = FetchMethod.post, **kwargs
59 | ) -> Union[Dict, None]:
60 | logger.debug(f"Fetch text: {url}")
61 | async with session.request(method, url, **kwargs) as resp:
62 | if not resp.ok:
63 | return None
64 | return await resp.json()
65 |
--------------------------------------------------------------------------------