triphuong57's picture
Update app.py
12f4f6d verified
raw
history blame
725 Bytes
import gradio as gr
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from peft import PeftModel
import spaces
device = "cuda"
@spaces.GPU
def greet(image, prompt):
base_model = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma-3b-mix-224").to(device)
processor = AutoProcessor.from_pretrained("google/paligemma-3b-mix-224")
model = PeftModel(base_model, "/folders").to(device)
inputs = processor(prompt, raw_image, return_tensors="pt")
output = model.generate(**inputs, max_new_tokens=20)
return output
demo = gr.Interface(fn=greet, inputs=[gr.Image(label="Upload image", sources=['upload', 'webcam'], type="pil"), gr.Text()], outputs="text")
demo.launch()