Update gradio_app.py

gradio_app.py

@@ -29,10 +29,7 @@ from mvdiffusion.pipelines.pipeline_mvdiffusion_image import MVDiffusionImagePip
 from diffusers import AutoencoderKL, DDPMScheduler, DDIMScheduler
 from einops import rearrange
 import numpy as np
-
-
-
-
+from transformers import SamModel
 
 def save_image(tensor):
     ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
@@ -60,10 +57,10 @@ 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}")
+    # 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 = SamModel.from_pretrained("facebook/sam-vit-huge").to(device=f"cuda:{_GPU_ID}")
     predictor = SamPredictor(sam)
     return predictor