Spaces:
Running
on
Zero
Running
on
Zero
import gradio as gr | |
from transformers import AutoProcessor, LlavaForConditionalGeneration | |
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer, TextIteratorStreamer | |
from threading import Thread | |
import re | |
import time | |
from PIL import Image | |
import torch | |
import spaces | |
import requests | |
model_id = "xtuner/llava-llama-3-8b-v1_1-transformers" | |
processor = AutoProcessor.from_pretrained(model_id) | |
model = LlavaForConditionalGeneration.from_pretrained( | |
model_id, | |
torch_dtype=torch.float16, | |
low_cpu_mem_usage=True, | |
) | |
model.to("cuda:0") | |
def bot_streaming(message, history): | |
print(message) | |
if message["files"]: | |
image = message["files"][-1]["path"] | |
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] | |
if image is None: | |
gr.Error("You need to upload an image for LLaVA to work.") | |
prompt=f"<|start_header_id|>user<|end_header_id|>\n\n<image>\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" | |
print(f"prompt: {prompt}") | |
image = Image.open(image) | |
inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16) | |
streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": True}) | |
generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=200) | |
generated_text = "" | |
thread = Thread(target=model.generate, kwargs=generation_kwargs) | |
thread.start() | |
text_prompt =f"<|start_header_id|>user<|end_header_id|>\n\n{message['text']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n" | |
print(f"text_prompt: {text_prompt}") | |
buffer = "" | |
for new_text in streamer: | |
buffer += new_text | |
generated_text_without_prompt = buffer[len(text_prompt):] | |
time.sleep(0.04) | |
yield generated_text_without_prompt | |
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) | |
demo.launch(debug=True) |