csgaobb commited on
Commit
a03a063
·
verified ·
1 Parent(s): 797bf1a

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -25,6 +25,7 @@ from torchvision.transforms.functional import pil_to_tensor
25
 
26
  from metauas import MetaUAS, set_random_seed, normalize, apply_ad_scoremap, safely_load_state_dict
27
  from huggingface_hub import hf_hub_download
 
28
 
29
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
30
 
@@ -50,6 +51,7 @@ model = MetaUAS(encoder_name,
50
  fusion_policy
51
  )
52
 
 
53
  def process_image(prompt_img, query_img, options):
54
  # Load the model based on selected options
55
  if 'model-512' in options:
 
25
 
26
  from metauas import MetaUAS, set_random_seed, normalize, apply_ad_scoremap, safely_load_state_dict
27
  from huggingface_hub import hf_hub_download
28
+ import spaces
29
 
30
  device = 'cuda' if torch.cuda.is_available() else 'cpu'
31
 
 
51
  fusion_policy
52
  )
53
 
54
+ @spaces.GPU
55
  def process_image(prompt_img, query_img, options):
56
  # Load the model based on selected options
57
  if 'model-512' in options: