Spaces:
Runtime error
Runtime error
import gradio as gr | |
from huggingface_hub import InferenceClient | |
import spaces | |
import os | |
import warnings | |
import shutil | |
import time | |
from threading import Thread | |
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig, AutoProcessor | |
from transformers import TextIteratorStreamer | |
import torch | |
from dc.model import * | |
from dc.constants import DEFAULT_IMAGE_PATCH_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN | |
from dc.conversation import conv_templates, SeparatorStyle | |
from PIL import Image | |
tokenizer = AutoTokenizer.from_pretrained('HuanjinYao/DenseConnector-v1.5-8B', use_fast=False) | |
model = LlavaLlamaForCausalLM.from_pretrained('HuanjinYao/DenseConnector-v1.5-8B', low_cpu_mem_usage=True,torch_dtype=torch.float16) | |
image_processor = model.get_vision_tower() | |
if not vision_tower.is_loaded: | |
vision_tower.load_model() | |
vision_tower.to(device=device, dtype=torch.float16) | |
image_processor = vision_tower.image_processor | |
model.to('cuda') | |
# model.generation_config.eos_token_id = 128009 | |
tokenizer.unk_token = "<|reserved_special_token_0|>" | |
tokenizer.pad_token = tokenizer.unk_token | |
terminators = [ | |
tokenizer.eos_token_id, | |
tokenizer.convert_tokens_to_ids("<|eot_id|>") | |
] | |
def bot_streaming(message, history): | |
print(message) | |
if message["files"]: | |
# message["files"][-1] is a Dict or just a string | |
if type(message["files"][-1]) == dict: | |
image = message["files"][-1]["path"] | |
else: | |
image = message["files"][-1] | |
else: | |
# if there's no image uploaded for this turn, look for images in the past turns | |
# kept inside tuples, take the last one | |
for hist in history: | |
if type(hist[0]) == tuple: | |
image = hist[0][0] | |
try: | |
if image is None: | |
# Handle the case where image is None | |
gr.Error("You need to upload an image for LLaVA to work.") | |
except NameError: | |
# Handle the case where 'image' is not defined at all | |
gr.Error("You need to upload an image for LLaVA to work.") | |
conv = conv_templates['llama_3'].copy() | |
if len(history) == 0: | |
user = DEFAULT_IMAGE_TOKEN + '\n' + message['text'] | |
else: | |
for idx, (user, assistant) in enumerate(history): | |
# conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}]) | |
if idx == 0: | |
user = DEFAULT_IMAGE_TOKEN + '\n' + user | |
conv.append_message(conv.roles[0], user) | |
conv.append_message(conv.roles[1], assistant) | |
conv.append_message(conv.roles[0], user) | |
conv.append_message(conv.roles[1], None) | |
prompt = conv.get_prompt() | |
image = Image.open(os.path.join(image, image_file)).convert('RGB') | |
image_tensor = image_processor([image], image_processor, self.model_config)[0] | |
inputs = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt') | |
streamer = TextIteratorStreamer(tokenizer, **{"skip_special_tokens": False, "skip_prompt": True}) | |
generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=False, eos_token_id = terminators) | |
thread = Thread(target=model.generate, kwargs=generation_kwargs) | |
thread.start() | |
buffer = "" | |
# time.sleep(0.5) | |
for new_text in streamer: | |
if "<|eot_id|>" in new_text: | |
new_text = new_text.split("<|eot_id|>")[0] | |
buffer += new_text | |
generated_text_without_prompt = buffer | |
# time.sleep(0.06) | |
yield generated_text_without_prompt | |
chatbot=gr.Chatbot(placeholder=PLACEHOLDER,scale=1) | |
chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False) | |
with gr.Blocks(fill_height=True, ) as demo: | |
gr.ChatInterface( | |
fn=bot_streaming, | |
title="LLaVA Llama-3-8B", | |
examples=[{"text": "What is on the flower?", "files": ["./bee.jpg"]}, | |
{"text": "How to make this pastry?", "files": ["./baklava.png"]}], | |
description="Try [LLaVA Llama-3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers). Upload an image and start chatting about it, or simply try one of the examples below. If you don't upload an image, you will receive an error.", | |
stop_btn="Stop Generation", | |
multimodal=True, | |
textbox=chat_input, | |
chatbot=chatbot, | |
) | |
if __name__ == "__main__": | |
demo.launch() |