Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,819 Bytes
f04732f d02b0d1 bdaba80 d02b0d1 f04732f 8ce485a f04732f cf67b8a 5ef8d1f 4d81d07 cf67b8a d02b0d1 a901f92 d02b0d1 a901f92 d02b0d1 a901f92 f04732f 340a6dd 47bb660 f04732f 47bb660 f04732f 47bb660 f04732f 47bb660 f04732f 47bb660 f04732f 47bb660 f04732f 47bb660 f04732f be7bc1f f04732f be7bc1f f04732f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
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
CSS ="""
.container { display: flex; flex-direction: column; height: 500vh; }
#chatbot { flex-grow: 1; }
"""
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")
# model.generation_config.eos_token_id = 128009
@spaces.GPU
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=1024)
# 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
with gr.Blocks(css=CSS) as demo:
chatbot = 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) |