File size: 910 Bytes
1b80e0f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import argparse
from PIL import Image
from moondream import VisionEncoder, TextModel
from huggingface_hub import snapshot_download

def main(image_path, prompt):
    model_path = snapshot_download("vikhyatk/moondream1")
    vision_encoder = VisionEncoder(model_path)
    text_model = TextModel(model_path)
    image = Image.open(image_path)
    image_embeds = vision_encoder(image)

    if prompt is None:
        while True:
            question = input("> ")
            print(text_model.answer_question(image_embeds, question))
    else:
        print(">", prompt)
        print(text_model.answer_question(image_embeds, prompt))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--image", type=str, required=True)
    parser.add_argument("--prompt", type=str, required=False)
    args = parser.parse_args()
    main(args.image, args.prompt)