File size: 1,717 Bytes
83c71a6
 
5756dad
83c71a6
 
5756dad
6dafc63
77b3326
 
83c71a6
 
bf22d27
6dafc63
83c71a6
77b3326
2f92f19
 
83c71a6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2f92f19
 
5756dad
60a2be3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83c71a6
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
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
from threading import Thread

import gradio as gr
from transformers import PreTrainedModel
from transformers import TextIteratorStreamer, AutoModelForCausalLM, AutoTokenizer

# Moondream does not support the HuggingFace pipeline system, so we have to do it manually
moondream_id = "vikhyatk/moondream2"
moondream_revision = "2024-04-02"
moondream_tokenizer = AutoTokenizer.from_pretrained(moondream_id, revision=moondream_revision)
moondream: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
    moondream_id, trust_remote_code=True, revision=moondream_revision
)
moondream.eval()


def answer_question(_img, _prompt):
    image_embeds = moondream.encode_image(_img)
    streamer = TextIteratorStreamer(moondream_tokenizer, skip_special_tokens=True)
    thread = Thread(
        target=moondream.answer_question,
        kwargs={
            "image_embeds": image_embeds,
            "question": _prompt,
            "tokenizer": moondream_tokenizer,
            "streamer": streamer,
        },
    )
    thread.start()

    buffer = ""
    for new_text in streamer:
        buffer += new_text
        yield buffer.strip()


if __name__ == "__main__":
    with gr.Blocks() as app:
        gr.Markdown(
            """
            # Food Identifier

            Final project for IAT 481 at Simon Fraser University, Spring 2024.
            """
        )
        with gr.Row():
            prompt = gr.Textbox(label="Input", value="Describe this image.")
            submit = gr.Button("Submit")
        with gr.Row():
            img = gr.Image(label="Image", type="pil")
            output = gr.TextArea(label="Output")

        submit.click(answer_question, [img, prompt], output)

    app.queue().launch()