Repository: huggingface/transformers
Pull request: [SAM3] Fix precompute vision_embeds or text_embeds for inference

# What does this PR do?

Some checks in the processing were preventing inference with precomputed vision_embeds or text_embeds.
This fixes it, add examples and tests


Changed files:
filename: docs/source/en/model_doc/sam3.md
patch: @@ -256,6 +256,104 @@ SAM3 also provides semantic segmentation alongside instance masks:
 >>> print(f"Semantic segmentation: {semantic_seg.shape}")
 ```
 
+### Efficient Multi-Prompt Inference on Single Image
+
+When running multiple text prompts on the same image, pre-compute vision embeddings to avoid redundant computation:
+
+```python
+>>> from transformers import Sam3Processor, Sam3Model
+>>> import torch
+>>> from PIL import Image
+>>> import requests
+
+>>> device = "cuda" if torch.cuda.is_available() else "cpu"
+
+>>> model = Sam3Model.from_pretrained("facebook/sam3").to(device)
+>>> processor = Sam3Processor.from_pretrained("facebook/sam3")
+
+>>> # Load image
+>>> image_url = "http://images.cocodataset.org/val2017/000000077595.jpg"
+>>> image = Image.open(requests.get(image_url, stream=True).raw).convert("RGB")
+
+>>> # Pre-process image and compute vision embeddings once
+>>> img_inputs = processor(images=image, return_tensors="pt").to(device)
+>>> with torch.no_grad():
+...     vision_embeds = model.get_vision_features(pixel_values=img_inputs.pixel_values)
+
+>>> # Run multiple text prompts efficiently
+>>> text_prompts = ["ear", "eye", "nose"]
+>>> all_results = []
+
+>>> for prompt in text_prompts:
+...     text_inputs = processor(text=prompt, return_tensors="pt").to(device)
+...     with torch.no_grad():
+...         outputs = model(vision_embeds=vision_embeds, **text_inputs)
+...
+...     results = processor.post_process_instance_segmentation(
+...         outputs,
+...         threshold=0.5,
+...         mask_threshold=0.5,
+...         target_sizes=img_inputs.get("original_sizes").tolist()
+...     )[0]
+...     all_results.append({"prompt": prompt, "results": results})
+
+>>> for item in all_results:
+...     print(f"Prompt '{item['prompt']}': {len(item['results']['masks'])} objects found")
+```
+
+### Efficient Single-Prompt Inference on Multiple Images
+
+When running the same text prompt on multiple images, pre-compute text embeddings to avoid redundant computation:
+
+```python
+>>> from transformers import Sam3Processor, Sam3Model
+>>> import torch
+>>> from PIL import Image
+>>> import requests
+
+>>> device = "cuda" if torch.cuda.is_available() else "cpu"
+
+>>> model = Sam3Model.from_pretrained("facebook/sam3").to(device)
+>>> processor = Sam3Processor.from_pretrained("facebook/sam3")
+
+>>> # Pre-compute text embeddings once
+>>> text_prompt = "ear"
+>>> text_inputs = processor(text=text_prompt, return_tensors="pt").to(device)
+>>> with torch.no_grad():
+...     text_embeds = model.get_text_features(**text_inputs)
+
+>>> # Load multiple images
+>>> image_urls = [
+...     "http://images.cocodataset.org/val2017/000000077595.jpg",
+...     "http://images.cocodataset.org/val2017/000000039769.jpg",
+... ]
+>>> images = [Image.open(requests.get(url, stream=True).raw).convert("RGB") for url in image_urls]
+
+>>> # Run inference on each image reusing text embeddings
+>>> # Note: attention_mask must be passed along with text_embeds for proper masking
+>>> all_results = []
+
+>>> for image in images:
+...     img_inputs = processor(images=image, return_tensors="pt").to(device)
+...     with torch.no_grad():
+...         outputs = model(
+...             pixel_values=img_inputs.pixel_values,
+...             text_embeds=text_embeds,
+...             attention_mask=text_inputs.attention_mask,
+...         )
+...
+...     results = processor.post_process_instance_segmentation(
+...         outputs,
+...         threshold=0.5,
+...         mask_threshold=0.5,
+...         target_sizes=img_inputs.get("original_sizes").tolist()
+...     )[0]
+...     all_results.append(results)
+
+>>> for i, results in enumerate(all_results):
+...     print(f"Image {i+1}: {len(results['masks'])} '{text_prompt}' objects found")
+```
+
 ### Prompt Label Conventions
 
 SAM3 uses the following label conventions:

filename: src/transformers/models/sam3/modeling_sam3.py
patch: @@ -2146,8 +2146,8 @@ def get_text_features(
         >>> from PIL import Image
         >>> import requests
 
-        >>> model = Sam3Model.from_pretrained("facebook/sam3-base")
-        >>> processor = Sam3Processor.from_pretrained("facebook/sam3-base")
+        >>> model = Sam3Model.from_pretrained("facebook/sam3")
+        >>> processor = Sam3Processor.from_pretrained("facebook/sam3")
 
         >>> # Pre-compute text embeddings
         >>> text_inputs = processor(text="cat", return_tensors="pt")
@@ -2184,8 +2184,8 @@ def get_vision_features(
         >>> from PIL import Image
         >>> import requests
 
-        >>> model = Sam3Model.from_pretrained("facebook/sam3-base")
-        >>> processor = Sam3Processor.from_pretrained("facebook/sam3-base")
+        >>> model = Sam3Model.from_pretrained("facebook/sam3")
+        >>> processor = Sam3Processor.from_pretrained("facebook/sam3")
 
         >>> # Pre-compute vision embeddings
         >>> img_url = "http://images.cocodataset.org/val2017/000000077595.jpg"
@@ -2233,8 +2233,8 @@ def forward(
         >>> import requests
         >>> from transformers import AutoModel, AutoProcessor
 
-        >>> model = AutoModel.from_pretrained("facebook/sam3-base")
-        >>> processor = AutoProcessor.from_pretrained("facebook/sam3-base")
+        >>> model = AutoModel.from_pretrained("facebook/sam3")
+        >>> processor = AutoProcessor.from_pretrained("facebook/sam3")
 
         >>> img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/sam-car.png"
         >>> raw_image = Image.open(requests.get(img_url, stream=True).raw).convert("RGB")

filename: src/transformers/models/sam3/processing_sam3.py
patch: @@ -151,8 +151,9 @@ def __call__(
             - `input_boxes_labels` (`torch.Tensor`): The processed labels for the bounding boxes.
             - `input_boxes` (`torch.Tensor`): The processed bounding boxes.
         """
+        encoding = None
         if images is not None:
-            encoding_image_processor = self.image_processor(
+            encoding = self.image_processor(
                 images,
                 segmentation_maps=segmentation_maps,
                 return_tensors=return_tensors,
@@ -161,24 +162,21 @@ def __call__(
         elif original_sizes is not None:
             if isinstance(original_sizes, torch.Tensor):
                 original_sizes = original_sizes.cpu().tolist()
-            encoding_image_processor = BatchEncoding({"original_sizes": original_sizes}, tensor_type=return_tensors)
-        else:
-            raise ValueError("Either images or original_sizes must be provided")
+            encoding = BatchEncoding({"original_sizes": original_sizes}, tensor_type=return_tensors)
+        elif input_boxes is not None:
+            raise ValueError("Either images or original_sizes must be provided if input_boxes is not None")
 
-        original_sizes = encoding_image_processor["original_sizes"]
-        # Check original_sizes is of length 1 or len(images)
-        if images is not None and len(original_sizes) != 1 and len(original_sizes) != len(images):
-            raise ValueError(
-                "original_sizes must be of length 1 or len(images). If you are passing a single image, you must pass a single original_size."
-            )
         text = self._resolve_text_prompts(text, input_boxes)
-
-        encoding_image_processor.update(
-            self.tokenizer(text, return_tensors=return_tensors, padding="max_length", max_length=32)
-        )
+        if text is not None:
+            text_inputs = self.tokenizer(text, return_tensors=return_tensors, padding="max_length", max_length=32)
+            if encoding is not None:
+                encoding.update(text_inputs)
+            else:
+                encoding = text_inputs
 
         # Process input boxes if provided
         if input_boxes is not None:
+            original_sizes = encoding["original_sizes"]
             # Validate and convert inputs to standardized format
             processed_boxes = self._validate_single_input(
                 input_boxes,
@@ -215,14 +213,14 @@ def __call__(
                     final_boxes, original_sizes, is_bounding_box=True, preserve_padding=True
                 )
                 final_boxes = box_xyxy_to_cxcywh(final_boxes)
-                encoding_image_processor.update({"input_boxes": final_boxes})
+                encoding.update({"input_boxes": final_boxes})
 
             if processed_boxes_labels is not None:
                 padded_boxes_labels = self._pad_nested_list(processed_boxes_labels, boxes_labels_max_dims)
                 final_boxes_labels = torch.tensor(padded_boxes_labels, dtype=torch.int64)
-                encoding_image_processor.update({"input_boxes_labels": final_boxes_labels})
+                encoding.update({"input_boxes_labels": final_boxes_labels})
 
-        return encoding_image_processor
+        return encoding
 
     def _normalize_coordinates(self, coords: "torch.Tensor", original_size, is_bounding_box=False) -> "torch.Tensor":
         """

filename: tests/models/sam3/test_modeling_sam3.py
patch: @@ -1447,3 +1447,100 @@ def test_semantic_segmentation_output(self):
         self.assertEqual(outputs.semantic_seg.shape, (1, 1, 288, 288))
         # Check that semantic seg has same spatial size as pred_masks
         self.assertEqual(outputs.semantic_seg.shape[-2:], outputs.pred_masks.shape[-2:])
+
+    def test_efficient_multi_prompt_single_image(self):
+        """Test efficient inference with multiple prompts on a single image using get_vision_features."""
+        raw_image = prepare_coco_cat_image()
+
+        # Pre-compute vision embeddings once
+        img_inputs = self.processor(images=raw_image, return_tensors="pt").to(torch_device)
+        with torch.no_grad():
+            vision_embeds = self.model.get_vision_features(pixel_values=img_inputs.pixel_values)
+
+        # Run multiple text prompts efficiently
+        text_prompts = ["ear", "eye"]
+        all_results = []
+
+        for prompt in text_prompts:
+            text_inputs = self.processor(text=prompt, return_tensors="pt").to(torch_device)
+            with torch.no_grad():
+                outputs = self.model(vision_embeds=vision_embeds, **text_inputs)
+
+            results = self.processor.post_process_instance_segmentation(
+                outputs,
+                threshold=0.5,
+                mask_threshold=0.5,
+                target_sizes=img_inputs.get("original_sizes").tolist(),
+            )[0]
+            all_results.append(results)
+
+        # Check that we get results for both prompts
+        self.assertEqual(len(all_results), 2)
+
+        # Verify outputs are equivalent to running with pixel_values directly
+        text_inputs = self.processor(text="ear", return_tensors="pt").to(torch_device)
+        with torch.no_grad():
+            outputs_with_embeds = self.model(vision_embeds=vision_embeds, **text_inputs)
+
+        inputs_direct = self.processor(images=raw_image, text="ear", return_tensors="pt").to(torch_device)
+        with torch.no_grad():
+            outputs_direct = self.model(**inputs_direct)
+
+        # Outputs should be identical
+        torch.testing.assert_close(outputs_with_embeds.pred_logits, outputs_direct.pred_logits, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(outputs_with_embeds.pred_boxes, outputs_direct.pred_boxes, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(outputs_with_embeds.pred_masks, outputs_direct.pred_masks, atol=1e-5, rtol=1e-5)
+
+    def test_efficient_single_prompt_multi_images(self):
+        """Test efficient inference with same prompt on multiple images using get_text_features."""
+        raw_image1 = prepare_coco_cat_image()
+        raw_image2 = prepare_coco_kitchen_image()
+
+        # Pre-compute text embeddings once
+        text_prompt = "handle"
+        text_inputs = self.processor(text=text_prompt, return_tensors="pt").to(torch_device)
+        with torch.no_grad():
+            text_embeds = self.model.get_text_features(**text_inputs)
+
+        # Run inference on multiple images reusing text embeddings
+        # Note: attention_mask must be passed along with text_embeds for proper masking
+        images = [raw_image1, raw_image2]
+        all_results = []
+
+        for image in images:
+            img_inputs = self.processor(images=image, return_tensors="pt").to(torch_device)
+            with torch.no_grad():
+                outputs = self.model(
+                    text_embeds=text_embeds,
+                    attention_mask=text_inputs.attention_mask,
+                    **img_inputs,
+                )
+
+            results = self.processor.post_process_instance_segmentation(
+                outputs,
+                threshold=0.5,
+                mask_threshold=0.5,
+                target_sizes=img_inputs.get("original_sizes").tolist(),
+            )[0]
+            all_results.append(results)
+
+        # Check that we get results for both images
+        self.assertEqual(len(all_results), 2)
+
+        # Verify outputs are equivalent to running with input_ids directly
+        img_inputs = self.processor(images=raw_image2, return_tensors="pt").to(torch_device)
+        with torch.no_grad():
+            outputs_with_embeds = self.model(
+                text_embeds=text_embeds,
+                attention_mask=text_inputs.attention_mask,
+                **img_inputs,
+            )
+
+        inputs_direct = self.processor(images=raw_image2, text=text_prompt, return_tensors="pt").to(torch_device)
+        with torch.no_grad():
+            outputs_direct = self.model(**inputs_direct)
+
+        # Outputs should be identical
+        torch.testing.assert_close(outputs_with_embeds.pred_logits, outputs_direct.pred_logits, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(outputs_with_embeds.pred_boxes, outputs_direct.pred_boxes, atol=1e-5, rtol=1e-5)
+        torch.testing.assert_close(outputs_with_embeds.pred_masks, outputs_direct.pred_masks, atol=1e-5, rtol=1e-5)