40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import gradio as gr
|
|
from transformers import AutoTokenizer, AutoModel
|
|
from gradio.themes.utils import sizes
|
|
|
|
|
|
theme = gr.themes.Default(radius_size=sizes.radius_none).set(
|
|
block_label_text_color = '#4D63FF',
|
|
block_title_text_color = '#4D63FF',
|
|
button_primary_text_color = '#4D63FF',
|
|
button_primary_background_fill='#FFFFFF',
|
|
button_primary_border_color='#4D63FF',
|
|
button_primary_background_fill_hover='#EDEFFF',
|
|
)
|
|
|
|
css = "footer {visibility: hidden}"
|
|
|
|
|
|
def respond(image_path, message, chat_history):
|
|
tokenizer = AutoTokenizer.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True)
|
|
model = AutoModel.from_pretrained("THUDM/visualglm-6b", trust_remote_code=True).half().cuda()
|
|
|
|
response, history = model.chat(tokenizer, image_path, message, history=chat_history)
|
|
|
|
return "", history
|
|
|
|
|
|
with gr.Blocks(theme=theme, css=css) as demo:
|
|
with gr.Row():
|
|
image = gr.Image(type="filepath")
|
|
with gr.Column():
|
|
chatbot = gr.Chatbot()
|
|
message = gr.Textbox()
|
|
clear = gr.Button("Clear")
|
|
|
|
message.submit(respond, [image, message, chatbot], [message, chatbot])
|
|
clear.click(lambda: None, None, chatbot, queue=False)
|
|
|
|
if __name__ == "__main__":
|
|
demo.launch(server_name="0.0.0.0")
|