Update app.py

app.py

@@ -7,12 +7,13 @@ from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 
 # Load model and processor
 model = Pix2StructForConditionalGeneration.from_pretrained("google/pix2struct-screen2words-large", dtype=torch.bfloat16).to("cuda")
+model.eval()
 processor = Pix2StructProcessor.from_pretrained("google/pix2struct-screen2words-large")
 
 # Define the function
 @spaces.GPU
 def describe_ui(image):
-    inputs = processor(images=image, text="describe this image: ", return_tensors="pt")
+    inputs = processor(images=image, text="describe this image: ", return_tensors="pt").to("cuda")
     predictions = model.generate(**inputs)
     return processor.decode(predictions[0], skip_special_tokens=True)