free_blackjack / app.py
davidr99's picture
Update app.py
1146671 verified
import gradio as gr
import torch
import requests
import spaces
from PIL import Image
from huggingface_hub import create_repo
from transformers import Trainer
from transformers import TrainingArguments
from transformers import PaliGemmaProcessor
from transformers import BitsAndBytesConfig
from transformers import AutoProcessor, PaliGemmaForConditionalGeneration
from huggingface_hub import whoami
model_id = "google/paligemma-3b-pt-224"
finetuned_model_id = "davidr99/finetuned_paligemma_blackjack"
import gradio as gr
model = PaliGemmaForConditionalGeneration.from_pretrained(model_id)
model.load_adapter(finetuned_model_id)
model.to('cuda')
processor = AutoProcessor.from_pretrained(model_id)
@spaces.GPU(duration=30)
def blackjack_ai(image):
prompt = "<image><bos>extract json\n"
image_file = Image.fromarray(image.astype('uint8'), 'RGB')
inputs = processor(images=[image_file], text=prompt, return_tensors="pt").to('cuda')
output = model.generate(**inputs, max_new_tokens=125).to('cpu')
return processor.decode(output[0], skip_special_tokens=True)[14:]
with gr.Blocks() as demo:
image = gr.Image()
submit = gr.Button("Submit")
output = gr.TextArea()
submit.click(blackjack_ai, inputs=[image], outputs=[output])
demo.launch()