Update app.py

app.py

@@ -1,81 +1,74 @@
-import torch
 import gradio as gr
-from transformers import AutoModelForVision2Seq, AutoProcessor
+import torch
+from PIL import Image
+from transformers import AutoTokenizer
 
-# ===============================
+# =========================
 # Model config
-# ===============================
+# =========================
 MODEL_ID = "vikhyatk/moondream2"
+DEVICE = "cpu"  
 
-device = "cpu"
-
-# Processor & Model
-processor = AutoProcessor.from_pretrained(
+# =========================
+# Load model (IMPORTANT)
+# =========================
+tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     trust_remote_code=True
 )
 
-model = AutoModelForVision2Seq.from_pretrained(
-    MODEL_ID,
-    trust_remote_code=True,
-    torch_dtype=torch.float32,  # CPU安全
-).to(device)
+model = torch.load(
+    torch.hub.load_state_dict_from_url(
+        f"https://huggingface.co/{MODEL_ID}/resolve/main/moondream.pt",
+        map_location=DEVICE
+    )
+)
 
 model.eval()
 
-# ===============================
+# =========================
 # Inference function
-# ===============================
-def infer(image, prompt):
+# =========================
+def understand_image(image, prompt):
     if image is None:
         return "Please upload an image."
 
-    if prompt is None or prompt.strip() == "":
-        prompt = "Describe this image."
+    image = image.convert("RGB")
 
     with torch.no_grad():
-        answer = model.answer(
-            image=image,
-            question=prompt
+        answer = model.answer_question(
+            image,
+            prompt,
+            tokenizer
         )
 
     return answer
 
 
-# ===============================
+# =========================
 # Gradio UI
-# ===============================
+# =========================
 with gr.Blocks() as demo:
-    gr.Markdown("# 🖼️ Vision Language Demo (moondream2 · CPU)")
+    gr.Markdown("# 🌓 Moondream2 Image Understanding (Free Tier)")
     gr.Markdown(
         "⚠️ Uploaded images are processed in memory and not stored permanently."
     )
 
     with gr.Row():
-        image = gr.Image(
-            type="pil",
-            label="Upload Image"
+        image_input = gr.Image(type="pil", label="Upload Image")
+        text_input = gr.Textbox(
+            label="Question",
+            placeholder="What is in this image?"
         )
 
-        with gr.Column():
-            textbox = gr.Textbox(
-                label="Prompt",
-                value="Describe this image."
-            )
-            btn = gr.Button("Run")
-
-    output = gr.Textbox(
-        label="Output",
-        lines=6
-    )
+    output = gr.Textbox(label="Answer")
 
+    btn = gr.Button("Run")
     btn.click(
-        fn=infer,
-        inputs=[image, textbox],
+        understand_image,
+        inputs=[image_input, text_input],
         outputs=output
     )
 
 demo.launch()
 
-
-