Spaces:
Sleeping
Sleeping
File size: 2,500 Bytes
c91d9f3 c580f5e b9c7982 c91d9f3 b9c7982 c91d9f3 dc3c0b2 dcd8e07 c580f5e dcd8e07 b9c7982 c91d9f3 dcd8e07 b9c7982 c91d9f3 dcd8e07 b9c7982 dcd8e07 b9c7982 dcd8e07 c91d9f3 dcd8e07 c91d9f3 dcd8e07 c91d9f3 33262af b9c7982 33262af b9c7982 33262af c91d9f3 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import gradio as gr
from transformers import (
PaliGemmaProcessor,
PaliGemmaForConditionalGeneration,
)
from transformers.image_utils import load_image
import torch
import os
import spaces # Import the spaces module
import requests
from io import BytesIO
from PIL import Image
def load_model():
"""Load PaliGemma2 model and processor with Hugging Face token."""
token = os.getenv("HUGGINGFACEHUB_API_TOKEN") # Retrieve token from environment variable
if not token:
raise ValueError(
"Hugging Face API token not found. Please set it in the environment variables."
)
# Load the processor and model using the correct identifier
model_id = "google/paligemma2-28b-pt-448"
processor = PaliGemmaProcessor.from_pretrained(model_id, token=token)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = PaliGemmaForConditionalGeneration.from_pretrained(
model_id, torch_dtype=torch.bfloat16, token=token
).to(device).eval()
return processor, model
@spaces.GPU(duration=120) # Increased timeout to 120 seconds
def process_image_and_text(image_pil, text_input):
"""Extract text from image using PaliGemma2."""
try:
processor, model = load_model()
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the image using load_image
image = load_image(image_pil)
# Use the provided text input
model_inputs = processor(text=text_input, images=image, return_tensors="pt").to(
device, dtype=torch.bfloat16
)
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
except Exception as e:
print(f"Error during GPU task: {e}")
raise gr.Error(f"GPU task failed: {e}")
if __name__ == "__main__":
iface = gr.Interface(
fn=process_image_and_text,
inputs=[
gr.Image(type="pil", label="Upload an image"),
gr.Textbox(label="Enter Text Prompt"),
],
outputs=gr.Textbox(label="Generated Text"),
title="PaliGemma2 Image and Text to Text",
description="Upload an image and enter a text prompt. The model will generate text based on both.",
)
iface.launch() |