diff --git a/Dockerfile b/Dockerfile index 934b1fb..a6154af 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.6 AS builder -WORKDIR /app +COPY requirements.txt . +RUN pip install --user -r requirements.txt +FROM python:3.10.6-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"] diff --git a/app/routers.py b/app/routers.py index e92ad30..ca1f160 100644 --- a/app/routers.py +++ b/app/routers.py @@ -12,6 +12,7 @@ TriggerResponse, UploadResponse, TriggerDescribeIn, + TriggerBlendIn, SendMessageResponse, SendMessageIn, ) @@ -25,7 +26,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) @@ -34,7 +41,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) @@ -43,7 +56,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) @@ -52,7 +71,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) @@ -61,7 +86,29 @@ 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("/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) @@ -97,4 +144,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} diff --git a/app/schema.py b/app/schema.py index 2d2b9f3..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 @@ -36,6 +42,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/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 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): 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),