Upload 2 files

app.py

@@ -6,7 +6,7 @@ from threading import Thread
 from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM
 from PIL import ImageDraw
 from torchvision.transforms.v2 import Resize
-
+from optimum.onnxruntime import ORTModelForImageClassification
 import subprocess
 #subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
 
@@ -17,14 +17,20 @@ model_id = "vikhyatk/moondream2"
 #revision = "2024-08-26"
 #tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-moondream = AutoModelForCausalLM.from_pretrained(
-    model_id, trust_remote_code=True, #revision=revision,
-    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, #device_map="auto",
+#moondream = AutoModelForCausalLM.from_pretrained(
+#    model_id, trust_remote_code=True, #revision=revision,
+#    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, #device_map="auto",
     #ignore_mismatched_sizes=True,
     #attn_implementation="flash_attention_2"
+#).to(device)
+moondream = ORTModelForImageClassification.from_pretrained(
+    model_id, trust_remote_code=True,
+    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+    from_transformers=True
 ).to(device)
+
 moondream.eval()
-moondream.to_bettertransformer()
+#moondream.to_bettertransformer()
 
 #@spaces.GPU
 def answer_question(img, prompt):

requirements.txt

@@ -5,5 +5,5 @@ accelerate>=0.32.1
 numpy<2
 torch==2.4.0
 torchvision
-optimum
+optimum[onnxruntime]
 intel_extension_for_pytorch