Add docker-compose.yaml

Dockerfile

@@ -1,13 +1,19 @@
-FROM python:3.10.6
-LABEL creator="yokon" email="[email protected]"
+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"]
+CMD ["http"]

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}

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):

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

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)

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
+python-multipart==0.0.6
+requests==2.31.0

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):

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),