File size: 1,720 Bytes
ffd7d69 |
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 |
import os
import torch
import spaces
import subprocess
import gradio as gr
from huggingface_hub import login
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
login(os.environ.get("HF_TOKEN"))
model_id = "google/paligemma-3b-mix-448"
model = PaliGemmaForConditionalGeneration.from_pretrained(
model_id, device_map={"": 0},
attn_implementation="flash_attention_2",
torch_dtype=torch.bfloat16,
)
processor = AutoProcessor.from_pretrained(model_id)
model.eval()
@spaces.GPU()
def answer_question(image, prompt):
model_inputs = processor(text=prompt, images=image, return_tensors="pt").to("cuda")
input_len = model_inputs["input_ids"].shape[-1]
with torch.inference_mode():
generation = model.generate(**model_inputs, max_new_tokens=100, do_sample=False)
generation = generation[0][input_len:]
decoded = processor.decode(generation, skip_special_tokens=True)
return decoded
with gr.Blocks() as demo:
gr.Markdown(
"""
# PaliGemma
Lightweight open vision-language model (VLM). [Model card](https://huggingface.co/google/paligemma-3b-mix-448)
"""
)
with gr.Row():
prompt = gr.Textbox(label="Input", value="Describe this image.", scale=4)
submit = gr.Button("Submit")
with gr.Row():
image = gr.Image(type="pil", label="Upload an Image")
output = gr.TextArea(label="Response")
submit.click(answer_question, [image, prompt], output)
prompt.submit(answer_question, [image, prompt], output)
demo.queue().launch()
|