Class SpatialVLAForConditionalGeneration

model/modeling_spatialvla.py:162–526

                module.weight.data[module.padding_idx].zero_()

class SpatialVLAForConditionalGeneration(SpatialVLAPreTrainedModel, GenerationMixin):
    def __init__(self, config: SpatialVLAConfig, vision_model=None, vision_zoe_model=None, projector_model=None, language_model=None):
        super().__init__(config)

        self.vision_tower = vision_model or AutoModel.from_config(config=config.vision_config)
        self.multi_modal_projector = projector_model or SpatialVLAMultiModalProjector(config)
        self.vocab_size = config.text_config.vocab_size
        if language_model is None:
            language_model = Gemma2ForCausalLM(config=config.text_config)
        if language_model._tied_weights_keys is not None:
            self._tied_weights_keys = [f"language_model.{k}" for k in language_model._tied_weights_keys]
        self.language_model = language_model

        if config.use_vision_zoe:
            self.vision_zoe_model = vision_zoe_model or ZoeDepthForDepthEstimation(config.vision_zoe_config)
            self.position_embedding_3d = Ego3DPositionEmbeddingMLP(
                config.ego3d_patch_reso**2 * 3, num_pos_feats=config.vision_config.hidden_size, n_freqs=config.n_freqs
            )
            # register buffer
            patch_size, reso, image_size = config.vision_config.patch_size, config.ego3d_patch_reso, config.vision_config.image_size
            y, x = torch.meshgrid(torch.arange(0, image_size, patch_size // reso), torch.arange(0, image_size, patch_size // reso), indexing="ij")  # (h//sp w//sp)
            y, x = y + patch_size / reso / 2, x + patch_size / reso / 2
            uv_h = torch.stack([x, y, torch.ones_like(x)], dim=0).reshape(3, -1)  # (3 hw)
            self.register_buffer("uv_h", uv_h, persistent=False)

        # shared spatial embeddings for <ACTION> <IMG>
        if config.use_spatial_token:
            self.spatial_embed_tokens = nn.Embedding(self.config.spatial_token_num, config.text_config.hidden_size)
        else:
            self.spatial_embed_tokens = None
        self.pad_token_id = self.config.pad_token_id if self.config.pad_token_id is not None else -1


    def backproject_patch(self, K: torch.Tensor, depth: torch.Tensor, patch_size=14, reso=2) -> torch.Tensor:
        """
        Backproject depth map to 3D points in camera coordinate.
        Args:
            K: camera intrinsic matrix (b 3 3)
            depth: depth map (b 1 h w)
            patch_size: patch size for siglip
            reso: reso^2 -> sample points in each patch
        patch sz = 14  ......          
        ┌────────┬────────┐       
        │ ─    ─ │ ─    ─ │       
        │ points │        ├─ ─ ─ 
        │ ─    ─ │ ─    ─ │       
        ├────────┼────────┤       
        │ ─    ─ │ ─    ─ │       
        │        │        │       
        │ ─    ─ │ ─    ─ │       
        └────────┴────────┘       
        reso=2───►points=4
            │                    
            │                        
        """
        b, c, h, w = depth.shape
        hp, wp = h // patch_size, w // patch_size
        sub_hp = sub_wp = reso