From c3f5a3a227aa2d6e98ab973295b091c686b216c4 Mon Sep 17 00:00:00 2001
From: eric <xiaoshuang.eric@gmail.com>
Date: Sat, 28 Oct 2023 21:33:32 +0800
Subject: [PATCH] update

---
 callback.py | 37 ++++++++++++++++++++++++++-----------
 1 file changed, 26 insertions(+), 11 deletions(-)

diff --git a/callback.py b/callback.py
index 9470d27..6df3f87 100644
--- a/callback.py
+++ b/callback.py
@@ -15,6 +15,17 @@ import json
 import base64
 import requests
 
+#使用int8的量化模型
+tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True)
+#model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).half().quantize(8).cuda()
+model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4",trust_remote_code=True).cuda()
+model.eval()
+
+#显存满配全开
+#tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
+#model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
+#model.eval()
+
 def submit_post(url: str, data: dict):
     """
     Submit a POST request to the given URL with the given data.
@@ -140,10 +151,10 @@ async def callback(
     try:
         query = unquote(query["query"], "utf-8")
         prompt = query #问题
-        
+
     except:
         prompt = ""
-   
+
     max_length=None
     top_p=None
     temperature=None
@@ -184,6 +195,10 @@ async def img(
     # print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid)
     print('query',query)
     print('query[query]',query["query"])
+    try:
+        sendTg(query["query"])
+    except:
+        print('发送tg失败')
 
     # 获取请求参数
     global model, tokenizer
@@ -202,17 +217,17 @@ async def img(
             query=query.replace('/img', '')
             prompt=query
 
-            txt2img_url = 'http://127.0.0.1:7861/sdapi/v1/txt2img'
+            txt2img_url = 'http://127.0.0.1:7860/sdapi/v1/txt2img'
             data = {'prompt': prompt}
             res = submit_post(txt2img_url, data)
             # 以时间戳命名
-            now =str(time.time()) 
+            now =str(time.time())
             save_encoded_image(res.json()['images'][0], f'img/{now}.png')
             # 网络图片
-            response=f'#web_img http://aiimg.yidooplanet.com/{now}.png'
+            response=f'#web_img http://aiimg.hackrobot.cn/{now}.png'
             print('response',response)
             # return responseimg
-        
+
         else:
             prompt = query #问题
             # 请求ai模型
@@ -227,7 +242,7 @@ async def img(
         print('ValueError',ValueError)
         print('ArithmeticError',ArithmeticError)
         prompt = ""
-   
+
     # log日志
     # now = datetime.datetime.now()
     # time = now.strftime("%Y-%m-%d %H:%M:%S")
@@ -242,9 +257,9 @@ async def img(
 
 
 
-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).half().quantize(8).cuda()
+#if __name__ == '__main__':
+    #tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
+    #model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()
 
     # 多显卡支持，使用下面三行代替上面两行，将num_gpus改为你实际的显卡数量
     # model_path = "THUDM/chatglm2-6b"
@@ -252,4 +267,4 @@ if __name__ == '__main__':
     # model = load_model_on_gpus(model_path, num_gpus=2)
 
     # model.eval()
-    uvicorn.run(app, host='0.0.0.0', port=8000, workers=1)
\ No newline at end of file
+   # uvicorn.run(app, host='0.0.0.0', port=8000, workers=1
\ No newline at end of file