nroggendorff commited on
Commit
aecd1c4
·
verified ·
1 Parent(s): 644c9ce

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +98 -0
app.py ADDED
@@ -0,0 +1,98 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import time
2
+ from threading import Thread
3
+
4
+ import gradio as gr
5
+ import torch
6
+ from PIL import Image
7
+ from transformers import AutoProcessor, LlavaForConditionalGeneration
8
+ from transformers import TextIteratorStreamer
9
+
10
+ import spaces
11
+
12
+
13
+ PLACEHOLDER = """
14
+ <div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
15
+ <img src="https://cdn-uploads.huggingface.co/production/uploads/64ccdc322e592905f922a06e/DDIW0kbWmdOQWwy4XMhwX.png" style="width: 80%; max-width: 550px; height: auto; opacity: 0.55; ">
16
+ <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">LLaVA-Llama-3-8B</h1>
17
+ <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Llava-Llama-3-8b is a LLaVA model fine-tuned from Meta-Llama-3-8B-Instruct and CLIP-ViT-Large-patch14-336 with ShareGPT4V-PT and InternVL-SFT by XTuner</p>
18
+ </div>
19
+ """
20
+
21
+
22
+ model_id = "xtuner/llava-llama-3-8b-v1_1-transformers"
23
+
24
+ processor = AutoProcessor.from_pretrained(model_id)
25
+
26
+ model = LlavaForConditionalGeneration.from_pretrained(
27
+ model_id,
28
+ torch_dtype=torch.float16,
29
+ low_cpu_mem_usage=True,
30
+ )
31
+
32
+ model.to("cuda:0")
33
+ model.generation_config.eos_token_id = 128009
34
+
35
+
36
+ @spaces.GPU
37
+ def bot_streaming(message, history):
38
+ print(message)
39
+ if message["files"]:
40
+ # message["files"][-1] is a Dict or just a string
41
+ if type(message["files"][-1]) == dict:
42
+ image = message["files"][-1]["path"]
43
+ else:
44
+ image = message["files"][-1]
45
+ else:
46
+ # if there's no image uploaded for this turn, look for images in the past turns
47
+ # kept inside tuples, take the last one
48
+ for hist in history:
49
+ if type(hist[0]) == tuple:
50
+ image = hist[0][0]
51
+ if image is None:
52
+ # Handle the case where image is None
53
+ image = "http://files.tatl.tv/ignore.png"
54
+
55
+ 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"
56
+ # print(f"prompt: {prompt}")
57
+ image = Image.open(image)
58
+ inputs = processor(prompt, image, return_tensors='pt').to(0, torch.float16)
59
+
60
+ streamer = TextIteratorStreamer(processor, **{"skip_special_tokens": False, "skip_prompt": True})
61
+ generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024, do_sample=False)
62
+
63
+ thread = Thread(target=model.generate, kwargs=generation_kwargs)
64
+ thread.start()
65
+
66
+ 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"
67
+ # print(f"text_prompt: {text_prompt}")
68
+
69
+ buffer = ""
70
+ time.sleep(0.5)
71
+ for new_text in streamer:
72
+ # find <|eot_id|> and remove it from the new_text
73
+ if "<|eot_id|>" in new_text:
74
+ new_text = new_text.split("<|eot_id|>")[0]
75
+ buffer += new_text
76
+
77
+ # generated_text_without_prompt = buffer[len(text_prompt):]
78
+ generated_text_without_prompt = buffer
79
+ # print(generated_text_without_prompt)
80
+ time.sleep(0.06)
81
+ # print(f"new_text: {generated_text_without_prompt}")
82
+ yield generated_text_without_prompt
83
+
84
+
85
+ chatbot=gr.Chatbot(placeholder=PLACEHOLDER,scale=1)
86
+ chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
87
+ with gr.Blocks(fill_height=True, ) as demo:
88
+ gr.ChatInterface(
89
+ fn=bot_streaming,
90
+ description="Try [LLaVA Llama-3-8B](https://huggingface.co/xtuner/llava-llama-3-8b-v1_1-transformers).",
91
+ stop_btn="Stop Generation",
92
+ multimodal=True,
93
+ textbox=chat_input,
94
+ chatbot=chatbot,
95
+ )
96
+
97
+ demo.queue(api_open=False)
98
+ demo.launch(show_api=False, share=False)