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()