cuda cpu if

app.py

@@ -12,11 +12,10 @@ from segment_anything import SamAutomaticMaskGenerator, SamPredictor, sam_model_
 
 matplotlib.pyplot.switch_backend('Agg') # for matplotlib to work in gradio
 
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # use GPU if available
 
 #setup model
 sam_checkpoint = "sam_vit_h_4b8939.pth"
-device = "cuda"
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # use GPU if available
 model_type = "default"
 sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
 sam.to(device=device)
@@ -56,10 +55,10 @@ with gr.blocks() as demo:
     gr.MArkdown("## Segment-anything Demo")
 
     with gr.Row():
-        image_input = gr.Image()
+        image = gr.Image()
         image_output = gr.Image()
 
     segment_image_button = gr.Button("Segment Image")
-    segment_image_button.click(segment_image, image_input, image_output)
+    segment_image_button.click(segment_image, inputs=[image], outputs=image_output)
 
-demo.launch()
+demo.launch()