device set

app.py

@@ -17,7 +17,7 @@ def extract_bbox_answer(content):
         return bbox
     return [0, 0, 0, 0]
 
-def process_image_and_text(image, text):
+def process_image_and_text(image, text, device):
     """Process image and text input, return thinking process and bbox"""
     question = f"Please provide the bounding box coordinate of the region this sentence describes: {text}."
     QUESTION_TEMPLATE = "{Question} First output the thinking process in <think> </think> tags and then output the final answer in <answer> </answer> tags. Output the final answer in JSON format."
@@ -45,7 +45,7 @@ def process_image_and_text(image, text):
         add_special_tokens=False,
     )
 
-    inputs = inputs.to("cuda")
+    inputs = inputs.to(device)
 
     with torch.no_grad():
         generated_ids = model.generate(**inputs, use_cache=True, max_new_tokens=256, do_sample=False)
@@ -76,11 +76,12 @@ if __name__ == "__main__":
     
     # model_path = "/data/shz/project/vlm-r1/VLM-R1/output/Qwen2.5-VL-3B-GRPO-REC/checkpoint-500"
     model_path = "SZhanZ/Qwen2.5VL-VLM-R1-REC-step500"
-    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path, torch_dtype=torch.float16, device_map="cuda")
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = Qwen2_5_VLForConditionalGeneration.from_pretrained(model_path, torch_dtype=torch.float16, device_map=device)
     processor = AutoProcessor.from_pretrained(model_path)
     
     def gradio_interface(image, text):
-        thinking, result_image = process_image_and_text(image, text)
+        thinking, result_image = process_image_and_text(image, text, device)
         return thinking, result_image
     
     demo = gr.Interface(