Change model

app.py

@@ -1,8 +1,7 @@
+import einops
 import streamlit as st
 from PIL import Image
-import torch
-
-from transformers import AutoProcessor, AutoModelForVision2Seq
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 DEVICE = "cuda:0"
 
@@ -15,22 +14,12 @@ if uploaded_file is not None:
     image = Image.open(uploaded_file)
     st.image(image, caption='Uploaded Image.', use_column_width=True)
 
-    processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics2-8b-base")
-    model = AutoModelForVision2Seq.from_pretrained(
-        "HuggingFaceM4/idefics2-8b-base",
-    ).to(DEVICE)
-
-    # Create inputs
-    prompts = [
-        "Describe the image in a few sentences.",
-    ]
-    images = [image]
-    inputs = processor(text=prompts, images=images, padding=True, return_tensors="pt")
-    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-
-
-    # Generate
-    generated_ids = model.generate(**inputs, max_new_tokens=500)
-    generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True)
+    model_id = "vikhyatk/moondream2"
+    revision = "2024-05-08"
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id, trust_remote_code=True, revision=revision
+    )
+    tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
 
-    print(generated_texts)
+    enc_image = model.encode_image(image)
+    print(model.answer_question(enc_image, "Describe this image in a few sentences.", tokenizer))

requirements.txt

@@ -1,4 +1,5 @@
+einops
 streamlit
 Pillow
-torch
-transformers
+transformers
+torch