#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           This file was automatically generated from src/transformers/models/sapiens2/modular_sapiens2.py.
#               Do NOT edit this file manually as any edits will be overwritten by the generation of
#             the file from the modular. If any change should be done, please apply the change to the
#                          modular_sapiens2.py file directly. One of our CI enforces this.
#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# Copyright 2026 Meta Platforms, Inc. and the HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Sapiens2 License. You may obtain a copy of the License at
#
#     https://github.com/facebookresearch/sapiens2/blob/main/LICENSE.md
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from huggingface_hub.dataclasses import strict

from ...backbone_utils import BackboneConfigMixin
from ...configuration_utils import PreTrainedConfig
from ...utils import auto_docstring


@auto_docstring(checkpoint="facebook/sapiens2-seg-0.4b")
@strict
class Sapiens2HeadConfig(PreTrainedConfig):
    r"""
    upsample_out_channels (`list[int]`, *optional*):
        Output channel counts for each upsample block.
        The first block takes `hidden_size` channels as input; subsequent blocks use the previous output.
    upsample_kernel_sizes (`list[int]`, *optional*):
        Kernel size for each upsample block. Auto-filled with `[4, ...]` when
        `upsample_out_channels` is set but this is `None`.
        Must have the same length as `upsample_out_channels`.
    upsample_kernel_size (`int`, defaults to 4):
        Default kernel size for upsample blocks when `upsample_kernel_sizes` is not set.
    use_pixel_shuffle (`bool`, *optional*):
        Whether the upsample head uses pixel-shuffle upsampling instead of transposed convolutions.
        When `None` (default), the head uses transposed convolutions.
    conv_out_channels (`list[int]`, *optional*):
        Output channel counts for the refinement conv layers that follow the upsample blocks.
    conv_kernel_sizes (`list[int]`, *optional*):
        Kernel size for each refinement conv layer. Auto-filled with `[1, ...]` when
        `conv_out_channels` is set but this is `None`.
        Must have the same length as `conv_out_channels`.
    conv_kernel_size (`int`, defaults to 1):
        Default kernel size for conv layers when `conv_kernel_sizes` is not set.
    scale_conv_out_channels (`list[int]`, *optional*):
        Output channel counts for the stride-2 conv layers used to predict the focal-length scale.
        When `None` (default), no scale branch is built.
    scale_conv_kernel_sizes (`list[int]`, *optional*):
        Kernel size for each scale conv layer. Auto-filled with `[1, ...]` when
        `scale_conv_out_channels` is set but this is `None`.
        Must have the same length as `scale_conv_out_channels`.
    scale_conv_kernel_size (`int`, defaults to 1):
        Default kernel size for scale conv layers when `scale_conv_kernel_sizes` is not set.
    scale_final_input_size (`int`, *optional*):
        Flattened feature size passed into the scale MLP.
        When `None` (default), it is automatically inferred from `image_size` and `patch_size`
        in the parent [`Sapiens2Config`].
    scale_final_hidden_sizes (`list[int]`, *optional*):
        Hidden-layer sizes for the MLP that maps flattened scale features to the scalar scale output.
        When `None` (default), no scale branch is built.
    """

    model_type = "sapiens2_head"
    base_config_key = "head_config"

    upsample_out_channels: list[int] | None = None
    upsample_kernel_sizes: list[int] | None = None
    upsample_kernel_size: int = 4
    use_pixel_shuffle: bool | None = None
    conv_out_channels: list[int] | None = None
    conv_kernel_sizes: list[int] | None = None
    conv_kernel_size: int = 1
    scale_conv_out_channels: list[int] | None = None
    scale_conv_kernel_sizes: list[int] | None = None
    scale_conv_kernel_size: int = 1
    scale_final_input_size: int | None = None
    scale_final_hidden_sizes: list[int] | None = None

    def __post_init__(self, **kwargs):
        if self.upsample_out_channels is not None and self.upsample_kernel_sizes is None:
            self.upsample_kernel_sizes = [self.upsample_kernel_size] * len(self.upsample_out_channels)
        if self.conv_out_channels is not None and self.conv_kernel_sizes is None:
            self.conv_kernel_sizes = [self.conv_kernel_size] * len(self.conv_out_channels)
        if self.scale_conv_out_channels is not None and self.scale_conv_kernel_sizes is None:
            self.scale_conv_kernel_sizes = [self.scale_conv_kernel_size] * len(self.scale_conv_out_channels)
        super().__post_init__(**kwargs)

    def _init_scale_final_input_size(
        self, image_size: int | list[int] | tuple[int, int], patch_size: int | list[int] | tuple[int, int]
    ) -> None:
        if (
            self.scale_final_input_size is not None
            or self.scale_conv_out_channels is None
            or self.scale_conv_kernel_sizes is None
        ):
            return
        image_height, image_width = image_size if isinstance(image_size, (list, tuple)) else (image_size, image_size)
        patch_height = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_width = patch_size if isinstance(patch_size, int) else patch_size[1]
        features_height = image_height // patch_height
        features_width = image_width // patch_width
        for kernel_size in self.scale_conv_kernel_sizes:
            padding = (kernel_size - 1) // 2
            features_height = (features_height + 2 * padding - kernel_size) // 2 + 1
            features_width = (features_width + 2 * padding - kernel_size) // 2 + 1
        self.scale_final_input_size = features_height * features_width * self.scale_conv_out_channels[-1]


@auto_docstring(checkpoint="facebook/sapiens2-pretrain-0.4b")
@strict
class Sapiens2Config(BackboneConfigMixin, PreTrainedConfig):
    r"""
    rope_theta (`float`, *optional*, defaults to 100.0):
        The base period of the RoPE embeddings.
    query_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to the query projection.
    key_bias (`bool`, *optional*, defaults to `False`):
        Whether to add a bias to the key projection.
    value_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to the value projection.
    proj_bias (`bool`, *optional*, defaults to `True`):
        Whether to add a bias to the output projection.
    layerscale_value (`float`, *optional*, defaults to 1.0):
        Initial value to use for layer scale.
    use_gated_mlp (`bool`, *optional*, defaults to `False`):
        Whether to use the SwiGLU feedforward neural network.
    num_register_tokens (`int`, *optional*, defaults to 0):
        The number of register tokens.
    pos_embed_shift (`float`, *optional*):
        Amount to randomly shift position embedding coordinates in [-shift, shift],
        applied only in training mode if not `None`.
    pos_embed_jitter (`float`, *optional*):
        Amount to randomly jitter position embedding coordinates in log-uniform value in [1/jitter, jitter],
        applied only in training mode if not `None`.
    pos_embed_rescale (`float`, *optional*, defaults to 2.0):
        Amount to randomly rescale position embedding coordinates in log-uniform value in [1/rescale, rescale],
        applied only in training mode if not `None`.
    reshape_hidden_states (`bool`, *optional*, defaults to `True`):
        Whether to reshape the hidden states to spatial dimensions when used as backbone.
    use_mask_token (`bool`, *optional*, defaults to `False`):
        Whether to use a mask token in the embeddings (needed for masked image modeling pretraining).
    rms_norm_eps (`float`, *optional*, defaults to 1e-6):
        Epsilon for the RMS normalization layers.
    normalize_backbone_outputs (`bool`, *optional*, defaults to `True`):
        Whether to apply RMSNorm to the backbone `feature_maps` and `cls_tokens` outputs before
        returning them from the forward pass. Only applies when the model is used as a backbone.
    use_qk_norm (`bool`, *optional*, defaults to `True`):
        Whether to apply RMSNorm to queries and keys before RoPE in attention layers.
    num_key_value_heads_per_layer (`list[int]`, *optional*):
        Number of key/value heads for each transformer layer. Setting a layer's value equal to
        `num_attention_heads` gives full multi-head attention; a smaller value gives grouped-query
        attention. Defaults to `num_attention_heads` for the first `num_first_full_attention_layers`
        and last `num_last_full_attention_layers` layers and `num_key_valueattention_heads` for all other
        layers.
    num_key_value_attention_heads (`int`):
        Number of key/value heads for layers that use grouped-query attention when `num_key_value_heads_per_layer`
        is not set. Ignored when `num_key_value_heads_per_layer` is set.
    num_first_full_attention_layers (`int`, *optional*, defaults to 8):
        Number of leading transformer layers that use full multi-head attention.
        Only used when `num_key_value_heads_per_layer` is `None`.
    num_last_full_attention_layers (`int`, *optional*, defaults to 8):
        Number of trailing transformer layers that use full multi-head attention.
        Only used when `num_key_value_heads_per_layer` is `None`.
    semantic_loss_ignore_index (`int`, *optional*, defaults to 255):
        Label index ignored when computing the segmentation loss.
    flip_pairs (`list[list[int]]`, *optional*):
        Pairs of keypoint indices that are mirrored horizontally (e.g., left ear ↔ right ear).
        Each pair is a two-element list `[left_index, right_index]`. Used for test-time
        horizontal flip augmentation in pose estimation: pass these pairs to the second
        forward call so the model flips heatmaps back before returning them.
    head_config (`Sapiens2HeadConfig`, *optional*):
        Configuration for the decode head. See [`Sapiens2HeadConfig`] for the available options.
    """

    model_type = "sapiens2"

    patch_size: int | list[int] | tuple[int, int] = 16

    hidden_size: int = 1024
    intermediate_size: int = 4096
    num_hidden_layers: int = 24
    num_attention_heads: int = 16
    hidden_act: str = "silu"
    attention_dropout: float | int = 0.0
    initializer_range: float = 0.02
    rope_theta: float = 100.0
    image_size: int | list[int] | tuple[int, int] = 224
    num_channels: int = 3
    query_bias: bool = True
    key_bias: bool = True
    value_bias: bool = True
    proj_bias: bool = True
    mlp_bias: bool = True
    layerscale_value: float = 1.0
    drop_path_rate: float | int = 0.0
    use_gated_mlp: bool = True
    num_register_tokens: int = 8
    pos_embed_shift: float | None = None
    pos_embed_jitter: float | None = None
    pos_embed_rescale: float | None = 2.0
    _out_features: list[str] | None = None
    _out_indices: list[int] | None = None
    reshape_hidden_states: bool = True
    sub_configs = {"head_config": Sapiens2HeadConfig}
    use_mask_token: bool = False
    rms_norm_eps: float = 1e-6
    normalize_backbone_outputs: bool = True
    use_qk_norm: bool = True
    num_key_value_heads_per_layer: list[int] | None = None
    num_key_value_attention_heads: int = 8
    num_first_full_attention_layers: int = 8
    num_last_full_attention_layers: int = 8
    semantic_loss_ignore_index: int = 255
    flip_pairs: list[list[int]] | None = None
    head_config: Sapiens2HeadConfig | dict | None = None

    def __post_init__(self, **kwargs):
        if self.num_key_value_heads_per_layer is None:
            self.num_key_value_heads_per_layer = [
                self.num_attention_heads
                if (
                    layer_index < self.num_first_full_attention_layers
                    or layer_index >= self.num_hidden_layers - self.num_last_full_attention_layers
                )
                else self.num_key_value_attention_heads
                for layer_index in range(self.num_hidden_layers)
            ]
        if isinstance(self.head_config, dict):
            self.head_config = Sapiens2HeadConfig(**self.head_config)
        if self.head_config is not None:
            self.head_config._init_scale_final_input_size(image_size=self.image_size, patch_size=self.patch_size)
        self.stage_names = ["stem"] + [f"stage{i}" for i in range(1, self.num_hidden_layers + 1)]
        self.set_output_features_output_indices(
            out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
        )
        super().__post_init__(**kwargs)


__all__ = ["Sapiens2Config", "Sapiens2HeadConfig"]
