Commit
•
4a5cedb
1
Parent(s):
e2a5026
Update app.py
Browse files
app.py
CHANGED
@@ -1,11 +1,27 @@
|
|
1 |
-
import spaces
|
2 |
import torch
|
3 |
import re
|
|
|
4 |
import gradio as gr
|
5 |
from threading import Thread
|
6 |
from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
|
7 |
-
|
8 |
import subprocess
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
9 |
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
|
10 |
|
11 |
model_id = "vikhyatk/moondream2"
|
@@ -18,11 +34,17 @@ moondream = AutoModelForCausalLM.from_pretrained(
|
|
18 |
)
|
19 |
moondream.eval()
|
20 |
|
21 |
-
|
22 |
-
|
23 |
-
|
|
|
|
|
|
|
|
|
24 |
image_embeds = moondream.encode_image(img)
|
|
|
25 |
streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
|
|
|
26 |
thread = Thread(
|
27 |
target=moondream.answer_question,
|
28 |
kwargs={
|
@@ -37,23 +59,24 @@ def answer_question(img, prompt):
|
|
37 |
buffer = ""
|
38 |
for new_text in streamer:
|
39 |
buffer += new_text
|
40 |
-
|
|
|
41 |
|
|
|
42 |
|
43 |
with gr.Blocks() as demo:
|
44 |
-
gr.
|
45 |
-
""
|
46 |
-
|
47 |
-
|
48 |
-
"""
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
53 |
-
|
54 |
-
|
55 |
-
|
56 |
-
submit.click(answer_question, [
|
57 |
-
prompt.submit(answer_question, [img, prompt], output)
|
58 |
|
59 |
demo.queue().launch()
|
|
|
|
|
1 |
import torch
|
2 |
import re
|
3 |
+
import os
|
4 |
import gradio as gr
|
5 |
from threading import Thread
|
6 |
from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
|
|
|
7 |
import subprocess
|
8 |
+
from PIL import Image
|
9 |
+
from io import BytesIO
|
10 |
+
import base64
|
11 |
+
|
12 |
+
SECRET_TOKEN = os.getenv('SECRET_TOKEN', 'default_secret')
|
13 |
+
|
14 |
+
# Regex pattern to match data URI scheme
|
15 |
+
data_uri_pattern = re.compile(r'data:image/(png|jpeg|jpg|webp);base64,')
|
16 |
+
|
17 |
+
def readb64(b64):
|
18 |
+
# Remove any data URI scheme prefix with regex
|
19 |
+
b64 = data_uri_pattern.sub("", b64)
|
20 |
+
# Decode and open the image with PIL
|
21 |
+
img = Image.open(BytesIO(base64.b64decode(b64)))
|
22 |
+
return img
|
23 |
+
|
24 |
+
|
25 |
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
|
26 |
|
27 |
model_id = "vikhyatk/moondream2"
|
|
|
34 |
)
|
35 |
moondream.eval()
|
36 |
|
37 |
+
def answer_question(secret_token, input, prompt):
|
38 |
+
if secret_token != SECRET_TOKEN:
|
39 |
+
raise gr.Error(
|
40 |
+
f'Invalid secret token. Please fork the original space if you want to use it for yourself.')
|
41 |
+
|
42 |
+
img = readb64(input)
|
43 |
+
|
44 |
image_embeds = moondream.encode_image(img)
|
45 |
+
|
46 |
streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
|
47 |
+
|
48 |
thread = Thread(
|
49 |
target=moondream.answer_question,
|
50 |
kwargs={
|
|
|
59 |
buffer = ""
|
60 |
for new_text in streamer:
|
61 |
buffer += new_text
|
62 |
+
|
63 |
+
buffer.strip()
|
64 |
|
65 |
+
return buffer
|
66 |
|
67 |
with gr.Blocks() as demo:
|
68 |
+
gr.HTML("""
|
69 |
+
<div style="z-index: 100; position: fixed; top: 0px; right: 0px; left: 0px; bottom: 0px; width: 100%; height: 100%; background: white; display: flex; align-items: center; justify-content: center; color: black;">
|
70 |
+
<div style="text-align: center; color: black;">
|
71 |
+
<p style="color: black;">This space is a headless component of the cloud rendering engine used by AiTube.</p>
|
72 |
+
<p style="color: black;">It is not available for public use, but you can use the <a href="https://huggingface.co/spaces/vikhyatk/moondream2" target="_blank">original space</a>.</p>
|
73 |
+
</div>
|
74 |
+
</div>""")
|
75 |
+
token = gr.Textbox()
|
76 |
+
input = gr.Textbox()
|
77 |
+
prompt = gr.Textbox()
|
78 |
+
submit = gr.Button()
|
79 |
+
output = gr.Textbox()
|
80 |
+
submit.click(answer_question, [token, input, prompt], output)
|
|
|
81 |
|
82 |
demo.queue().launch()
|