test transformers sam

app.py

@@ -17,14 +17,16 @@ import colorsys
 
 sam_checkpoint = "sam_vit_h_4b8939.pth"
 model_type = "vit_h"
-device = "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 
-sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
-sam.to(device=device)
-predictor = SamPredictor(sam)
-mask_generator = SamAutomaticMaskGenerator(sam)
+#sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+#sam.to(device=device)
+#predictor = SamPredictor(sam)
+#mask_generator = SamAutomaticMaskGenerator(sam)
 
+generator = pipeline(model="facebook/sam-vit-base", task="mask-generation", points_per_batch=256)
+#image_url = "https://huggingface.co/ybelkada/segment-anything/resolve/main/assets/car.png"
 
 controlnet, controlnet_params = FlaxControlNetModel.from_pretrained(
     "SAMControlNet/sd-controlnet-sam-seg", dtype=jnp.float32
@@ -70,18 +72,27 @@ with gr.Blocks() as demo:
         clear = gr.Button("Clear")
 
     def generate_mask(image):
-        predictor.set_image(image)
-        input_point = np.array([120, 21])
-        input_label = np.ones(input_point.shape[0])
-        mask, _, _ = predictor.predict(
-            point_coords=input_point,
-            point_labels=input_label,
-            multimask_output=False,
-        )
+        outputs = generator(image, points_per_batch=256)
+        
+        for mask in outputs["masks"]:
+            color = np.concatenate([np.random.random(3), np.array([1.0])], axis=0)
+            h, w = mask.shape[-2:]
+            mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+        
+        return mask_image
+
+        # predictor.set_image(image)
+        # input_point = np.array([120, 21])
+        # input_label = np.ones(input_point.shape[0])
+        # mask, _, _ = predictor.predict(
+        #     point_coords=input_point,
+        #     point_labels=input_label,
+        #     multimask_output=False,
+        # )
 
         # clear torch cache
-        torch.cuda.empty_cache()
-        mask = Image.fromarray(mask[0, :, :])
+        # torch.cuda.empty_cache()
+        # mask = Image.fromarray(mask[0, :, :])
         # segs = mask_generator.generate(image)
         # boolean_masks = [s["segmentation"] for s in segs]
         # finseg = np.zeros(
@@ -99,9 +110,9 @@ with gr.Blocks() as demo:
         #     rgb_mask[:, :, 2] = boolean_mask * rgb[2]
         #     finseg += rgb_mask
 
-        torch.cuda.empty_cache()
+        # torch.cuda.empty_cache()
 
-        return mask
+        # return mask
 
     def infer(
         image, prompts, negative_prompts, num_inference_steps=50, seed=4, num_samples=4

requirements.txt

@@ -1,7 +1,7 @@
 torch
 torchvision
 git+https://github.com/facebookresearch/segment-anything.git
-transformers
+git+https://github.com/huggingface/transformers@main
 flax
 jax[cuda11_pip]
 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html