pengHTYX commited on
Commit
844ef15
1 Parent(s): 501a6cc
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -4,7 +4,7 @@ import fire
4
  import gradio as gr
5
  from PIL import Image
6
  from functools import partial
7
- +import spaces
8
  import cv2
9
  import time
10
  import numpy as np
@@ -62,7 +62,7 @@ _GPU_ID = 0
62
  if not hasattr(Image, 'Resampling'):
63
  Image.Resampling = Image
64
 
65
- +@spaces.GPU
66
  def sam_init():
67
  sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
68
  model_type = "vit_h"
@@ -71,7 +71,7 @@ def sam_init():
71
  predictor = SamPredictor(sam)
72
  return predictor
73
 
74
- +@spaces.GPU
75
  def sam_segment(predictor, input_image, *bbox_coords):
76
  bbox = np.array(bbox_coords)
77
  image = np.asarray(input_image)
@@ -143,7 +143,7 @@ def preprocess(predictor, input_image, chk_group=None, segment=True, rescale=Fal
143
  input_image = expand2square(input_image, (127, 127, 127, 0))
144
  return input_image, input_image.resize((320, 320), Image.Resampling.LANCZOS)
145
 
146
- +@spaces.GPU
147
  def load_era3d_pipeline(cfg):
148
  # Load scheduler, tokenizer and models.
149
 
@@ -168,7 +168,7 @@ def prepare_data(single_image, crop_size, cfg):
168
  return dataset[0]
169
 
170
  scene = 'scene'
171
- +@spaces.GPU
172
  def run_pipeline(pipeline, cfg, single_image, guidance_scale, steps, seed, crop_size, chk_group=None):
173
  import pdb
174
  global scene
 
4
  import gradio as gr
5
  from PIL import Image
6
  from functools import partial
7
+ import spaces
8
  import cv2
9
  import time
10
  import numpy as np
 
62
  if not hasattr(Image, 'Resampling'):
63
  Image.Resampling = Image
64
 
65
+ @spaces.GPU
66
  def sam_init():
67
  sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
68
  model_type = "vit_h"
 
71
  predictor = SamPredictor(sam)
72
  return predictor
73
 
74
+ @spaces.GPU
75
  def sam_segment(predictor, input_image, *bbox_coords):
76
  bbox = np.array(bbox_coords)
77
  image = np.asarray(input_image)
 
143
  input_image = expand2square(input_image, (127, 127, 127, 0))
144
  return input_image, input_image.resize((320, 320), Image.Resampling.LANCZOS)
145
 
146
+ @spaces.GPU
147
  def load_era3d_pipeline(cfg):
148
  # Load scheduler, tokenizer and models.
149
 
 
168
  return dataset[0]
169
 
170
  scene = 'scene'
171
+ @spaces.GPU
172
  def run_pipeline(pipeline, cfg, single_image, guidance_scale, steps, seed, crop_size, chk_group=None):
173
  import pdb
174
  global scene