GPT4Tools / app.py
stevengrove
change layout
dc0dba4
raw
history blame contribute delete
No virus
6.55 kB
import os
import re
import pickle
import base64
import requests
import argparse
import numpy as np
import gradio as gr
from functools import partial
from PIL import Image
SERVER_URL = os.getenv('SERVER_URL')
def get_images(state):
history = ''
for i in range(len(state)):
for j in range(len(state[i])):
history += state[i][j] + '\n'
for image_path in re.findall('image/[0-9,a-z]+\.png', history):
if os.path.exists(image_path):
continue
data = {'method': 'get_image', 'args': [image_path], 'kwargs': {}}
data = base64.b64encode(pickle.dumps(data)).decode('utf-8')
response = requests.post(SERVER_URL, json=data)
image = pickle.loads(base64.b64decode(response.json().encode('utf-8')))
image.save(image_path)
def bot_request(method, *args, **kwargs):
data = {'method': method, 'args': args, 'kwargs': kwargs}
data = base64.b64encode(pickle.dumps(data)).decode('utf-8')
response = requests.post(SERVER_URL, json=data)
response = pickle.loads(base64.b64decode(response.json().encode('utf-8')))
if response is not None:
state = response[0]
get_images(state)
return response
def run_image(image, *args, **kwargs):
if image is not None:
width, height = image.size
ratio = min(512 / width, 512 / height)
width_new, height_new = (round(width * ratio), round(height * ratio))
width_new = int(np.round(width_new / 64.0)) * 64
height_new = int(np.round(height_new / 64.0)) * 64
image = image.resize((width_new, height_new))
image = image.convert('RGB')
return bot_request('run_image', image, *args, **kwargs)
def predict_example(temperature, top_p, max_new_token, keep_last_n_paragraphs, image, text):
state = []
buffer = ''
chatbot, state, text, buffer = run_image(image, state, text, buffer)
chatbot, state, text, buffer = bot_request(
'run_text', text, state, temperature, top_p,
max_new_token, keep_last_n_paragraphs, buffer)
return chatbot, state, text, None, buffer
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--temperature', type=float, default=0.0, help='temperature for the llm model')
parser.add_argument('--max_new_tokens', type=int, default=256, help='max number of new tokens to generate')
parser.add_argument('--top_p', type=float, default=1.0, help='top_p for the llm model')
parser.add_argument('--top_k', type=int, default=40, help='top_k for the llm model')
parser.add_argument('--keep_last_n_paragraphs', type=int, default=0, help='keep last n paragraphs in the memory')
args = parser.parse_args()
examples = [
['images/example-1.jpg', 'What is unusual about this image?'],
['images/example-2.jpg', 'Make the image look like a cartoon.'],
['images/example-3.jpg', 'Segment the tie in the image.'],
['images/example-4.jpg', 'Generate a man watching a sea based on the pose of the woman.'],
['images/example-5.jpg', 'Replace the dog with a monkey.'],
]
if not os.path.exists('image'):
os.makedirs('image')
with gr.Blocks() as demo:
state = gr.Chatbot([], visible=False)
buffer = gr.Textbox('', visible=False)
with gr.Row():
with gr.Column(scale=0.3):
with gr.Row():
image = gr.Image(type='pil', label='input image')
with gr.Row():
txt = gr.Textbox(lines=7, show_label=False, elem_id='textbox',
placeholder='Enter text and press submit, or upload an image').style(container=False)
with gr.Row():
submit = gr.Button('Submit')
with gr.Row():
clear = gr.Button('Clear')
with gr.Row():
llm_name = gr.Radio(
["Vicuna-13B"],
label="LLM Backend",
value="Vicuna-13B",
interactive=True)
keep_last_n_paragraphs = gr.Slider(
minimum=0,
maximum=3,
value=args.keep_last_n_paragraphs,
step=1,
interactive=True,
label='Remember Last N Paragraphs')
max_new_token = gr.Slider(
minimum=64,
maximum=512,
value=args.max_new_tokens,
step=1,
interactive=True,
label='Max New Tokens')
temperature = gr.Slider(
minimum=0.0,
maximum=1.0,
value=args.temperature,
step=0.1,
interactive=True,
visible=False,
label='Temperature')
top_p = gr.Slider(
minimum=0.0,
maximum=1.0,
value=args.top_p,
step=0.1,
interactive=True,
visible=False,
label='Top P')
with gr.Column(scale=0.7):
chatbot = gr.Chatbot(elem_id='chatbot', label='πŸ¦™ GPT4Tools').style(height=690)
image.upload(lambda: '', None, txt)
submit.click(run_image,
[image, state, txt, buffer],
[chatbot, state, txt, buffer]).then(
partial(bot_request, 'run_text'),
[txt, state, temperature, top_p, max_new_token, keep_last_n_paragraphs, buffer],
[chatbot, state, txt, buffer]).then(
lambda: None, None, image)
clear.click(partial(bot_request, 'clear'))
clear.click(lambda: [[], [], '', ''], None, [chatbot, state, txt, buffer])
with gr.Row():
gr.Examples(
examples=examples,
fn=partial(predict_example, args.temperature, args.top_p,
args.max_new_tokens, args.keep_last_n_paragraphs),
inputs=[image, txt],
outputs=[chatbot, state, txt, image, buffer],
cache_examples=True,
)
demo.queue(concurrency_count=6)
demo.launch()