space-doctor run 20260506T132411Z-acca0806

app.py

@@ -6,10 +6,9 @@ from transformers import SamModel, SamProcessor
 from gradio_image_prompter import ImagePrompter
 import spaces
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to("cuda")
+sam_model = SamModel.from_pretrained("facebook/sam-vit-huge")
 sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
-slimsam_model = SamModel.from_pretrained("nielsr/slimsam-50-uniform").to("cuda")
+slimsam_model = SamModel.from_pretrained("nielsr/slimsam-50-uniform")
 slimsam_processor = SamProcessor.from_pretrained("nielsr/slimsam-50-uniform")
 
 def get_processor_and_model(slim: bool):
@@ -21,12 +20,13 @@ def get_processor_and_model(slim: bool):
 def sam_box_inference(image, x_min, y_min, x_max, y_max, *, slim=False):
 
     processor, model = get_processor_and_model(slim)
-    
+    model.to("cuda")
+
     inputs = processor(
         Image.fromarray(image),
         input_boxes=[[[[x_min, y_min, x_max, y_max]]]],
         return_tensors="pt"
-    ).to(device)
+    ).to("cuda")
 
     with torch.no_grad():
         outputs = model(**inputs)
@@ -45,11 +45,12 @@ def sam_box_inference(image, x_min, y_min, x_max, y_max, *, slim=False):
 def sam_point_inference(image, x, y, *, slim=False):
 
     processor, model = get_processor_and_model(slim)
-    
+    model.to("cuda")
+
     inputs = processor(
         image,
         input_points=[[[x, y]]],
-        return_tensors="pt").to(device)
+        return_tensors="pt").to("cuda")
 
     with torch.no_grad():
         outputs = model(**inputs)
@@ -88,7 +89,10 @@ def infer_point(img):
 
 def infer_box(prompts):
     # background (original image) layers[0] ( point prompt) composite (total image)
-    image = prompts["image"]
+    image = None
+    if prompts is None:
+      raise gr.Error("Please upload an image and draw a box before submitting")
+    image = prompts.get("image")
     if image is None:
       raise gr.Error("Please upload an image and draw a box before submitting")
     points = prompts["points"][0]
@@ -135,4 +139,4 @@ with gr.Blocks(title="SlimSAM") as demo:
           output_sam = gr.AnnotatedImage(label="SAM Output")
 
     im.change(infer_point, inputs=im, outputs=[output_slimsam, output_sam])
-demo.launch(debug=True)
+demo.launch()