| """Concerto (Pointcept/Concerto @ 5b411171) as a transformers ``PreTrainedModel``. |
| |
| ``ConcertoModel.model`` is the package's ``PointTransformerV3`` (vendored verbatim in ``concerto_model.py``), built |
| from the checkpoint's stored config the way ``concerto.load`` builds it; ``forward`` runs ``PointTransformerV3.forward`` |
| and returns a ``ConcertoModelOutput``: ``last_hidden_state`` (the output Point's ``feat``) and ``point``, the whole |
| output Point (an ``addict.Dict``: ``coord``, ``offset``, the pooling_parent / pooling_inverse chain, ...). |
| Runtime packages beyond torch / transformers: ``spconv``, ``torch_scatter``, ``timm``, ``addict``, ``scipy`` |
| (as the original); ``flash_attn`` when ``enable_flash`` is on. |
| """ |
|
|
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import torch |
| from transformers import PreTrainedModel |
| from transformers.utils import ModelOutput |
|
|
| from .concerto_model import PointTransformerV3 |
| from .configuration_concerto import ConcertoConfig |
|
|
|
|
| @dataclass |
| class ConcertoModelOutput(ModelOutput): |
| """last_hidden_state (M, C): features of the encoder's output points (``point.feat``); point: the output Point, whose |
| pooling_parent / pooling_inverse chain and ``inverse`` map features back to the input points.""" |
|
|
| last_hidden_state: torch.Tensor | None = None |
| point: Any = None |
|
|
|
|
| class ConcertoPreTrainedModel(PreTrainedModel): |
| config_class = ConcertoConfig |
| config: ConcertoConfig |
| base_model_prefix = "model" |
| main_input_name = "feat" |
|
|
|
|
| class ConcertoModel(ConcertoPreTrainedModel): |
| def __init__(self, config: ConcertoConfig): |
| super().__init__(config) |
| |
| ptv3_config = dict(config.ptv3_config) |
| if config.enable_flash is not None: |
| ptv3_config["enable_flash"] = config.enable_flash |
| if config.enc_patch_size is not None: |
| ptv3_config["enc_patch_size"] = list(config.enc_patch_size) |
| self.model = PointTransformerV3(**ptv3_config) |
| self.post_init() |
|
|
| def forward( |
| self, |
| coord: torch.Tensor | None = None, |
| grid_coord: torch.Tensor | None = None, |
| feat: torch.Tensor | None = None, |
| offset: torch.Tensor | None = None, |
| batch: torch.Tensor | None = None, |
| return_dict: bool | None = None, |
| **kwargs, |
| ) -> tuple | ConcertoModelOutput: |
| """A collected point cloud (``ConcertoProcessor``): coord / grid_coord (N, 3), feat (N, C_in), and offset or batch; |
| other keys the processor returns (color, inverse, ...) are kept in the Point as the original pipeline does. |
| ``return_dict=False`` returns a tuple.""" |
| return_dict = return_dict if return_dict is not None else self.config.return_dict |
| named = {"coord": coord, "grid_coord": grid_coord, "feat": feat, "offset": offset, "batch": batch} |
| point = self.model({**{k: v for k, v in named.items() if v is not None}, **kwargs}) |
| output = ConcertoModelOutput(last_hidden_state=point.feat, point=point) |
| return output if return_dict else output.to_tuple() |
|
|
|
|
| __all__ = ["ConcertoModel", "ConcertoModelOutput", "ConcertoPreTrainedModel"] |
|
|
