'update'

app.py

@@ -4,7 +4,7 @@ import fire
 import gradio as gr
 from PIL import Image
 from functools import partial
-
++import spaces 
 import cv2
 import time
 import numpy as np
@@ -62,16 +62,16 @@ _GPU_ID = 0
 if not hasattr(Image, 'Resampling'):
     Image.Resampling = Image
 
-
++@spaces.GPU
 def sam_init():
     sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
     model_type = "vit_h"
 
-    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda:{_GPU_ID}")
+    sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device="cuda")
     predictor = SamPredictor(sam)
     return predictor
 
-
++@spaces.GPU
 def sam_segment(predictor, input_image, *bbox_coords):
     bbox = np.array(bbox_coords)
     image = np.asarray(input_image)
@@ -143,7 +143,7 @@ def preprocess(predictor, input_image, chk_group=None, segment=True, rescale=Fal
         input_image = expand2square(input_image, (127, 127, 127, 0))
     return input_image, input_image.resize((320, 320), Image.Resampling.LANCZOS)
 
-
++@spaces.GPU
 def load_era3d_pipeline(cfg):
     # Load scheduler, tokenizer and models.
     
@@ -153,7 +153,7 @@ def load_era3d_pipeline(cfg):
     )
 
     if torch.cuda.is_available():
-        pipeline.to('cuda:0')
+        pipeline.to('cuda')
         pipeline.unet.enable_xformers_memory_efficient_attention()
     # sys.main_lock = threading.Lock()
     return pipeline
@@ -168,7 +168,7 @@ def prepare_data(single_image, crop_size, cfg):
     return dataset[0]
 
 scene = 'scene'
-
++@spaces.GPU
 def run_pipeline(pipeline, cfg, single_image, guidance_scale, steps, seed, crop_size, chk_group=None):
     import pdb
     global scene
@@ -302,7 +302,7 @@ def run_demo():
 
     pipeline = load_era3d_pipeline(cfg)
     torch.set_grad_enabled(False)
-    pipeline.to(f'cuda:{_GPU_ID}')
+    pipeline.to('cuda')
 
     predictor = sam_init()