From bda7252a7d575a77c77abba5fdc52e5215712fcd Mon Sep 17 00:00:00 2001
From: yokon <15152347277@163.com>
Date: Mon, 5 Jun 2023 17:47:26 +0800
Subject: [PATCH] fix: prompt error #28

---
 README.md      | 23 +++++++++++++++++++++--
 app/handler.py |  9 +++++++--
 app/routers.py |  5 +++--
 app/schema.py  |  3 +++
 util/_queue.py |  5 ++---
 5 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/README.md b/README.md
index 65cd029..5f73799 100644
--- 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
index 462b832..450387e 100644
--- 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
index 44042e4..e92ad30 100644
--- a/app/routers.py
+++ b/app/routers.py
@@ -12,7 +12,8 @@
     TriggerResponse,
     UploadResponse,
     TriggerDescribeIn,
-    SendMessageResponse, SendMessageIn,
+    SendMessageResponse,
+    SendMessageIn,
 )
 
 router = APIRouter()
@@ -20,7 +21,7 @@
 
 @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
index be722df..2d2b9f3 100644
--- 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
index 7856ecd..1ed0639 100644
--- 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),
 )