diff --git a/callback.py b/callback.py index 99f394a..0035518 100644 --- a/callback.py +++ b/callback.py @@ -182,31 +182,36 @@ async def img( max_length=None top_p=None temperature=None + response=None try: query = unquote(query["query"], "utf-8") # 判断有没有'/img'标识符 - if ('/img' in query): + if ('/img' in query) is True: 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() + now =str(time.time()) save_encoded_image(res.json()['images'][0], f'img/{now}.png') response=f'http://aiimg.yidooplanet.com/{now}.png' + print('response',response) + # return responseimg 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) + # 请求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: + except(ValueError, ArithmeticError): + print('ValueError',ValueError) + print('ArithmeticError',ArithmeticError) prompt = "" @@ -228,10 +233,12 @@ 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).cuda() + model = AutoModel.from_pretrained("THUDM/chatglm2-6b-int4", trust_remote_code=True).half().quantize(8).cuda() + # 多显卡支持,使用下面三行代替上面两行,将num_gpus改为你实际的显卡数量 # model_path = "THUDM/chatglm2-6b" # tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) # 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) \ No newline at end of file