Model_Manager/Model_api.py at main · linsixi/Model_Manager · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import requests
import json
from io import BytesIO
import matplotlib.pyplot as plt
from PIL import Image
import base64
import urllib


def get_access_token_qianfan():
    url = "https://aip.baidubce.com/oauth/2.0/token?grant_type=client_credentials&client_id=DTNX9fIESYPHsOLYRSgGEPBJ&client_secret=PUfU7ew9kXJOBpZRzlmkY3OBRZ7ksB1m"
    payload = json.dumps("")
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)
    return response.json().get("access_token")


def get_zidongtaichu( question ):
    # 紫东太初文本生成
    api = 'https://ai-maas.wair.ac.cn/maas/v1/chat/completions'
    headers = {'Authorization': 'Bearer tlf3tc8sltk89etyx0p16u5p'}
    params = {
        'model': 'taichu_llm',
        'messages': [{"role": "user", "content": f"{question}"}],
        'stream': False
    }

    response = requests.post(api, json=params, headers=headers, stream=True)
    if response.status_code == 200:
        response_json = response.json()
        content = response_json['choices'][0]['message']['content']
        return content
    else:
        body = response.content.decode('utf-8')
        print(f'request failed,status_code:{response.status_code},body:{body}')


def get_qianfan_text( question ):
    # 千帆文本生成
    access_token = get_access_token_qianfan()
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/completions_pro?access_token=" + str(access_token)
    payload = json.dumps({
        "messages": [
            {
                "role": "user",
                "content": f"{question}"
            }
        ]
    })
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data = payload )
    if response.status_code == 200:
        response_json = response.json()
        content = response_json[ 'result' ]
        return content
    else:
        body = response.content.decode('utf-8')
        print(f'request failed,status_code:{response.status_code},body:{body}')


def get_qianfan_graph( question ):
    access_token = get_access_token_qianfan()
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/text2image/sd_xl?access_token=" + str( access_token )
    payload = json.dumps({
        "prompt": f"{question}",
        "size": "1024x1024",
        "n": 1,
        "steps": 20,
        "sampler_index": "Euler a"
    })
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json'
    }

    response = requests.request("POST", url, headers=headers, data=payload)
    if response.status_code != 200:
        body = response.content.decode('utf-8')
        print(f'request failed,status_code:{response.status_code},body:{body}')
    response_dict = response.json()
    b64_image = response_dict['data'][0]['b64_image']

    # plt显示图片
    image_data = base64.b64decode( b64_image )
    image = Image.open(BytesIO(image_data))
    plt.imshow( image )
    plt.axis('off')
    plt.show()


def get_qianfan_read( question ):  # 这里只能传入图片的地址
    access_token = get_access_token_qianfan()
    url = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/image2text/fuyu_8b?access_token=" + access_token
    question = image_to_base64( question )

    payload = json.dumps( {
        "prompt": "解释一下这张图片",
        "image": f"{question}"
    } )
    headers = {
        'Content-Type': 'application/json'
    }

    response = requests.request( "POST", url, headers = headers, data = payload )
    response_dict = response.json()
    return response_dict[ 'result' ]


def image_to_base64(image_path):
    with open( image_path, "rb" ) as image_file:
        image_data = image_file.read()
        base64_encoded = base64.b64encode( image_data ).decode( 'utf-8' )
        return base64_encoded


def api_check(n, question,question_text):#question_text是大模型提示词

    if n == '1':
        result = get_qianfan_text( question,question_text )
    elif n == '2':
        result = get_zidongtaichu( question,question_text )
    elif n == '3':
        result = get_qianfan_graph( question,question_text )
    elif n == '4':
        result = get_qianfan_read( question,question_text )
    print( result )