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 pipeline
+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,16 +57,18 @@ if not hasattr(Image, 'Resampling'):
 
 
 def sam_init():
-    predictor =  pipeline("mask-generation", device = f"cuda:{_GPU_ID}", points_per_batch = 256)
-    return predictor
+    model = SamModel.from_pretrained("facebook/sam-vit-huge")
+    processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+    return model, processor
 
-def sam_segment(predictor, input_image, *bbox_coords):
+def sam_segment(sam_model, sam_processor, input_image, *bbox_coords):
     bbox = np.array(bbox_coords)
     image = np.asarray(input_image)
 
     start_time = time.time()
 
-    outputs = predictor(input_image, points_per_batch = 256)
+    inputs = sam_processor(raw_image, input_boxes=bbox, return_tensors="pt").to("cuda")
+    outputs = sam_model(**inputs)
 
     print(f"SAM Time: {time.time() - start_time:.3f}s")
     out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
@@ -89,7 +91,7 @@ def expand2square(pil_img, background_color):
         result.paste(pil_img, ((height - width) // 2, 0))
         return result
 
-def preprocess(predictor, input_image, chk_group=None, segment=True, rescale=False):
+def preprocess(sam_model, sam_processor, input_image, chk_group=None, segment=True, rescale=False):
     RES = 1024
     input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
     if chk_group is not None:
@@ -105,7 +107,7 @@ def preprocess(predictor, input_image, chk_group=None, segment=True, rescale=Fal
         y_min = int(y_nonzero[0].min())
         x_max = int(x_nonzero[0].max())
         y_max = int(y_nonzero[0].max())
-        input_image = sam_segment(predictor, input_image.convert('RGB'), x_min, y_min, x_max, y_max)
+        input_image = sam_segment(sam_model, sam_processor, input_image.convert('RGB'), x_min, y_min, x_max, y_max)
     # Rescale and recenter
     if rescale:
         image_arr = np.array(input_image)
@@ -253,7 +255,7 @@ def run_demo():
     torch.set_grad_enabled(False)
     pipeline.to(f'cuda:{_GPU_ID}')
 
-    predictor = sam_init()
+    sam_model, sam_processor = sam_init()
 
     custom_theme = gr.themes.Soft(primary_hue="blue").set(
                     button_secondary_background_fill="*neutral_100",
@@ -328,7 +330,7 @@ def run_demo():
             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), 
+        run_btn.click(fn=partial(preprocess, sam_model, sam_processor), 
                         inputs=[input_image, input_processing], 
                         outputs=[processed_image_highres, processed_image], queue=True
             ).success(fn=partial(run_pipeline, pipeline, cfg),