Csplk commited on
Commit
6a8ca1f
β€’
1 Parent(s): 93b2eb7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -0
app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import spaces
2
+ import torch
3
+ import re
4
+ import gradio as gr
5
+ from transformers import AutoTokenizer, AutoModelForCausalLM
6
+
7
+ if torch.cuda.is_available():
8
+ device, dtype = "cuda", torch.float16
9
+ else:
10
+ device, dtype = "cpu", torch.float32
11
+
12
+ model_id = "vikhyatk/moondream2"
13
+ revision = "2024-04-02"
14
+ tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
15
+ moondream = AutoModelForCausalLM.from_pretrained(
16
+ model_id, trust_remote_code=True, revision=revision
17
+ ).to(device=device, dtype=dtype)
18
+ moondream.eval()
19
+
20
+ @spaces.GPU(duration=10)
21
+ def answer_questions(images, prompt_text):
22
+ prompts = [p.strip() for p in prompt_text.split(',')] # Splitting and cleaning prompts
23
+ image_embeds = [moondream.encode_image(img) for img in images]
24
+ answers = moondream.batch_answer(
25
+ images=image_embeds,
26
+ prompts=prompts,
27
+ tokenizer=tokenizer,
28
+ )
29
+ return ["\n".join(ans) for ans in answers]
30
+
31
+ with gr.Blocks() as demo:
32
+ gr.Markdown("# πŸŒ” moondream2\nA tiny vision language model. [GitHub](https://github.com/vikhyatk/moondream)")
33
+ with gr.Row():
34
+ img = gr.Gallery(label="Upload Images", type="pil")
35
+ prompt = gr.Textbox(label="Input Prompts", placeholder="Enter prompts separated by commas. Ex: Describe this image, What is in this image?", lines=2)
36
+ submit = gr.Button("Submit")
37
+ output = gr.TextArea(label="Responses", lines=4)
38
+ submit.click(answer_questions, [img, prompt], output)
39
+
40
+ demo.queue().launch()