30 lines
880 B
Python
30 lines
880 B
Python
from moondream import VisionEncoder, TextModel
|
|
from PIL import Image
|
|
from huggingface_hub import snapshot_download
|
|
import argparse
|
|
|
|
model_path = snapshot_download("vikhyatk/moondream0")
|
|
|
|
vision_encoder = VisionEncoder(model_path)
|
|
text_model = TextModel(model_path)
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--image", type=str, required=True)
|
|
parser.add_argument("--interactive", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
image = Image.open(args.image)
|
|
image_embeds = vision_encoder(image)
|
|
|
|
if args.interactive:
|
|
while True:
|
|
question = input("> ")
|
|
print(text_model.answer_question(image_embeds, question))
|
|
print()
|
|
else:
|
|
suggestions = text_model.suggest_questions(image_embeds)
|
|
for suggestion in suggestions:
|
|
print("> ", suggestion)
|
|
print(text_model.answer_question(image_embeds, suggestion))
|
|
print()
|