|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
import base64 |
|
import datetime |
|
import hashlib |
|
import json |
|
import os |
|
import random |
|
import re |
|
import sys |
|
|
|
from functools import partial |
|
from io import BytesIO |
|
|
|
import cv2 |
|
import numpy as np |
|
import requests |
|
import streamlit as st |
|
from constants import LOGDIR, server_error_msg |
|
from library import Library |
|
from PIL import Image, ImageDraw, ImageFont |
|
from streamlit_image_select import image_select |
|
|
|
custom_args = sys.argv[1:] |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('--controller_url', type=str, default='http://10.140.60.209:10075', help='url of the controller') |
|
parser.add_argument('--sd_worker_url', type=str, default='http://0.0.0.0:40006', help='url of the stable diffusion worker') |
|
parser.add_argument('--max_image_limit', type=int, default=4, help='maximum number of images') |
|
args = parser.parse_args(custom_args) |
|
controller_url = args.controller_url |
|
sd_worker_url = args.sd_worker_url |
|
max_image_limit = args.max_image_limit |
|
print('args:', args) |
|
|
|
|
|
def get_model_list(): |
|
ret = requests.post(controller_url + '/refresh_all_workers') |
|
assert ret.status_code == 200 |
|
ret = requests.post(controller_url + '/list_models') |
|
models = ret.json()['models'] |
|
return models |
|
|
|
|
|
def load_upload_file_and_show(): |
|
if uploaded_files is not None: |
|
images = [] |
|
for file in uploaded_files: |
|
file_bytes = np.asarray(bytearray(file.read()), dtype=np.uint8) |
|
img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR) |
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
|
img = Image.fromarray(img) |
|
images.append(img) |
|
with upload_image_preview.container(): |
|
Library(images) |
|
|
|
image_hashes = [hashlib.md5(image.tobytes()).hexdigest() for image in images] |
|
for image, hash in zip(images, image_hashes): |
|
t = datetime.datetime.now() |
|
filename = os.path.join(LOGDIR, 'serve_images', f'{t.year}-{t.month:02d}-{t.day:02d}', f'{hash}.jpg') |
|
if not os.path.isfile(filename): |
|
os.makedirs(os.path.dirname(filename), exist_ok=True) |
|
image.save(filename) |
|
return images |
|
|
|
|
|
def get_selected_worker_ip(): |
|
ret = requests.post(controller_url + '/get_worker_address', |
|
json={'model': selected_model}) |
|
worker_addr = ret.json()['address'] |
|
return worker_addr |
|
|
|
|
|
def generate_response(messages): |
|
send_messages = [{'role': 'system', 'content': persona_rec}] |
|
for message in messages: |
|
if message['role'] == 'user': |
|
user_message = {'role': 'user', 'content': message['content']} |
|
if 'image' in message and len('image') > 0: |
|
user_message['image'] = [] |
|
for image in message['image']: |
|
user_message['image'].append(pil_image_to_base64(image)) |
|
send_messages.append(user_message) |
|
else: |
|
send_messages.append({'role': 'assistant', 'content': message['content']}) |
|
pload = { |
|
'model': selected_model, |
|
'prompt': send_messages, |
|
'temperature': float(temperature), |
|
'top_p': float(top_p), |
|
'max_new_tokens': max_length, |
|
'max_input_tiles': max_input_tiles, |
|
'repetition_penalty': float(repetition_penalty), |
|
} |
|
worker_addr = get_selected_worker_ip() |
|
headers = {'User-Agent': 'InternVL-Chat Client'} |
|
placeholder, output = st.empty(), '' |
|
try: |
|
response = requests.post(worker_addr + '/worker_generate_stream', |
|
headers=headers, json=pload, stream=True, timeout=10) |
|
for chunk in response.iter_lines(decode_unicode=False, delimiter=b'\0'): |
|
if chunk: |
|
data = json.loads(chunk.decode()) |
|
if data['error_code'] == 0: |
|
output = data['text'] |
|
|
|
if '4B' in selected_model and '�' in output[-2:]: |
|
output = output.replace('�', '') |
|
break |
|
placeholder.markdown(output + '▌') |
|
else: |
|
output = data['text'] + f" (error_code: {data['error_code']})" |
|
placeholder.markdown(output) |
|
placeholder.markdown(output) |
|
except requests.exceptions.RequestException as e: |
|
placeholder.markdown(server_error_msg) |
|
return output |
|
|
|
|
|
def pil_image_to_base64(image): |
|
buffered = BytesIO() |
|
image.save(buffered, format='PNG') |
|
return base64.b64encode(buffered.getvalue()).decode('utf-8') |
|
|
|
|
|
def clear_chat_history(): |
|
st.session_state.messages = [] |
|
st.session_state['image_select'] = -1 |
|
|
|
|
|
def clear_file_uploader(): |
|
st.session_state.uploader_key += 1 |
|
st.rerun() |
|
|
|
|
|
def combined_func(func_list): |
|
for func in func_list: |
|
func() |
|
|
|
|
|
def show_one_or_multiple_images(message, total_image_num, is_input=True): |
|
if 'image' in message: |
|
if is_input: |
|
total_image_num = total_image_num + len(message['image']) |
|
if lan == 'English': |
|
if len(message['image']) == 1 and total_image_num == 1: |
|
label = f"(In this conversation, {len(message['image'])} image was uploaded, {total_image_num} image in total)" |
|
elif len(message['image']) == 1 and total_image_num > 1: |
|
label = f"(In this conversation, {len(message['image'])} image was uploaded, {total_image_num} images in total)" |
|
else: |
|
label = f"(In this conversation, {len(message['image'])} images were uploaded, {total_image_num} images in total)" |
|
else: |
|
label = f"(在本次对话中，上传了{len(message['image'])}张图片，总共上传了{total_image_num}张图片)" |
|
upload_image_preview = st.empty() |
|
with upload_image_preview.container(): |
|
Library(message['image']) |
|
if is_input and len(message['image']) > 0: |
|
st.markdown(label) |
|
|
|
|
|
def find_bounding_boxes(response): |
|
pattern = re.compile(r'<ref>\s*(.*?)\s*</ref>\s*<box>\s*(\[\[.*?\]\])\s*</box>') |
|
matches = pattern.findall(response) |
|
results = [] |
|
for match in matches: |
|
results.append((match[0], eval(match[1]))) |
|
returned_image = None |
|
for message in st.session_state.messages: |
|
if message['role'] == 'user' and 'image' in message and len(message['image']) > 0: |
|
last_image = message['image'][-1] |
|
width, height = last_image.size |
|
returned_image = last_image.copy() |
|
draw = ImageDraw.Draw(returned_image) |
|
for result in results: |
|
line_width = max(1, int(min(width, height) / 200)) |
|
random_color = (random.randint(0, 128), random.randint(0, 128), random.randint(0, 128)) |
|
category_name, coordinates = result |
|
coordinates = [(float(x[0]) / 1000, float(x[1]) / 1000, float(x[2]) / 1000, float(x[3]) / 1000) for x in coordinates] |
|
coordinates = [(int(x[0] * width), int(x[1] * height), int(x[2] * width), int(x[3] * height)) for x in coordinates] |
|
for box in coordinates: |
|
draw.rectangle(box, outline=random_color, width=line_width) |
|
font = ImageFont.truetype('static/SimHei.ttf', int(20 * line_width / 2)) |
|
text_size = font.getbbox(category_name) |
|
text_width, text_height = text_size[2] - text_size[0], text_size[3] - text_size[1] |
|
text_position = (box[0], max(0, box[1] - text_height)) |
|
draw.rectangle( |
|
[text_position, (text_position[0] + text_width, text_position[1] + text_height)], |
|
fill=random_color |
|
) |
|
draw.text(text_position, category_name, fill='white', font=font) |
|
return returned_image if len(matches) > 0 else None |
|
|
|
|
|
def query_image_generation(response, sd_worker_url, timeout=15): |
|
sd_worker_url = f'{sd_worker_url}/generate_image/' |
|
pattern = r'```drawing-instruction\n(.*?)\n```' |
|
match = re.search(pattern, response, re.DOTALL) |
|
if match: |
|
payload = {'caption': match.group(1)} |
|
print('drawing-instruction:', payload) |
|
response = requests.post(sd_worker_url, json=payload, timeout=timeout) |
|
response.raise_for_status() |
|
image = Image.open(BytesIO(response.content)) |
|
return image |
|
else: |
|
return None |
|
|
|
|
|
def regenerate(): |
|
st.session_state.messages = st.session_state.messages[:-1] |
|
st.rerun() |
|
|
|
|
|
logo_code = """ |
|
<svg width="1700" height="200" xmlns="http://www.w3.org/2000/svg"> |
|
<defs> |
|
<linearGradient id="gradient1" x1="0%" y1="0%" x2="100%" y2="0%"> |
|
<stop offset="0%" style="stop-color: red; stop-opacity: 1" /> |
|
<stop offset="100%" style="stop-color: orange; stop-opacity: 1" /> |
|
</linearGradient> |
|
</defs> |
|
<text x="000" y="160" font-size="180" font-weight="bold" fill="url(#gradient1)" style="font-family: Arial, sans-serif;"> |
|
InternVL2 Demo |
|
</text> |
|
</svg> |
|
""" |
|
|
|
|
|
st.set_page_config(page_title='InternVL2') |
|
|
|
if 'uploader_key' not in st.session_state: |
|
st.session_state.uploader_key = 0 |
|
|
|
|
|
system_message = """我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。人工智能实验室致力于原始技术创新，开源开放，共享共创，推动科技进步和产业发展。 |
|
|
|
对于目标检测任务，请按照以下格式输出坐标框：<ref>某类物体</ref><box>[[x1, y1, x2, y2], ...]</box> |
|
|
|
对于画画任务，请按照以下格式输出绘图指令，注意指令需要英文：```drawing-instruction\n{instruction}\n``` |
|
|
|
在处理输入包含多张图像的情况下，请严格按照以下规则区分和处理每一张图像，并小心区分用户的提问针对的是哪一张图片： |
|
|
|
1. 图像编号和标记：每张图像都将使用明确的编号标记，例如 "Image-1: <img></img>"，"Image-2: <img></img>"，"Image-3: <img></img>" 等等。 |
|
|
|
2. 用户提问关联：用户的提问可能会具体指向某一张编号的图像，请仔细辨别用户问题中提到的图像编号。 |
|
|
|
请尽可能详细地回答用户的问题。""" |
|
|
|
|
|
with st.sidebar: |
|
model_list = get_model_list() |
|
|
|
lan = st.selectbox('#### Language / 语言', ['English', '中文'], on_change=st.rerun) |
|
if lan == 'English': |
|
st.logo(logo_code, link='https://github.com/OpenGVLab/InternVL', icon_image=logo_code) |
|
st.subheader('Models and parameters') |
|
selected_model = st.sidebar.selectbox('Choose a InternVL2 chat model', model_list, key='selected_model', on_change=clear_chat_history) |
|
with st.expander('🤖 System Prompt'): |
|
persona_rec = st.text_area('System Prompt', value=system_message, |
|
help='System prompt is a pre-defined message used to instruct the assistant at the beginning of a conversation.', |
|
height=200) |
|
with st.expander('🔥 Advanced Options'): |
|
temperature = st.slider('temperature', min_value=0.0, max_value=1.0, value=0.8, step=0.1) |
|
top_p = st.slider('top_p', min_value=0.0, max_value=1.0, value=0.7, step=0.1) |
|
repetition_penalty = st.slider('repetition_penalty', min_value=1.0, max_value=1.5, value=1.1, step=0.02) |
|
max_length = st.slider('max_length', min_value=0, max_value=4096, value=2048, step=128) |
|
max_input_tiles = st.slider('max_input_tiles (control image resolution)', min_value=1, max_value=24, value=12, step=1) |
|
upload_image_preview = st.empty() |
|
uploaded_files = st.file_uploader('Upload files', accept_multiple_files=True, |
|
type=['png', 'jpg', 'jpeg', 'webp'], |
|
help='You can upload multiple images (max to 4) or a single video.', |
|
key=f'uploader_{st.session_state.uploader_key}', |
|
on_change=st.rerun) |
|
uploaded_pil_images = load_upload_file_and_show() |
|
else: |
|
st.subheader('模型和参数') |
|
selected_model = st.sidebar.selectbox('选择一个 InternVL2 对话模型', model_list, key='selected_model', on_change=clear_chat_history) |
|
with st.expander('🤖 系统提示'): |
|
persona_rec = st.text_area('系统提示', value=system_message, |
|
help='系统提示是在对话开始时用于指示助手的预定义消息。', |
|
height=200) |
|
with st.expander('🔥 高级选项'): |
|
temperature = st.slider('temperature', min_value=0.0, max_value=1.0, value=0.8, step=0.1) |
|
top_p = st.slider('top_p', min_value=0.0, max_value=1.0, value=0.7, step=0.1) |
|
repetition_penalty = st.slider('重复惩罚', min_value=1.0, max_value=1.5, value=1.1, step=0.02) |
|
max_length = st.slider('最大输出长度', min_value=0, max_value=4096, value=2048, step=128) |
|
max_input_tiles = st.slider('最大图像块数 (控制图像分辨率)', min_value=1, max_value=24, value=12, step=1) |
|
upload_image_preview = st.empty() |
|
uploaded_files = st.file_uploader('上传文件', accept_multiple_files=True, |
|
type=['png', 'jpg', 'jpeg', 'webp'], |
|
help='你可以上传多张图像（最多4张）或者一个视频。', |
|
key=f'uploader_{st.session_state.uploader_key}', |
|
on_change=st.rerun) |
|
uploaded_pil_images = load_upload_file_and_show() |
|
|
|
|
|
gradient_text_html = """ |
|
<style> |
|
.gradient-text { |
|
font-weight: bold; |
|
background: -webkit-linear-gradient(left, red, orange); |
|
background: linear-gradient(to right, red, orange); |
|
-webkit-background-clip: text; |
|
-webkit-text-fill-color: transparent; |
|
display: inline; |
|
font-size: 3em; |
|
} |
|
</style> |
|
<div class="gradient-text">InternVL2</div> |
|
""" |
|
if lan == 'English': |
|
st.markdown(gradient_text_html, unsafe_allow_html=True) |
|
st.caption('Expanding Performance Boundaries of Open-Source Multimodal Large Language Models') |
|
else: |
|
st.markdown(gradient_text_html.replace('InternVL2', '书生·万象'), unsafe_allow_html=True) |
|
st.caption('扩展开源多模态大语言模型的性能边界') |
|
|
|
|
|
if 'messages' not in st.session_state.keys(): |
|
clear_chat_history() |
|
|
|
gallery_placeholder = st.empty() |
|
with gallery_placeholder.container(): |
|
images = ['gallery/prod_en_17.png', 'gallery/astro_on_unicorn.png', |
|
'gallery/prod_12.png', 'gallery/prod_9.jpg', |
|
'gallery/prod_4.png', 'gallery/cheetah.png', 'gallery/prod_1.jpeg'] |
|
images = [Image.open(image) for image in images] |
|
if lan == 'English': |
|
captions = ["I'm on a diet, but I really want to eat them.", |
|
'Could you help me draw a picture like this one?', |
|
'What are the consequences of the easy decisions shown in this image?', |
|
"What's at the far end of the image?", |
|
'Is this a real plant? Analyze the reasons.', |
|
'Detect the <ref>the middle leopard</ref> in the image with its bounding box.', |
|
'What is the atmosphere of this image?'] |
|
else: |
|
captions = ['我在减肥，但我真的很想吃这个。', |
|
'请画一张类似这样的画', |
|
'这张图上 easy decisions 导致了什么后果?', |
|
'画面最远处是什么?', |
|
'这是真的植物吗？分析原因', |
|
'在以下图像中进行目标检测，并标出所有物体。', |
|
'这幅图的氛围如何？'] |
|
img_idx = image_select( |
|
label='', |
|
images=images, |
|
captions=captions, |
|
use_container_width=True, |
|
index=-1, |
|
return_value='index', |
|
key='image_select' |
|
) |
|
if lan == 'English': |
|
st.caption('Note: For non-commercial research use only. AI responses may contain errors. Users should not spread or allow others to spread hate speech, violence, pornography, or fraud-related harmful information.') |
|
else: |
|
st.caption('注意：仅限非商业研究使用。用户应不传播或允许他人传播仇恨言论、暴力、色情内容或与欺诈相关的有害信息。') |
|
if img_idx != -1 and len(st.session_state.messages) == 0 and selected_model is not None: |
|
gallery_placeholder.empty() |
|
st.session_state.messages.append({'role': 'user', 'content': captions[img_idx], 'image': [images[img_idx]]}) |
|
st.rerun() |
|
|
|
if len(st.session_state.messages) > 0: |
|
gallery_placeholder.empty() |
|
|
|
|
|
total_image_num = 0 |
|
for message in st.session_state.messages: |
|
with st.chat_message(message['role']): |
|
st.markdown(message['content']) |
|
show_one_or_multiple_images(message, total_image_num, is_input=message['role'] == 'user') |
|
if 'image' in message and message['role'] == 'user': |
|
total_image_num += len(message['image']) |
|
|
|
input_disable_flag = (len(model_list) == 0) or total_image_num + len(uploaded_files) > max_image_limit |
|
if lan == 'English': |
|
st.sidebar.button('Clear Chat History', on_click=partial(combined_func, func_list=[clear_chat_history, clear_file_uploader])) |
|
if input_disable_flag: |
|
prompt = st.chat_input('Too many images have been uploaded. Please clear the history.', disabled=input_disable_flag) |
|
else: |
|
prompt = st.chat_input('Send messages to InternVL', disabled=input_disable_flag) |
|
else: |
|
st.sidebar.button('清空聊天记录', on_click=partial(combined_func, func_list=[clear_chat_history, clear_file_uploader])) |
|
if input_disable_flag: |
|
prompt = st.chat_input('输入的图片太多了，请清空历史记录。', disabled=input_disable_flag) |
|
else: |
|
prompt = st.chat_input('给 “InternVL” 发送消息', disabled=input_disable_flag) |
|
|
|
alias_instructions = { |
|
'目标检测': '在以下图像中进行目标检测，并标出所有物体。', |
|
'检测': '在以下图像中进行目标检测，并标出所有物体。', |
|
'object detection': 'Please identify and label all objects in the following image.', |
|
'detection': 'Please identify and label all objects in the following image.' |
|
} |
|
|
|
if prompt: |
|
prompt = alias_instructions[prompt] if prompt in alias_instructions else prompt |
|
gallery_placeholder.empty() |
|
image_list = uploaded_pil_images |
|
st.session_state.messages.append({'role': 'user', 'content': prompt, 'image': image_list}) |
|
with st.chat_message('user'): |
|
st.write(prompt) |
|
show_one_or_multiple_images(st.session_state.messages[-1], total_image_num, is_input=True) |
|
if image_list: |
|
clear_file_uploader() |
|
|
|
|
|
if len(st.session_state.messages) > 0 and st.session_state.messages[-1]['role'] != 'assistant': |
|
with st.chat_message('assistant'): |
|
with st.spinner('Thinking...'): |
|
if not prompt: |
|
prompt = st.session_state.messages[-1]['content'] |
|
response = generate_response(st.session_state.messages) |
|
message = {'role': 'assistant', 'content': response} |
|
with st.spinner('Drawing...'): |
|
if '<ref>' in response: |
|
has_returned_image = find_bounding_boxes(response) |
|
message['image'] = [has_returned_image] if has_returned_image else [] |
|
if '```drawing-instruction' in response: |
|
has_returned_image = query_image_generation(response, sd_worker_url=sd_worker_url) |
|
message['image'] = [has_returned_image] if has_returned_image else [] |
|
st.session_state.messages.append(message) |
|
show_one_or_multiple_images(message, total_image_num, is_input=False) |
|
|
|
if len(st.session_state.messages) > 0: |
|
col1, col2, col3, col4 = st.columns([1, 1, 1, 1.3]) |
|
text1 = 'Clear Chat History' if lan == 'English' else '清空聊天记录' |
|
text2 = 'Regenerate' if lan == 'English' else '重新生成' |
|
text3 = 'Copy' if lan == 'English' else '复制回答' |
|
with col1: |
|
st.button(text1, on_click=partial(combined_func, func_list=[clear_chat_history, clear_file_uploader]), |
|
key='clear_chat_history_button') |
|
with col2: |
|
st.button(text2, on_click=regenerate, key='regenerate_button') |
|
|
|
print(st.session_state.messages) |
|
|
