muxingyin/VisualGLM-6B

1import os2import json3from transformers import AutoTokenizer, AutoModel4import uvicorn5from fastapi import FastAPI, Request6import datetime7from model import process_image8import torch9 10tokenizer = AutoTokenizer.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True)11model = AutoModel.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True).half().cuda()12 13 14app = FastAPI()15@app.post('/')16async def visual_glm(request: Request):17    json_post_raw = await request.json()18    print("Start to process request")19 20    json_post = json.dumps(json_post_raw)21    request_data = json.loads(json_post)22 23    history = request_data.get("history")24    image_encoded = request_data.get("image")25    query = request_data.get("text")26    image_path = process_image(image_encoded)27 28    with torch.no_grad():    29        result = model.stream_chat(tokenizer, image_path, query, history=history)30    last_result = None31    for value in result:32        last_result = value33    answer = last_result[0]34 35    if os.path.isfile(image_path):36        os.remove(image_path)37    now = datetime.datetime.now()38    time = now.strftime("%Y-%m-%d %H:%M:%S")39    response = {40        "result": answer,41        "history": history,42        "status": 200,43        "time": time44    }45    return response46 47 48if __name__ == "__main__":49   uvicorn.run(app, host='0.0.0.0', port=8080, workers=1)