Spaces:
Paused
Paused
Update app.py
Browse filesupdate to match impl from vikhyatk that uses flash_attention_2. also using ZeroGPU.
app.py
CHANGED
@@ -4,23 +4,24 @@ import re
|
|
4 |
import gradio as gr
|
5 |
from threading import Thread
|
6 |
from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
|
|
|
|
|
7 |
|
8 |
import subprocess
|
9 |
-
|
10 |
-
|
11 |
-
device = torch.device("cuda")
|
12 |
-
dtype = torch.float32
|
13 |
|
14 |
model_id = "vikhyatk/moondream2"
|
15 |
-
revision = "2024-05-
|
16 |
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
|
17 |
moondream = AutoModelForCausalLM.from_pretrained(
|
18 |
-
model_id, trust_remote_code=True, revision=revision
|
19 |
-
|
|
|
|
|
20 |
moondream.eval()
|
21 |
|
22 |
|
23 |
-
@spaces.GPU(duration=
|
24 |
def answer_question(img, prompt):
|
25 |
image_embeds = moondream.encode_image(img)
|
26 |
streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
|
@@ -40,6 +41,35 @@ def answer_question(img, prompt):
|
|
40 |
buffer += new_text
|
41 |
yield buffer.strip()
|
42 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
43 |
|
44 |
with gr.Blocks() as demo:
|
45 |
gr.Markdown(
|
@@ -53,8 +83,12 @@ with gr.Blocks() as demo:
|
|
53 |
submit = gr.Button("Submit")
|
54 |
with gr.Row():
|
55 |
img = gr.Image(type="pil", label="Upload an Image")
|
56 |
-
|
|
|
|
|
|
|
57 |
submit.click(answer_question, [img, prompt], output)
|
58 |
prompt.submit(answer_question, [img, prompt], output)
|
|
|
59 |
|
60 |
-
demo.queue().launch()
|
|
|
4 |
import gradio as gr
|
5 |
from threading import Thread
|
6 |
from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
|
7 |
+
from PIL import ImageDraw
|
8 |
+
from torchvision.transforms.v2 import Resize
|
9 |
|
10 |
import subprocess
|
11 |
+
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
|
|
|
|
|
|
|
12 |
|
13 |
model_id = "vikhyatk/moondream2"
|
14 |
+
revision = "2024-05-20"
|
15 |
tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
|
16 |
moondream = AutoModelForCausalLM.from_pretrained(
|
17 |
+
model_id, trust_remote_code=True, revision=revision,
|
18 |
+
torch_dtype=torch.bfloat16, device_map={"": "cuda"},
|
19 |
+
attn_implementation="flash_attention_2"
|
20 |
+
)
|
21 |
moondream.eval()
|
22 |
|
23 |
|
24 |
+
@spaces.GPU(duration=10)
|
25 |
def answer_question(img, prompt):
|
26 |
image_embeds = moondream.encode_image(img)
|
27 |
streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
|
|
|
41 |
buffer += new_text
|
42 |
yield buffer.strip()
|
43 |
|
44 |
+
def extract_floats(text):
|
45 |
+
# Regular expression to match an array of four floating point numbers
|
46 |
+
pattern = r"\[\s*(-?\d+\.\d+)\s*,\s*(-?\d+\.\d+)\s*,\s*(-?\d+\.\d+)\s*,\s*(-?\d+\.\d+)\s*\]"
|
47 |
+
match = re.search(pattern, text)
|
48 |
+
if match:
|
49 |
+
# Extract the numbers and convert them to floats
|
50 |
+
return [float(num) for num in match.groups()]
|
51 |
+
return None # Return None if no match is found
|
52 |
+
|
53 |
+
|
54 |
+
def extract_bbox(text):
|
55 |
+
bbox = None
|
56 |
+
if extract_floats(text) is not None:
|
57 |
+
x1, y1, x2, y2 = extract_floats(text)
|
58 |
+
bbox = (x1, y1, x2, y2)
|
59 |
+
return bbox
|
60 |
+
|
61 |
+
def process_answer(img, answer):
|
62 |
+
if extract_bbox(answer) is not None:
|
63 |
+
x1, y1, x2, y2 = extract_bbox(answer)
|
64 |
+
draw_image = Resize(768)(img)
|
65 |
+
width, height = draw_image.size
|
66 |
+
x1, x2 = int(x1 * width), int(x2 * width)
|
67 |
+
y1, y2 = int(y1 * height), int(y2 * height)
|
68 |
+
bbox = (x1, y1, x2, y2)
|
69 |
+
ImageDraw.Draw(draw_image).rectangle(bbox, outline="red", width=3)
|
70 |
+
return gr.update(visible=True, value=draw_image)
|
71 |
+
|
72 |
+
return gr.update(visible=False, value=None)
|
73 |
|
74 |
with gr.Blocks() as demo:
|
75 |
gr.Markdown(
|
|
|
83 |
submit = gr.Button("Submit")
|
84 |
with gr.Row():
|
85 |
img = gr.Image(type="pil", label="Upload an Image")
|
86 |
+
with gr.Column():
|
87 |
+
output = gr.Markdown(label="Response")
|
88 |
+
ann = gr.Image(visible=False, label="Annotated Image")
|
89 |
+
|
90 |
submit.click(answer_question, [img, prompt], output)
|
91 |
prompt.submit(answer_question, [img, prompt], output)
|
92 |
+
output.change(process_answer, [img, output], ann, show_progress=False)
|
93 |
|
94 |
+
demo.queue().launch()
|