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