fix: prompt error yokonsan#28

vinfai · Jun 5, 2023 · bda7252 · bda7252
1 parent ecbfee4
commit bda7252
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 9 deletions.
diff --git a/README.md b/README.md
@@ -2,8 +2,6 @@
 
 基于 Discord 的 Midjourney API。
 
-使用教程参考：[Midjourney｜如何集成到自己的平台](https://mp.weixin.qq.com/s?__biz=Mzg4MjkzMzc1Mg==&mid=2247484029&idx=1&sn=d3c458bba9459f19f05d13ab23f5f67e&chksm=cf4e68eaf839e1fc2db025bd9940d0f5e57862f1788c88215b4a66cb23f553a30c5f37ac3ae8&token=79614426&lang=zh_CN#rd)
-
 **添加 Midjourney 违禁词入口 [issue](https://github.com/yokonsan/midjourney-api/issues/new?assignees=&labels=banned+prompt&projects=&template=banned_prompt_report.yml&title=Banned+prompt%3A+)**
 
 
@@ -26,6 +24,12 @@ sequenceDiagram
     DiscordAPI-->>APIServer: 清除队列任务
 ```
 
+## 使用条件
+
+1. 确保程序启动环境能访问 Discord
+2. 已有 Midjourney、Discord 账户
+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)
+
 
 ## 安装启动
 
@@ -105,6 +109,8 @@ sh start.sh
 
 ### imagine
 
+文生图
+
 ```bash
 curl -X 'POST' \
   'http://127.0.0.1:8062/v1/api/trigger/imagine' \
@@ -115,6 +121,19 @@ curl -X 'POST' \
 }'
 ```
 
+图生图，需带上图片 URL
+
+```bash
+curl -X 'POST' \
+  'http://127.0.0.1:8062/v1/api/trigger/imagine' \
+  -H 'accept: application/json' \
+  -H 'Content-Type: application/json' \
+  -d '{
+  "prompt": "a cute cat",
+  "picurl": "https://xxxxxx/xxxxxxxxxxxx.jpg"
+}'
+```
+
 ### upscale
 
 ```bash

diff --git a/app/handler.py b/app/handler.py
@@ -1,6 +1,7 @@
 import hashlib
 import time
 from functools import wraps
+from typing import Union
 
 from fastapi import status
 from fastapi.responses import JSONResponse
@@ -23,14 +24,18 @@ def unique_id():
     return int(hashlib.sha256(str(time.time()).encode("utf-8")).hexdigest(), 16) % 10**10
 
 
-def prompt_handler(prompt: str):
+def prompt_handler(prompt: str, picurl: Union[str, None] = None):
     """
     拼接 Prompt 形如: <#1234567890#>a cute cat
     """
     check_banned(prompt)
 
     trigger_id = str(unique_id())
-    return trigger_id, f"{PROMPT_PREFIX}{trigger_id}{PROMPT_SUFFIX}{prompt}"
+
+    if not picurl and prompt.startswith(("http://", "https://")):
+        picurl, _, prompt = prompt.partition(" ")
+
+    return trigger_id, f"{picurl+' ' if picurl else ''}{PROMPT_PREFIX}{trigger_id}{PROMPT_SUFFIX}{prompt}"
 
 
 def http_response(func):

diff --git a/app/routers.py b/app/routers.py
@@ -12,15 +12,16 @@
     TriggerResponse,
     UploadResponse,
     TriggerDescribeIn,
-    SendMessageResponse, SendMessageIn,
+    SendMessageResponse,
+    SendMessageIn,
 )
 
 router = APIRouter()
 
 
 @router.post("/imagine", response_model=TriggerResponse)
 async def imagine(body: TriggerImagineIn):
-    trigger_id, prompt = prompt_handler(body.prompt)
+    trigger_id, prompt = prompt_handler(body.prompt, body.picurl)
     trigger_type = TriggerType.generate.value
 
     taskqueue.put(trigger_id, discord.generate, prompt)

diff --git a/app/schema.py b/app/schema.py
@@ -1,8 +1,11 @@
+from typing import Optional
+
 from pydantic import BaseModel
 
 
 class TriggerImagineIn(BaseModel):
     prompt: str
+    picurl: Optional[str]
 
 
 class TriggerUVIn(BaseModel):

diff --git a/util/_queue.py b/util/_queue.py
@@ -64,7 +64,6 @@ def _exec(self):
         #     lambda t: print(t.result())
         # )  # todo
 
-
     def concur_size(self):
         return self._concur_size
 
@@ -79,6 +78,6 @@ def clear_concur(self):
 
 
 taskqueue = TaskQueue(
-    int(getenv("CONCUR_SIZE") or 3),
-    int(getenv("WAIT_SIZE") or 10),
+    int(getenv("CONCUR_SIZE") or 9999),
+    int(getenv("WAIT_SIZE") or 9999),
 )