feat: ai生成图片

callback.py

@@ -11,6 +11,26 @@ import time
 from transformers import AutoTokenizer, AutoModel
 import uvicorn, json, datetime
 import torch
+import json
+import base64
+import requests
+
+def submit_post(url: str, data: dict):
+    """
+    Submit a POST request to the given URL with the given data.
+    """
+    return requests.post(url, data=json.dumps(data))
+
+
+def save_encoded_image(b64_image: str, output_path: str):
+    """
+    Save the given image to the given output path.
+    """
+    with open(output_path, "wb") as image_file:
+        image_file.write(base64.b64decode(b64_image))
+
+
+
 
 DEVICE = "cuda"
 DEVICE_ID = "0"
@@ -86,7 +106,7 @@ async def create_item(request: Request):
     return answer
 
 
-# 通过网址发送闪信
+# AI智能对话
 @app.post("/callback")
 async def callback(
     # item: Item,
@@ -137,6 +157,75 @@ async def callback(
     headers = {"Content-Type": "application/json"}
     return JSONResponse(content=content, headers=headers)
 
+
+
+# AI生成图片
+@app.post("/img")
+async def img(
+    # item: Item,
+    # inputArguments: Union[str, None] = None,
+    query=Body(None),
+    appid: Union[str, None] = Header(default=None),
+    bid: Union[str, None] = Header(default=None),
+    requestid: Union[str, None] = Header(default=None),
+    uid: Union[str, None] = Header(default=None),
+):
+    print('query',query)
+    print('query[query]',query["query"])
+
+    # 获取请求参数
+    global model, tokenizer
+    # time.sleep(10)
+    # print("item", item)
+    history = []  #默认空数组
+    prompt=''
+    max_length=None
+    top_p=None
+    temperature=None
+    try:
+        query = unquote(query["query"], "utf-8")
+        # 判断有没有'/img'标识符
+        if ('/img' in query):
+            query=query.replace('/img', '')
+            prompt=query
+            txt2img_url = 'http://127.0.0.1:7861/sdapi/v1/txt2img'
+            data = {'prompt': prompt}
+            res = submit_post(txt2img_url, data)
+            # 以时间戳命名
+            now = datetime.datetime.now()
+            save_encoded_image(res.json()['images'][0], f'img/{now}.png')
+            response=f'http://aiimg.yidooplanet.com/{now}.png'
+        
+        else:
+            prompt = query #问题
+        # 请求ai模型
+        response, history = model.chat(tokenizer,
+                                   prompt,
+                                   history=history,
+                                   max_length=max_length if max_length else 2048,
+                                   top_p=top_p if top_p else 0.7,
+                                   temperature=temperature if temperature else 0.95)
+
+    except:
+        prompt = ""
+   
+
+
+
+    # log日志
+    # now = datetime.datetime.now()
+    # time = now.strftime("%Y-%m-%d %H:%M:%S")
+    # log = "[" + time + "] " + '", prompt:"' + prompt + '", response:"' + repr(response) + '"'
+    # print(log)
+    content = {
+        "answer": response,
+        "answer_type": "text",
+    }
+    headers = {"Content-Type": "application/json"}
+    return JSONResponse(content=content, headers=headers)
+
+
+
 if __name__ == '__main__':
     tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
     model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).cuda()