Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 import torch
+import spaces
 
 # Dictionary of model names and their corresponding HuggingFace model IDs
 MODEL_OPTIONS = {
@@ -31,21 +32,20 @@ def load_model(model_name):
         current_model = VisionEncoderDecoderModel.from_pretrained(model_id)
         current_model_name = model_name
         
-        # Move model to GPU if available
-        if torch.cuda.is_available():
-            current_model = current_model.to('cuda')
+        # Move model to GPU
+        current_model = current_model.to('cuda')
     
     return current_processor, current_model
 
+@spaces.GPU
 def process_image(image, model_name):
     processor, model = load_model(model_name)
     
     # Prepare image
     pixel_values = processor(image, return_tensors="pt").pixel_values
     
-    # Move input to GPU if model is on GPU
-    if next(model.parameters()).is_cuda:
-        pixel_values = pixel_values.to('cuda')
+    # Move input to GPU
+    pixel_values = pixel_values.to('cuda')
     
     # Generate (no beam search)
     with torch.no_grad():