48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
from moondream import VisionEncoder, TextModel
|
|
from PIL import Image
|
|
from huggingface_hub import snapshot_download
|
|
import argparse
|
|
from threading import Thread
|
|
from transformers import TextIteratorStreamer
|
|
import re
|
|
|
|
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 question in suggestions:
|
|
print(">", question)
|
|
|
|
streamer = TextIteratorStreamer(text_model.tokenizer, skip_special_tokens=True)
|
|
generation_kwargs = dict(
|
|
image_embeds=image_embeds, question=question, streamer=streamer
|
|
)
|
|
thread = Thread(target=text_model.answer_question, kwargs=generation_kwargs)
|
|
thread.start()
|
|
|
|
buffer = ""
|
|
for new_text in streamer:
|
|
buffer += new_text
|
|
if not new_text.endswith("Human"):
|
|
print(buffer, end="", flush=True)
|
|
buffer = ""
|
|
print(re.sub("Human$", "", buffer))
|
|
|
|
print()
|