sergiopaniego HF Staff commited on
Commit
ebbe0c2
·
verified ·
1 Parent(s): e0683ff

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -18,19 +18,19 @@ from gradio_image_prompter import ImagePrompter
18
 
19
  #sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-huge")
20
  #sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-huge")
21
- sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-base")
22
  sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-base")
23
 
24
  #sam_model = SamModel.from_pretrained("facebook/sam-vit-huge")
25
  #sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
26
- sam_model = SamModel.from_pretrained("facebook/sam-vit-base")
27
  sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
28
 
29
  @spaces.GPU
30
  def predict_masks_and_scores(model, processor, raw_image, input_points=None, input_boxes=None):
31
  if input_boxes is not None:
32
  input_boxes = [input_boxes]
33
- inputs = processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt")
34
  with torch.no_grad():
35
  outputs = model(**inputs)
36
 
@@ -118,6 +118,7 @@ with gr.Blocks(theme=theme, title="🔍 Compare SAM vs SAM-HQ") as demo:
118
  gr.Interface(
119
  fn=process_inputs,
120
  examples=example_paths,
 
121
  inputs=ImagePrompter(show_label=False),
122
  outputs=result_html,
123
  )
 
18
 
19
  #sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-huge")
20
  #sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-huge")
21
+ sam_hq_model = SamHQModel.from_pretrained("syscv-community/sam-hq-vit-base", device_map="auto", torch_dtype="auto")
22
  sam_hq_processor = SamHQProcessor.from_pretrained("syscv-community/sam-hq-vit-base")
23
 
24
  #sam_model = SamModel.from_pretrained("facebook/sam-vit-huge")
25
  #sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
26
+ sam_model = SamModel.from_pretrained("facebook/sam-vit-base", device_map="auto", torch_dtype="auto")
27
  sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
28
 
29
  @spaces.GPU
30
  def predict_masks_and_scores(model, processor, raw_image, input_points=None, input_boxes=None):
31
  if input_boxes is not None:
32
  input_boxes = [input_boxes]
33
+ inputs = processor(raw_image, input_boxes=input_boxes, input_points=input_points, return_tensors="pt").to(model.device)
34
  with torch.no_grad():
35
  outputs = model(**inputs)
36
 
 
118
  gr.Interface(
119
  fn=process_inputs,
120
  examples=example_paths,
121
+ cache_examples=False,
122
  inputs=ImagePrompter(show_label=False),
123
  outputs=result_html,
124
  )