Csplk/moondream2-batch-processing

1import spaces2import torch3import re4import gradio as gr5from threading import Thread6from transformers import AutoTokenizer, AutoModelForCausalLM7#from PIL import ImageDraw8#from torchvision.transforms.v2 import Resize9#import subprocess10 11#subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)12 13#subprocess.run('cp -r moondream/torch clients/python/moondream/torch')14#subprocess.run('pip install moondream[gpu]')15 16#def load_moondream():17#    """Load Moondream model and tokenizer."""18#    model = AutoModelForCausalLM.from_pretrained(19#        "vikhyatk/moondream2", trust_remote_code=True, device_map={"": "cuda"}20#    )21#    tokenizer = AutoTokenizer.from_pretrained("vikhyatk/moondream2")22#    return model, tokenizer23 24#"""Load Moondream model and tokenizer."""25#moondream = AutoModelForCausalLM.from_pretrained(26#    "vikhyatk/moondream2", trust_remote_code=True, device_map={"": "cuda"}27#)28#tokenizer = AutoTokenizer.from_pretrained("vikhyatk/moondream2")29 30#model_id = "vikhyatk/moondream2"31#revision = "2025-01-09"32#tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)33#moondream = AutoModelForCausalLM.from_pretrained(34#    model_id, trust_remote_code=True, revision=revision,35#    torch_dtype=torch.bfloat16, device_map={"": "cuda"},36#)37 38#moondream.eval()39 40moondream = AutoModelForCausalLM.from_pretrained(41    "vikhyatk/moondream2",42    revision="2025-06-21",43    trust_remote_code=True,44    device_map={"": "cuda"}  # ...or 'mps', on Apple Silicon45)46 47@spaces.GPU(durtion="150")48def answer_questions(image_tuples, prompt_text):49    result = ""50    Q_and_A = ""51    prompts = [p.strip() for p in prompt_text.split('?')]52    image_embeds = [img[0] for img in image_tuples if img[0] is not None]53    answers = []54 55    for prompt in prompts:56        answers.append(moondream.batch_answer(57            images=[img.convert("RGB") for img in image_embeds],58            prompts=[prompt] * len(image_embeds),59            tokenizer=tokenizer60        ))61 62    for i, prompt in enumerate(prompts):63        Q_and_A += f"### Q: {prompt}\n"64        for j, image_tuple in enumerate(image_tuples):65            image_name = f"image{j+1}"66            answer_text = answers[i][j]67            Q_and_A += f"**{image_name} A:** \n {answer_text} \n"68 69    result = {'headers': prompts, 'data': answers}70    #print("result\n{}\n\nQ_and_A\n{}\n\n".format(result, Q_and_A))71    return Q_and_A, result72 73with gr.Blocks() as demo:74    gr.Markdown("# moondream2 unofficial batch processing demo")75    gr.Markdown("1. Select images\n2. Enter one or more prompts separated by commas. Ex: Describe this image, What is in this image?\n\n")76    gr.Markdown("**Currently each image will be sent as a batch with the prompts thus asking each prompt on each image**")77    gr.Markdown("*Running on free CPU space tier currently so results may take a bit to process compared to duplicating space and using GPU space hardware*")78    gr.Markdown("A tiny vision language model. [moondream2](https://huggingface.co/vikhyatk/moondream2)")79    with gr.Row():80        img = gr.Gallery(label="Upload Images", type="pil", preview=True, columns=4)81    with gr.Row():82        prompt = gr.Textbox(label="Input Prompts", placeholder="Enter prompts (one prompt for each image provided) separated by question marks. Ex: Describe this image? What is in this image?", lines=8)83    with gr.Row():84        submit = gr.Button("Submit")85    with gr.Row():86        output = gr.Markdown(label="Questions and Answers", line_breaks=True)87    with gr.Row():88        output2 = gr.Dataframe(label="Structured Dataframe", type="array", wrap=True)89    submit.click(answer_questions, inputs=[img, prompt], outputs=[output, output2])90 91demo.queue().launch()92