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 =""" .contain { display: flex; flex-direction: column; } #component-0 { height: 100%; } #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") @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\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 demo = gr.ChatInterface(fn=bot_streaming, css=CSS, 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)