File size: 3,167 Bytes
f04732f
d02b0d1
 
bdaba80
d02b0d1
f04732f
 
 
 
 
 
8ce485a
f04732f
50def22
 
 
d02b0d1
 
6c67d55
d02b0d1
6c67d55
 
 
 
 
d02b0d1
6c67d55
 
f04732f
 
 
 
 
 
 
 
 
 
 
 
340a6dd
1117f0e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f04732f
 
50def22
f04732f
cdaf5e3
 
50def22
 
 
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
77
78
79
80
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 ="""
#component-3 { height: 400px; }
"""
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]

    try:
        if image is None:
            # Handle the case where image is None
            gr.Error("You need to upload an image for LLaVA to work.")
        else:
            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.08)
            yield generated_text_without_prompt
    except NameError:
        gr.Error("You need to upload an image for LLaVA to work.")


demo = gr.ChatInterface(fn=bot_streaming, css=CSS, fill_height=True, 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.queue(default_concurrency_limit=20, max_size=20, api_open=False)
demo.launch(show_api=False, share=False)