Update gradio_app.py

gradio_app.py

@@ -29,7 +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
+from transformers import SamModel, SamProcessor
 
 def save_image(tensor):
     ndarr = tensor.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
@@ -57,11 +57,7 @@ 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 = SamModel.from_pretrained("facebook/sam-vit-huge").to(device=f"cuda:{_GPU_ID}")
-    predictor = SamPredictor(sam)
+    predictor =  pipeline("mask-generation", device = f"cuda:{_GPU_ID}", points_per_batch = 256)
     return predictor
 
 def sam_segment(predictor, input_image, *bbox_coords):
@@ -71,16 +67,13 @@ def sam_segment(predictor, input_image, *bbox_coords):
     start_time = time.time()
     predictor.set_image(image)
 
-    masks_bbox, scores_bbox, logits_bbox = predictor.predict(
-        box=bbox,
-        multimask_output=True
-    )
+    generator(image_url, points_per_batch = 256)
 
     print(f"SAM Time: {time.time() - start_time:.3f}s")
     out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
     out_image[:, :, :3] = image
     out_image_bbox = out_image.copy()
-    out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255
+    out_image_bbox[:, :, 3] = outputs["masks"][-1].astype(np.uint8) * 255
     torch.cuda.empty_cache()
     return Image.fromarray(out_image_bbox, mode='RGBA')