From dd6fb9291d0fedb48eaee9bb3c21850a4a914e5c Mon Sep 17 00:00:00 2001
From: Yuvraj Sharma <48665385+yvrjsharma@users.noreply.github.com>
Date: Mon, 22 Jan 2024 19:13:53 +0530
Subject: [PATCH] Adding and updating the files for a gradio demo

---
 README.md        |  9 ++++++++
 gradio_demo.py   | 54 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  3 ++-
 3 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 gradio_demo.py

diff --git a/README.md b/README.md
index 373fa96..464e9c0 100644
--- a/README.md
+++ b/README.md
@@ -44,6 +44,15 @@ python sample.py --image [IMAGE_PATH] --prompt [PROMPT]
 When the `--prompt` argument is not provided, the script will allow you to ask
 questions interactively.
 
+**Gradio demo**
+
+Use the `grado_demo.py` script to run the gradio app:
+
+```python
+python gradio_app.py
+```
+
+
 **Limitations**
 
 * The model may generate inaccurate statements.
diff --git a/gradio_demo.py b/gradio_demo.py
new file mode 100644
index 0000000..954d884
--- /dev/null
+++ b/gradio_demo.py
@@ -0,0 +1,54 @@
+import re
+import gradio as gr
+from moondream import VisionEncoder, TextModel
+from PIL import Image
+from huggingface_hub import snapshot_download
+from threading import Thread
+from transformers import TextIteratorStreamer
+
+model_path = snapshot_download("vikhyatk/moondream1")
+
+vision_encoder = VisionEncoder(model_path)
+text_model = TextModel(model_path)
+
+# model inference
+def moondream(img, prompt):
+    
+    image_embeds = vision_encoder(img)
+
+    streamer = TextIteratorStreamer(text_model.tokenizer, skip_special_tokens=True)
+    generation_kwargs = dict(
+        image_embeds=image_embeds, question=prompt, streamer=streamer
+    )
+    thread = Thread(target=text_model.answer_question, kwargs=generation_kwargs)
+    thread.start()
+
+    buffer = ""
+    for new_text in streamer:
+        # check for the end of generated text and yield the generated token
+        if not new_text.endswith("<") and not new_text.endswith("END"):
+          buffer += new_text
+          yield buffer
+        else:
+          new_text = re.sub("<$", "", re.sub("END$", "", new_text))
+          buffer += new_text
+          yield buffer
+
+# Using Gradio Blocks API
+with gr.Blocks() as demo:
+  gr.HTML("<h1><center>🌔MoonDream</center></h1>")
+  gr.HTML("<h3><center>A tiny vision language model. <a href='https://github.com/vikhyat/moondream' target='blank_'>GitHub</a></center></h3>")
+  with gr.Group():
+    with gr.Row():
+      prompt = gr.Textbox(label='Input Prompt for the model',placeholder='Type whatever you want to ask about the image',scale=4 )
+      submit = gr.Button('Submit', scale=1,)
+    with gr.Row():
+      img = gr.Image(type='pil', label='Upload or Drag an Image')
+      output = gr.TextArea(label="Bot's response to the user query-", info='The response might take a few seconds..' )
+  
+  # handling events
+  submit.click(moondream, [img, prompt], output)
+  prompt.submit(moondream, [img, prompt], output)
+
+# launch gradio demo with debug mode on
+demo.queue().launch(debug=True)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 408d944..ef66319 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,4 +4,5 @@ Pillow==10.1.0
 torch==2.1.2
 torchvision==0.16.2
 transformers==4.36.2
-einops==0.7.0
\ No newline at end of file
+einops==0.7.0
+gradio==4.15.0
\ No newline at end of file