SAM 3 text-prompted segmentation for WebGPU: Standard (fp16) and Compact (8-bit weights)
7eee2d7 verified | """SAM 3's text-prompted segmentation, split into the three graphs the browser runs. | |
| The split follows `Sam3Model.forward` (Hugging Face Transformers, `models/sam3/modeling_sam3.py`) | |
| for a text prompt with no box prompts: | |
| - **vision**: the Perception Encoder ViT and its neck, run once per image. Its outputs are the | |
| first three feature levels (288², 144² and 72², 256 channels); `Sam3Model` drops the fourth. | |
| - **text**: the CLIP text encoder and `text_projection`, run once per phrase. | |
| - **detector**: the DETR encoder and decoder, the box head, the dot-product scoring and the mask | |
| decoder, run once per phrase over the vision graph's features. | |
| The 72² level's sine position encoding depends only on the input size, which is fixed at | |
| 1008 × 1008, so the detector carries it as a constant instead of taking it as an input. | |
| """ | |
| import torch | |
| from torch import nn | |
| from transformers.models.sam3.modeling_sam3 import box_cxcywh_to_xyxy, inverse_sigmoid | |
| IMAGE_SIZE = 1008 | |
| TEXT_TOKENS = 32 | |
| class Sam3Vision(nn.Module): | |
| def __init__(self, model): | |
| super().__init__() | |
| self.vision_encoder = model.vision_encoder | |
| def forward(self, pixel_values: torch.Tensor): | |
| levels = self.vision_encoder(pixel_values).fpn_hidden_states | |
| return levels[0], levels[1], levels[2] | |
| class Sam3Text(nn.Module): | |
| def __init__(self, model): | |
| super().__init__() | |
| self.text_encoder = model.text_encoder | |
| self.text_projection = model.text_projection | |
| def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor): | |
| hidden = self.text_encoder( | |
| input_ids=input_ids, attention_mask=attention_mask, return_dict=True | |
| ).last_hidden_state | |
| return self.text_projection(hidden) | |
| class Sam3Detector(nn.Module): | |
| def __init__(self, model, position_encoding: torch.Tensor): | |
| super().__init__() | |
| self.detr_encoder = model.detr_encoder | |
| self.detr_decoder = model.detr_decoder | |
| self.dot_product_scoring = model.dot_product_scoring | |
| self.mask_decoder = model.mask_decoder | |
| self.register_buffer("position_encoding", position_encoding) | |
| def forward( | |
| self, | |
| features_288: torch.Tensor, | |
| features_144: torch.Tensor, | |
| features_72: torch.Tensor, | |
| text_features: torch.Tensor, | |
| text_mask: torch.Tensor, | |
| ): | |
| mask = text_mask.bool() | |
| encoder = self.detr_encoder( | |
| vision_features=[features_72], | |
| text_features=text_features, | |
| vision_pos_embeds=[self.position_encoding], | |
| text_mask=mask, | |
| ) | |
| decoder = self.detr_decoder( | |
| vision_features=encoder.last_hidden_state, | |
| text_features=encoder.text_features, | |
| vision_pos_encoding=encoder.pos_embeds_flattened, | |
| text_mask=mask, | |
| spatial_shapes=encoder.spatial_shapes, | |
| ) | |
| hidden = decoder.intermediate_hidden_states[-1:] | |
| offsets = self.detr_decoder.box_head(hidden) | |
| boxes = (inverse_sigmoid(decoder.reference_boxes[-1:]) + offsets).sigmoid() | |
| logits = self.dot_product_scoring( | |
| decoder_hidden_states=hidden, | |
| text_features=encoder.text_features, | |
| text_mask=mask, | |
| ).squeeze(-1) | |
| masks = self.mask_decoder( | |
| decoder_queries=hidden[0], | |
| backbone_features=[features_288, features_144, features_72], | |
| encoder_hidden_states=encoder.last_hidden_state, | |
| prompt_features=text_features, | |
| prompt_mask=mask, | |
| ) | |
| return ( | |
| logits[0], | |
| box_cxcywh_to_xyxy(boxes)[0], | |
| decoder.presence_logits[-1], | |
| masks.pred_masks, | |
| masks.semantic_seg, | |
| ) | |
| def position_encoding_72(model) -> torch.Tensor: | |
| """The 72² level's position encoding for a 1008² image, as the vision model computes it.""" | |
| with torch.no_grad(): | |
| dummy = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE) | |
| return model.vision_encoder(dummy).fpn_position_encoding[2].clone() | |