update

app.py

@@ -17,7 +17,7 @@ import cv2
 import time
 import numpy as np
 import trimesh
-from segment_anything import sam_model_registry, SamPredictor
+from segment_anything import build_sam, SamPredictor
 
 import random
 from pytorch3d import transforms
@@ -56,9 +56,9 @@ if not hasattr(Image, 'Resampling'):
 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}")
-    predictor = SamPredictor(sam)
+    # sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda:{_GPU_ID}")
+    # predictor = SamPredictor(sam)
+    predictor = SamPredictor(build_sam(checkpoint=sam_checkpoint).to("cuda"))
     return predictor
 
 
@@ -518,7 +518,7 @@ def run_demo():
     torch.cuda.set_device(_GPU_ID)
     args.rank = _GPU_ID
     args.world_size = 1
-    args.gpu = os.environ['CUDA_VISIBLE_DEVICES']
+    args.gpu = f'{_GPU_ID}'
     device = f'cuda:{_GPU_ID}'
 
     resolution = (256, 256)
@@ -607,11 +607,6 @@ def run_demo():
             shape_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="Reconstructed Model")
             shape_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  label="Bank Base Shape Model")
 
-        with gr.Row():
-            view_gallery = gr.Gallery(interactive=False, show_label=False, container=True, preview=True,  allow_preview=False, height=1200)
-            normal_gallery = gr.Gallery(interactive=False, show_label=False, container=True, preview=True, allow_preview=False, height=1200)
-        
-
         run_btn.click(fn=partial(preprocess, predictor), 
                         inputs=[input_image, input_processing], 
                         outputs=[processed_image_highres, processed_image], queue=True
@@ -620,6 +615,8 @@ def run_demo():
                         outputs=[view_1, view_2, shape_1, shape_2]
                         )
         demo.queue().launch(share=True, max_threads=80)
+        # _, local_url, share_url = demo.launch(share=True, server_name="0.0.0.0", server_port=23425)
+        # print('local_url: ', local_url)
 
 
 if __name__ == '__main__':