This commit is contained in:
eric
2023-10-28 21:33:32 +08:00
parent 1eff7fd145
commit c3f5a3a227

View File

@@ -15,6 +15,17 @@ import json
import base64 import base64
import requests 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): def submit_post(url: str, data: dict):
""" """
Submit a POST request to the given URL with the given data. Submit a POST request to the given URL with the given data.
@@ -140,10 +151,10 @@ async def callback(
try: try:
query = unquote(query["query"], "utf-8") query = unquote(query["query"], "utf-8")
prompt = query #问题 prompt = query #问题
except: except:
prompt = "" prompt = ""
max_length=None max_length=None
top_p=None top_p=None
temperature=None temperature=None
@@ -184,6 +195,10 @@ async def img(
# print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid) # print("console :,", "appid", appid, "bid", bid, "requestid", requestid, "uid", uid)
print('query',query) print('query',query)
print('query[query]',query["query"]) print('query[query]',query["query"])
try:
sendTg(query["query"])
except:
print('发送tg失败')
# 获取请求参数 # 获取请求参数
global model, tokenizer global model, tokenizer
@@ -202,17 +217,17 @@ async def img(
query=query.replace('/img', '') query=query.replace('/img', '')
prompt=query 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} data = {'prompt': prompt}
res = submit_post(txt2img_url, data) 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') 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) print('response',response)
# return responseimg # return responseimg
else: else:
prompt = query #问题 prompt = query #问题
# 请求ai模型 # 请求ai模型
@@ -227,7 +242,7 @@ async def img(
print('ValueError',ValueError) print('ValueError',ValueError)
print('ArithmeticError',ArithmeticError) print('ArithmeticError',ArithmeticError)
prompt = "" prompt = ""
# log日志 # log日志
# now = datetime.datetime.now() # now = datetime.datetime.now()
# time = now.strftime("%Y-%m-%d %H:%M:%S") # time = now.strftime("%Y-%m-%d %H:%M:%S")
@@ -242,9 +257,9 @@ async def img(
if __name__ == '__main__': #if __name__ == '__main__':
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True) #tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", 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", trust_remote_code=True).cuda()
# 多显卡支持使用下面三行代替上面两行将num_gpus改为你实际的显卡数量 # 多显卡支持使用下面三行代替上面两行将num_gpus改为你实际的显卡数量
# model_path = "THUDM/chatglm2-6b" # model_path = "THUDM/chatglm2-6b"
@@ -252,4 +267,4 @@ if __name__ == '__main__':
# model = load_model_on_gpus(model_path, num_gpus=2) # model = load_model_on_gpus(model_path, num_gpus=2)
# model.eval() # model.eval()
uvicorn.run(app, host='0.0.0.0', port=8000, workers=1) # uvicorn.run(app, host='0.0.0.0', port=8000, workers=1