Update app.py

app.py

@@ -1,3 +1,13 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForVision2Seq, AutoImageProcessor
+from PIL import Image
+import spaces
+import os
+
+from huggingface_hub import login
+login(os.environ["HF_KEY"])
+
 # Load the model and tokenizer
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModelForVision2Seq.from_pretrained("stabilityai/japanese-stable-vlm", trust_remote_code=True, device_map='auto')
@@ -32,6 +42,7 @@ def build_prompt(task="caption", input=None, sep="\n\n### "):
     return p
 
 # Define the function to generate text from the image and prompt
+@spaces.GPU(duration=120)
 def generate_text(image, task, input_text=None):
     prompt = build_prompt(task=task, input=input_text)
     inputs = processor(images=image, return_tensors="pt")