Merge pull request #8 from yvrjsharma/add-webui-demo

Adding and updating the files for a gradio demo

README.md

@@ -45,6 +45,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.

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)

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
+einops==0.7.0
+gradio==4.15.0