Update app.py

app.py

@@ -1,15 +1,16 @@
-import spaces
 from transformers import AutoProcessor, AutoModelForImageTextToText
 import torch
 import gradio as gr
+import spaces
 
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 MODEL_PATH = "zai-org/GLM-OCR"
 processor = AutoProcessor.from_pretrained(MODEL_PATH)
 model = AutoModelForImageTextToText.from_pretrained(
     pretrained_model_name_or_path=MODEL_PATH,
     torch_dtype="auto",
     device_map="auto",
-)
+).to(device)
 
 
 
@@ -39,7 +40,7 @@ def read_img(img):
     add_generation_prompt=True,
     return_dict=True,
     return_tensors="pt"
-  )
+  ).to(device)
 
   inputs.pop("token_type_ids", None)
   generated_ids = model.generate(**inputs, max_new_tokens=8192)