Update app.py

app.py

@@ -25,8 +25,6 @@ except ImportError:
 
     IN_SPACES = False
 
-print("IN_SPACES", IN_SPACES)
-
 import torch
 import os
 import gradio as gr
@@ -52,8 +50,6 @@ if IN_SPACES:
     )
 
 auth_token = os.environ.get("TOKEN_FROM_SECRET") or True
-os.environ["HF_TOKEN"] = auth_token
-
 tokenizer = AutoTokenizer.from_pretrained("vikhyatk/moondream-next")
 moondream = AutoModelForCausalLM.from_pretrained(
     "vikhyatk/moondream-next",
@@ -61,6 +57,7 @@ moondream = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.float16,
     device_map={"": "cuda"},
     attn_implementation="flash_attention_2",
+    token=auth_token if IN_SPACES else None,
 )
 
 # CKPT_DIRS = ["/tmp/md-ckpt/ckpt/ft/song-moon-4c-s15/s72001/"]
@@ -217,6 +214,32 @@ def detect(img, object):
     )
 
 
+@spaces.GPU(duration=10)
+def point(img, object):
+    if img is None:
+        yield "", gr.update(visible=False, value=None)
+        return
+
+    w, h = img.size
+    if w > 768 or h > 768:
+        img = Resize(768)(img)
+        w, h = img.size
+
+    objs = moondream.point(img, object, tokenizer)
+    draw_image = ImageDraw.Draw(img)
+    for o in objs:
+        draw_image.ellipse(
+            (o["x"] * w - 5, o["y"] * h - 5, o["x"] * w + 5, o["y"] * h + 5),
+            fill="red",
+            outline="blue",
+            width=2,
+        )
+
+    yield {"text": f"{len(objs)} detected", "entities": []}, gr.update(
+        visible=True, value=img
+    )
+
+
 js = """
     function createBgAnimation() {
         var canvas = document.createElement('canvas');
@@ -375,7 +398,7 @@ with gr.Blocks(title="moondream vl (new)", css=css, js=js) as demo:
         """
     )
     mode_radio = gr.Radio(
-        ["Caption", "Query", "Detect"],
+        ["Caption", "Query", "Detect", "Point"],
         show_label=False,
         value=lambda: "Caption",
     )
@@ -427,6 +450,19 @@ with gr.Blocks(title="moondream vl (new)", css=css, js=js) as demo:
                     submit.click(detect, [img, prompt], [thought, ann])
                     prompt.submit(detect, [img, prompt], [thought, ann])
                     img.change(detect, [img, prompt], [thought, ann])
+                elif mode == "Point":
+                    with gr.Group():
+                        with gr.Row():
+                            prompt = gr.Textbox(
+                                label="Object",
+                                value="Cat",
+                                scale=4,
+                            )
+                            submit = gr.Button("Submit")
+                        img = gr.Image(type="pil", label="Upload an Image")
+                    submit.click(point, [img, prompt], [thought, ann])
+                    prompt.submit(point, [img, prompt], [thought, ann])
+                    img.change(point, [img, prompt], [thought, ann])
                 else:
                     gr.Markdown("Coming soon!")
 
@@ -473,4 +509,4 @@ with gr.Blocks(title="moondream vl (new)", css=css, js=js) as demo:
         [output, thought, ann],
     )
 
-demo.queue().launch()
+demo.queue().launch(share=True)