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 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.
@@ -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,14 +217,14 @@ 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())
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
@@ -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)
# uvicorn.run(app, host='0.0.0.0', port=8000, workers=1