From 47efb50694044cd51f35226cbc776b860de19af8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A4=D0=B8=D0=BB=D0=B0=D1=82=D0=BE=D0=B2=20=D0=90=D0=BD?= =?UTF-8?q?=D1=82=D0=BE=D0=BD=20=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D0=BC=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B8=D1=87?= Date: Fri, 28 Jul 2023 12:51:08 +0300 Subject: [PATCH 1/6] Add docker-compose.yaml --- docker-compose.yaml | 31 +++++++++++++++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 docker-compose.yaml diff --git a/docker-compose.yaml b/docker-compose.yaml new file mode 100644 index 0000000..bb39e9e --- /dev/null +++ b/docker-compose.yaml @@ -0,0 +1,31 @@ +networks: + mj-net: + name: mj-net + +services: + mj-server: + image: kunyu/midjourney-api:1.0 + ports: + - "8062:8062" + environment: + - LOG_LEVEL=${LOG_LEVEL} + - USER_TOKEN=${USER_TOKEN} + - GUILD_ID=${GUILD_ID} + - CHANNEL_ID=${CHANNEL_ID} + - CONCUR_SIZE=${CONCUR_SIZE} + - WAIT_SIZE=${WAIT_SIZE} + command: [http] + restart: unless-stopped + + mj-bot: + image: kunyu/midjourney-api:1.0 + environment: + - LOG_LEVEL=${LOG_LEVEL} + - USER_TOKEN=${USER_TOKEN} + - GUILD_ID=${GUILD_ID} + - CHANNEL_ID=${CHANNEL_ID} + - BOT_TOKEN=${BOT_TOKEN} + - CALLBACK_URL=${CALLBACK_URL} + - QUEUE_RELEASE_API=${QUEUE_RELEASE_API} + command: [bot] + restart: unless-stopped \ No newline at end of file From 47ea68404ed9c9f4ba8d144bf1828f81be9048f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A4=D0=B8=D0=BB=D0=B0=D1=82=D0=BE=D0=B2=20=D0=90=D0=BD?= =?UTF-8?q?=D1=82=D0=BE=D0=BD=20=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D0=BC=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B8=D1=87?= Date: Sun, 30 Jul 2023 14:52:48 +0300 Subject: [PATCH 2/6] Imagine requet also returns queue data --- app/routers.py | 8 +++++++- app/schema.py | 2 ++ util/_queue.py | 6 ++++++ 3 files changed, 15 insertions(+), 1 deletion(-) diff --git a/app/routers.py b/app/routers.py index e92ad30..0001448 100644 --- a/app/routers.py +++ b/app/routers.py @@ -25,7 +25,13 @@ async def imagine(body: TriggerImagineIn): trigger_type = TriggerType.generate.value taskqueue.put(trigger_id, discord.generate, prompt) - return {"trigger_id": trigger_id, "trigger_type": trigger_type} + + queue_wait_size = taskqueue.wait_queue_size() + queue_concur_size = taskqueue.concur_queue_size() + + return {"trigger_id": trigger_id, "trigger_type": trigger_type, + "wait_size": queue_wait_size, + "concur_size": queue_concur_size} @router.post("/upscale", response_model=TriggerResponse) diff --git a/app/schema.py b/app/schema.py index 2d2b9f3..a11c8c4 100644 --- a/app/schema.py +++ b/app/schema.py @@ -36,6 +36,8 @@ class TriggerResponse(BaseModel): message: str = "success" trigger_id: str trigger_type: str = "" + wait_size: int + concur_size: int class UploadResponse(BaseModel): diff --git a/util/_queue.py b/util/_queue.py index 1ed0639..e0d9951 100644 --- a/util/_queue.py +++ b/util/_queue.py @@ -76,6 +76,12 @@ def clear_wait(self): def clear_concur(self): self._concur_queue.clear() + def wait_queue_size(self): + return len(self._wait_queue) + + def concur_queue_size(self): + return len(self._concur_queue) + taskqueue = TaskQueue( int(getenv("CONCUR_SIZE") or 9999), From e5030c620c9dec519bbacbcce7cc650b420106ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A4=D0=B8=D0=BB=D0=B0=D1=82=D0=BE=D0=B2=20=D0=90=D0=BD?= =?UTF-8?q?=D1=82=D0=BE=D0=BD=20=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D0=BC=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B8=D1=87?= Date: Sun, 30 Jul 2023 15:22:09 +0300 Subject: [PATCH 3/6] New Dockerfile (image size 153 mb) --- Dockerfile | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/Dockerfile b/Dockerfile index 934b1fb..ffedb22 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,13 +1,19 @@ -FROM python:3.10.6 -LABEL creator="yokon" email="944682328@qq.com" +FROM python:3.10 AS builder -WORKDIR /app +COPY requirements.txt . +RUN pip install --user -r requirements.txt +FROM python:3.10-slim + +WORKDIR /code + +COPY --from=builder /root/.local /root/.local COPY . . -RUN pip install --upgrade pip \ - && pip install -i https://pypi.douban.com/simple/ -r requirements.txt \ - && chmod +x entrypoint.sh + +ENV PATH=/root/.local:$PATH + +RUN chmod +x entrypoint.sh ENTRYPOINT ["bash", "entrypoint.sh"] EXPOSE 8062 -CMD ["http"] \ No newline at end of file +CMD ["http"] From 0f357ed411b00e16670748e12ba7934fa73ad4b3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A4=D0=B8=D0=BB=D0=B0=D1=82=D0=BE=D0=B2=20=D0=90=D0=BD?= =?UTF-8?q?=D1=82=D0=BE=D0=BD=20=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D0=BC=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B8=D1=87?= Date: Sun, 30 Jul 2023 17:48:02 +0300 Subject: [PATCH 4/6] Fix docker base image --- Dockerfile | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Dockerfile b/Dockerfile index ffedb22..a6154af 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,9 +1,9 @@ -FROM python:3.10 AS builder +FROM python:3.10.6 AS builder COPY requirements.txt . RUN pip install --user -r requirements.txt -FROM python:3.10-slim +FROM python:3.10.6-slim WORKDIR /code From 4c7fc027946acd2bbb797d53f9438f8c6c130ac4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A4=D0=B8=D0=BB=D0=B0=D1=82=D0=BE=D0=B2=20=D0=90=D0=BD?= =?UTF-8?q?=D1=82=D0=BE=D0=BD=20=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D0=BC=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B8=D1=87?= Date: Sun, 30 Jul 2023 18:15:25 +0300 Subject: [PATCH 5/6] All responses return queue size --- app/routers.py | 39 ++++++++++++++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/app/routers.py b/app/routers.py index 0001448..04411aa 100644 --- a/app/routers.py +++ b/app/routers.py @@ -40,7 +40,13 @@ async def upscale(body: TriggerUVIn): trigger_type = TriggerType.upscale.value taskqueue.put(trigger_id, discord.upscale, **body.dict()) - return {"trigger_id": trigger_id, "trigger_type": trigger_type} + + queue_wait_size = taskqueue.wait_queue_size() + queue_concur_size = taskqueue.concur_queue_size() + + return {"trigger_id": trigger_id, "trigger_type": trigger_type, + "wait_size": queue_wait_size, + "concur_size": queue_concur_size} @router.post("/variation", response_model=TriggerResponse) @@ -49,7 +55,13 @@ async def variation(body: TriggerUVIn): trigger_type = TriggerType.variation.value taskqueue.put(trigger_id, discord.variation, **body.dict()) - return {"trigger_id": trigger_id, "trigger_type": trigger_type} + + queue_wait_size = taskqueue.wait_queue_size() + queue_concur_size = taskqueue.concur_queue_size() + + return {"trigger_id": trigger_id, "trigger_type": trigger_type, + "wait_size": queue_wait_size, + "concur_size": queue_concur_size} @router.post("/reset", response_model=TriggerResponse) @@ -58,7 +70,13 @@ async def reset(body: TriggerResetIn): trigger_type = TriggerType.reset.value taskqueue.put(trigger_id, discord.reset, **body.dict()) - return {"trigger_id": trigger_id, "trigger_type": trigger_type} + + queue_wait_size = taskqueue.wait_queue_size() + queue_concur_size = taskqueue.concur_queue_size() + + return {"trigger_id": trigger_id, "trigger_type": trigger_type, + "wait_size": queue_wait_size, + "concur_size": queue_concur_size} @router.post("/describe", response_model=TriggerResponse) @@ -67,7 +85,13 @@ async def describe(body: TriggerDescribeIn): trigger_type = TriggerType.describe.value taskqueue.put(trigger_id, discord.describe, **body.dict()) - return {"trigger_id": trigger_id, "trigger_type": trigger_type} + + queue_wait_size = taskqueue.wait_queue_size() + queue_concur_size = taskqueue.concur_queue_size() + + return {"trigger_id": trigger_id, "trigger_type": trigger_type, + "wait_size": queue_wait_size, + "concur_size": queue_concur_size} @router.post("/upload", response_model=UploadResponse) @@ -103,4 +127,9 @@ async def queue_release(body: QueueReleaseIn): """bot 清除队列任务""" taskqueue.pop(body.trigger_id) - return body + queue_wait_size = taskqueue.wait_queue_size() + queue_concur_size = taskqueue.concur_queue_size() + + return {"trigger_id": body.trigger_id, "trigger_type": "release_queue", + "wait_size": queue_wait_size, + "concur_size": queue_concur_size} From feed62cf2e0adeb85fe311e305c307bf1c9ef7b5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=A4=D0=B8=D0=BB=D0=B0=D1=82=D0=BE=D0=B2=20=D0=90=D0=BD?= =?UTF-8?q?=D1=82=D0=BE=D0=BD=20=D0=92=D0=BB=D0=B0=D0=B4=D0=B8=D0=BC=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B8=D1=87?= Date: Fri, 4 Aug 2023 14:37:18 +0300 Subject: [PATCH 6/6] Implemented blend command --- app/routers.py | 17 +++++++++++++++++ app/schema.py | 6 ++++++ lib/api/discord.py | 33 +++++++++++++++++++++++++++++++++ requirements.txt | 3 ++- task/bot/handler.py | 12 +++++++++++- 5 files changed, 69 insertions(+), 2 deletions(-) diff --git a/app/routers.py b/app/routers.py index 04411aa..ca1f160 100644 --- a/app/routers.py +++ b/app/routers.py @@ -12,6 +12,7 @@ TriggerResponse, UploadResponse, TriggerDescribeIn, + TriggerBlendIn, SendMessageResponse, SendMessageIn, ) @@ -94,6 +95,22 @@ async def describe(body: TriggerDescribeIn): "concur_size": queue_concur_size} +@router.post("/blend", response_model=TriggerResponse) +async def describe(body: TriggerBlendIn): + trigger_id = body.trigger_id + trigger_type = TriggerType.blend.value + + taskqueue.put(trigger_id, discord.blend, **body.dict()) + + queue_wait_size = taskqueue.wait_queue_size() + queue_concur_size = taskqueue.concur_queue_size() + + return {"trigger_id": trigger_id, "trigger_type": trigger_type, + "wait_size": queue_wait_size, + "concur_size": queue_concur_size + } + + @router.post("/upload", response_model=UploadResponse) async def upload_attachment(file: UploadFile): if not file.content_type.startswith("image/"): diff --git a/app/schema.py b/app/schema.py index a11c8c4..c8749a0 100644 --- a/app/schema.py +++ b/app/schema.py @@ -28,6 +28,12 @@ class TriggerDescribeIn(BaseModel): trigger_id: str +class TriggerBlendIn(BaseModel): + upload_filename_1: str + upload_filename_2: str + trigger_id: str + + class QueueReleaseIn(BaseModel): trigger_id: str diff --git a/lib/api/discord.py b/lib/api/discord.py index 7e8b37d..e8bffd4 100644 --- a/lib/api/discord.py +++ b/lib/api/discord.py @@ -23,6 +23,7 @@ class TriggerType(str, Enum): max_upscale = "max_upscale" reset = "reset" describe = "describe" + blend = "blend" async def trigger(payload: Dict[str, Any]): @@ -188,3 +189,35 @@ async def describe(upload_filename: str, **kwargs): }] }) return await trigger(payload) + + +async def blend(upload_filename_1: str,upload_filename_2: str, **kwargs): + payload = _trigger_payload(2, { + "version": "1118961510123847773", + "id": "1062880104792997970", + "name": "blend", + "type": 1, + "options": [ + { + "type": 11, + "name": "image1", + "value": 0 + }, + { + "type": 11, + "name": "image2", + "value": 1 + } + ], + "attachments": [{ + "id": "0", + "filename": upload_filename_1.split("/")[-1], + "uploaded_filename": upload_filename_1, + }, + { + "id": "1", + "filename": upload_filename_2.split("/")[-1], + "uploaded_filename": upload_filename_2, + }] + }) + return await trigger(payload) diff --git a/requirements.txt b/requirements.txt index 23c1e4c..6837f9b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,4 +6,5 @@ pydantic~=1.10.7 aiohttp~=3.8.4 loguru==0.7.0 aiofiles==23.1.0 -python-multipart==0.0.6 \ No newline at end of file +python-multipart==0.0.6 +requests==2.31.0 \ No newline at end of file diff --git a/task/bot/handler.py b/task/bot/handler.py index 41c752c..74ba566 100644 --- a/task/bot/handler.py +++ b/task/bot/handler.py @@ -2,12 +2,15 @@ import re from typing import Dict, Union, Any +import requests from discord import Message from app.handler import PROMPT_PREFIX, PROMPT_SUFFIX from lib.api.callback import queue_release, callback from task.bot._typing import CallbackData, Attachment, Embed +from requests import request + TRIGGER_ID_PATTERN = f"{PROMPT_PREFIX}(\w+?){PROMPT_SUFFIX}" # 消息 ID 正则 TEMP_MAP: Dict[str, bool] = {} # 临时存储消息流转信息 @@ -31,7 +34,14 @@ def pop_temp(trigger_id: str): def match_trigger_id(content: str) -> Union[str, None]: match = re.findall(TRIGGER_ID_PATTERN, content) - return match[0] if match else None + if match: + return match[0] + if 'https://s.mj.run' in content: + match2 = re.findall('(?<=\*\*<).+(?=> <)', content) + r = requests.get(match2[0]) + return r.url.split("/")[-1].split(".")[0] + + else: return None async def callback_trigger(trigger_id: str, trigger_status: str, message: Message):