gokaygokay commited on
Commit
69c66b9
·
1 Parent(s): efbc3fe
Files changed (1) hide show
  1. matte_processor.py +2 -0
matte_processor.py CHANGED
@@ -1,3 +1,4 @@
 
1
  import gradio as gr
2
  from PIL import Image
3
  import torch
@@ -19,6 +20,7 @@ def load_VITMatte_model(local_files_only=False):
19
  # Load model globally
20
  vit_matte_model = load_VITMatte_model(local_files_only=False)
21
 
 
22
  def generate_VITMatte(image, trimap, local_files_only=False, max_megapixels=2.0):
23
  if image is None or trimap is None:
24
  return None
 
1
+ import spaces
2
  import gradio as gr
3
  from PIL import Image
4
  import torch
 
20
  # Load model globally
21
  vit_matte_model = load_VITMatte_model(local_files_only=False)
22
 
23
+ @spaces.GPU
24
  def generate_VITMatte(image, trimap, local_files_only=False, max_megapixels=2.0):
25
  if image is None or trimap is None:
26
  return None