migrate to zero

app.py

@@ -2,7 +2,9 @@ import gradio as gr
 import requests
 from PIL import Image
 from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
+import spaces
 
+@spaces.GPU
 def infer_infographics(image, question):
     model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-ai2d-base").to("cuda")
     processor = Pix2StructProcessor.from_pretrained("google/pix2struct-ai2d-base")
@@ -12,6 +14,7 @@ def infer_infographics(image, question):
     predictions = model.generate(**inputs)
     return processor.decode(predictions[0], skip_special_tokens=True)
 
+@spaces.GPU
 def infer_ui(image, question):
     model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-screen2words-base").to("cuda")
     processor = Pix2StructProcessor.from_pretrained("google/pix2struct-screen2words-base")
@@ -21,6 +24,7 @@ def infer_ui(image, question):
     predictions = model.generate(**inputs)
     return processor.decode(predictions[0], skip_special_tokens=True)
 
+@spaces.GPU
 def infer_chart(image, question):
   model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-chartqa-base").to("cuda")
   processor = Pix2StructProcessor.from_pretrained("google/pix2struct-chartqa-base")
@@ -30,6 +34,7 @@ def infer_chart(image, question):
   predictions = model.generate(**inputs)
   return processor.decode(predictions[0], skip_special_tokens=True)
 
+@spaces.GPU
 def infer_doc(image, question):
   model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-docvqa-base").to("cuda")
   processor = Pix2StructProcessor.from_pretrained("google/pix2struct-docvqa-base")