add transformers moondream

Dockerfile

@@ -13,6 +13,5 @@ COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
 COPY --chown=user . /app
-RUN chmod +x /app/moondream
 
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

app.py

@@ -3,7 +3,9 @@ import subprocess
 from tempfile import NamedTemporaryFile
 
 from fastapi import FastAPI, HTTPException
+from PIL import Image
 from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 app = FastAPI()
 
@@ -14,6 +16,20 @@ class RequestData(BaseModel):
     image: str
 
 
+def load_model():
+    model_id = "vikhyatk/moondream2"
+    revision = "2024-08-26"
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, trust_remote_code=True, revision=revision
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
+    return model, tokenizer
+
+
+MODEL, TOKENIZER = load_model()
+print("INFO: Model loaded successfully!")
+
+
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
@@ -32,15 +48,10 @@ def query(data: RequestData):
             temp_image.write(image)
             temp_image.flush()
 
-            result = subprocess.run(
-                ["./moondream", "--prompt", prompt, "--image", temp_image.name],
-                capture_output=True,
-                text=True,
-            )
-
-            if result.returncode != 0:
-                raise Exception(result.stderr)
+            image = Image.open(temp_image.name)
+            enc_image = MODEL.encode_image(image)
+            response = MODEL.answer_question(enc_image, str(prompt), TOKENIZER)
 
-            return {"response": result.stdout}
+            return {"response": str(response)}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

requirements.txt

@@ -1,3 +1,5 @@
 fastapi
 uvicorn[standard]
 pydantic
+transformers
+torch