#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
#           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.

import math
from collections.abc import Callable, Iterable
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn.functional as F
from torch import Tensor, nn

from ... import initialization as init
from ...activations import ACT2FN
from ...backbone_utils import BackboneMixin, filter_output_hidden_states
from ...integrations import use_kernel_forward_from_hub
from ...modeling_layers import GradientCheckpointingLayer
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutput,
    BaseModelOutputWithPooling,
    ModelOutput,
    SemanticSegmenterOutput,
)
from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
from ...processing_utils import Unpack
from ...pytorch_utils import compile_compatible_method_lru_cache
from ...utils import TransformersKwargs, auto_docstring
from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults
from ...utils.output_capturing import capture_outputs
from .configuration_sapiens2 import Sapiens2Config


@auto_docstring(
    custom_intro="""
    Output type of [`Sapiens2Backbone`], extending [`BackboneOutput`] with optional CLS tokens from
    each selected feature stage (used when `config.return_class_token=True`).
    """
)
@dataclass
class Sapiens2BackboneOutput(BackboneOutput):
    r"""
    cls_tokens (`tuple(torch.FloatTensor)`, *optional*):
        CLS token from each selected feature stage, each of shape `(batch_size, hidden_size)`.
        Only present when `config.return_class_token=True`.
    """

    cls_tokens: tuple[torch.FloatTensor] | None = None


# General docstring


@auto_docstring(
    custom_intro="""
    Class for outputs of pose estimation models.
    """
)
@dataclass
class Sapiens2PoseEstimatorOutput(ModelOutput):
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Pose estimation loss.
    heatmaps (`torch.FloatTensor` of shape `(batch_size, num_keypoints, height, width)`):
        Heatmaps as predicted by the model.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
        one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states
        (also called feature maps) of the model at the output of each stage.
    """

    loss: torch.FloatTensor | None = None
    heatmaps: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None


@auto_docstring(
    custom_intro="""
    Class for outputs of normal estimation models.
    """
)
@dataclass
class Sapiens2NormalEstimatorOutput(ModelOutput):
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Normal estimation loss.
    normals (`torch.FloatTensor` of shape `(batch_size, num_labels, height, width)`):
        Raw normal map predictions as output by the model (unnormalized).
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage)
        of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of
        each layer plus the initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one per layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`. Attentions weights after the attention softmax.
    """

    loss: torch.FloatTensor | None = None
    normals: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None


@auto_docstring(
    custom_intro="""
    Class for outputs of pointmap estimation models.
    """
)
@dataclass
class Sapiens2PointmapEstimatorOutput(ModelOutput):
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Pointmap estimation loss.
    pointmaps (`torch.FloatTensor` of shape `(batch_size, 3, height, width)`):
        Per-pixel 3D XYZ coordinate predictions in canonical camera space.
    scales (`torch.FloatTensor` of shape `(batch_size, 1)`, *optional*):
        Canonical focal length / actual focal length ratio. `None` when no scale branch is configured.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for the output of each stage)
        of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of
        each layer plus the initial embedding outputs.
    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
        Tuple of `torch.FloatTensor` (one per layer) of shape `(batch_size, num_heads, sequence_length,
        sequence_length)`. Attentions weights after the attention softmax.
    """

    loss: torch.FloatTensor | None = None
    pointmaps: torch.FloatTensor | None = None
    scales: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor, ...] | None = None
    attentions: tuple[torch.FloatTensor, ...] | None = None


@auto_docstring(
    custom_intro="""
    Class for outputs of image matting models.
    """
)
@dataclass
class Sapiens2ImageMattingOutput(ModelOutput):
    r"""
    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
        Loss.
    alphas (`torch.FloatTensor` of shape `(batch_size, 1, height, width)`):
        Estimated alpha values.
    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
        one for the output of each stage) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states
        (also called feature maps) of the model at the output of each stage.
    foregrounds (`torch.FloatTensor` of shape `(batch_size, 3, height, width)`):
        Pre-multiplied RGB foreground predictions in `[0, 1]` (sigmoid-activated).
    """

    loss: torch.FloatTensor | None = None
    alphas: torch.FloatTensor | None = None
    hidden_states: tuple[torch.FloatTensor] | None = None
    attentions: tuple[torch.FloatTensor] | None = None

    foregrounds: torch.FloatTensor | None = None


class Sapiens2Embeddings(nn.Module):
    """
    Construct the CLS token, mask token, position and patch embeddings.
    """

    def __init__(self, config: Sapiens2Config):
        super().__init__()
        self.config = config
        self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
        self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) if config.use_mask_token else None
        self.register_tokens = nn.Parameter(torch.empty(1, config.num_register_tokens, config.hidden_size))
        self.patch_embeddings = nn.Conv2d(
            config.num_channels, config.hidden_size, kernel_size=config.patch_size, stride=config.patch_size
        )

    def forward(self, pixel_values: torch.Tensor, bool_masked_pos: torch.Tensor | None = None) -> torch.Tensor:
        if bool_masked_pos is not None and self.mask_token is None:
            raise ValueError("bool_masked_pos requires use_mask_token=True in the config")
        batch_size = pixel_values.shape[0]
        target_dtype = self.patch_embeddings.weight.dtype

        # (batch_size, num_channels, height, width) -> (batch_size, num_patches, hidden_size)
        patch_embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))
        patch_embeddings = patch_embeddings.flatten(2).transpose(1, 2)

        if bool_masked_pos is not None:
            mask_token = self.mask_token.to(patch_embeddings.dtype)
            patch_embeddings = torch.where(bool_masked_pos.unsqueeze(-1), mask_token, patch_embeddings)

        # Add CLS and register tokens
        cls_token = self.cls_token.expand(batch_size, -1, -1)
        register_tokens = self.register_tokens.expand(batch_size, -1, -1)
        embeddings = torch.cat([cls_token, register_tokens, patch_embeddings], dim=1)

        return embeddings


@compile_compatible_method_lru_cache(maxsize=32)
def get_patches_center_coordinates(
    num_patches_h: int, num_patches_w: int, dtype: torch.dtype, device: torch.device
) -> torch.Tensor:
    """
    Computes the 2D coordinates of the centers of image patches, normalized to the range [-1, +1].
    The center of each patch is exactly halfway between its top-left and bottom-right corners.

    Args:
        num_patches_h (int): Number of patches along the vertical (height) axis.
        num_patches_w (int): Number of patches along the horizontal (width) axis.
        dtype (torch.dtype): The desired data type of the returned tensor.

    Returns:
        torch.Tensor: A tensor of shape (height * width, 2), where each row contains the (y, x)
            coordinates of a patch center, normalized to [-1, +1].
    """
    coords_h = torch.arange(0.5, num_patches_h, dtype=dtype, device=device)
    coords_w = torch.arange(0.5, num_patches_w, dtype=dtype, device=device)
    coords_h = coords_h / num_patches_h
    coords_w = coords_w / num_patches_w
    # (height, width, 2) -> (height * width, 2)
    coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1)
    coords = coords.flatten(0, 1)
    # Shift range [0, 1] to [-1, +1]
    coords = 2.0 * coords - 1.0
    return coords


def augment_patches_center_coordinates(
    coords: torch.Tensor,
    shift: float | None = None,
    jitter: float | None = None,
    rescale: float | None = None,
) -> torch.Tensor:
    # Shift coords by adding a uniform value in [-shift, shift]
    if shift is not None:
        shift_hw = torch.empty((1, 2), device=coords.device, dtype=coords.dtype)
        shift_hw = shift_hw.uniform_(-shift, shift)
        coords = coords + shift_hw

    # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter]
    if jitter is not None:
        jitter_range = np.log(jitter)
        jitter_hw = torch.empty((1, 2), device=coords.device, dtype=coords.dtype)
        jitter_hw = jitter_hw.uniform_(-jitter_range, jitter_range).exp()
        coords = coords * jitter_hw

    # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale]
    if rescale is not None:
        rescale_range = np.log(rescale)
        rescale_hw = torch.empty(1, device=coords.device, dtype=coords.dtype)
        rescale_hw = rescale_hw.uniform_(-rescale_range, rescale_range).exp()
        coords = coords * rescale_hw

    return coords


class Sapiens2RopePositionEmbedding(nn.Module):
    inv_freq: torch.Tensor

    def __init__(self, config: Sapiens2Config):
        super().__init__()

        self.config = config
        self.base = config.rope_theta
        self.head_dim = config.hidden_size // config.num_attention_heads

        inv_freq = 1 / self.base ** torch.arange(0, 1, 4 / self.head_dim, dtype=torch.float32)  # (head_dim / 4,)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        image_size = config.image_size
        image_h, image_w = image_size if isinstance(image_size, Iterable) else (image_size, image_size)
        patch_size = config.patch_size
        patch_size_h = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_size_w = patch_size if isinstance(patch_size, int) else patch_size[1]
        self.num_patches_h = image_h // patch_size_h
        self.num_patches_w = image_w // patch_size_w

    def forward(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        _, _, height, width = pixel_values.shape
        num_patches_h = height // self.config.patch_size
        num_patches_w = width // self.config.patch_size

        device = pixel_values.device
        device_type = device.type if isinstance(device.type, str) and device.type != "mps" else "cpu"

        with maybe_autocast(device_type=device_type, enabled=False):  # Force float32
            # Although we could precompute static patch_coords from image_size and patch_size in the config,
            # the model was trained with random_scale, so it can process images of varying sizes.
            # Therefore, it's better to compute patch_coords dynamically (with lru_cache).
            patch_coords = get_patches_center_coordinates(
                num_patches_h, num_patches_w, dtype=torch.float32, device=device
            )
            if self.training:
                patch_coords = augment_patches_center_coordinates(
                    patch_coords,
                    shift=self.config.pos_embed_shift,
                    jitter=self.config.pos_embed_jitter,
                    rescale=self.config.pos_embed_rescale,
                )

            # (height * width, 2, head_dim / 4) -> (height * width, head_dim / 2) -> (height * width, head_dim)
            angles = 2 * math.pi * patch_coords[:, :, None] * self.inv_freq[None, None, :]
            angles = angles.flatten(1, 2)
            angles = angles.tile(2)

            cos = torch.cos(angles)
            sin = torch.sin(angles)

        dtype = pixel_values.dtype
        return cos.to(dtype=dtype), sin.to(dtype=dtype)


@use_kernel_forward_from_hub("RMSNorm")
class Sapiens2RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
        """
        Sapiens2RMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)

    def extra_repr(self):
        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"


def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    dropout: float | int = 0.0,
    scaling: float | None = None,
    softcap: float | None = None,
    **kwargs,
) -> tuple[torch.Tensor, torch.Tensor]:
    if scaling is None:
        scaling = module.head_dim**-0.5

    key_states = repeat_kv(key, module.num_key_value_groups)
    value_states = repeat_kv(value, module.num_key_value_groups)

    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling

    if softcap is not None:
        attn_weights = attn_weights / softcap
        attn_weights = torch.tanh(attn_weights)
        attn_weights = attn_weights * softcap
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask

    # upcast attention to fp32
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
    attn_output = torch.matmul(attn_weights, value_states)
    attn_output = attn_output.transpose(1, 2).contiguous()
    return attn_output, attn_weights


def apply_rotary_pos_emb(
    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, **kwargs
) -> tuple[torch.Tensor, torch.Tensor]:
    """Applies Rotary Position Embedding to the query and key tensors, but only to the patch tokens,
    ignoring the prefix tokens (cls token and register tokens).

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.

    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    """

    num_tokens = q.shape[-2]
    num_patches = sin.shape[-2]
    num_prefix_tokens = num_tokens - num_patches  # cls token + register tokens

    q_prefix_tokens, q_patches = q.split((num_prefix_tokens, num_patches), dim=-2)
    k_prefix_tokens, k_patches = k.split((num_prefix_tokens, num_patches), dim=-2)

    # apply rope only to patch tokens
    q_patches = (q_patches * cos) + (rotate_half(q_patches) * sin)
    k_patches = (k_patches * cos) + (rotate_half(k_patches) * sin)

    q = torch.cat((q_prefix_tokens, q_patches), dim=-2)
    k = torch.cat((k_prefix_tokens, k_patches), dim=-2)

    return q, k


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
    """
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    """
    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
    if n_rep == 1:
        return hidden_states
    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)


class Sapiens2Attention(nn.Module):
    """
    Multi-headed attention compatible with ALL_ATTENTION_FUNCTIONS.
    """

    def __init__(self, config: Sapiens2Config, layer_idx: int):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        self.is_causal = False

        self.scaling = self.head_dim**-0.5
        self.is_causal = False

        self.dropout = config.attention_dropout

        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.query_bias)
        self.o_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.proj_bias)
        self.num_key_value_heads = config.num_key_value_heads_per_layer[layer_idx]
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.k_proj = nn.Linear(self.embed_dim, self.num_key_value_heads * self.head_dim, bias=config.key_bias)
        self.v_proj = nn.Linear(self.embed_dim, self.num_key_value_heads * self.head_dim, bias=config.value_bias)
        self.q_norm = Sapiens2RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.use_qk_norm else nn.Identity()
        self.k_norm = Sapiens2RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.use_qk_norm else nn.Identity()

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Input shape: Batch x Time x Channel"""
        input_shape = hidden_states.shape[:-1]
        hidden_shape = (*input_shape, -1, self.head_dim)

        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
        query_states = self.q_norm(query_states)
        key_states = self.k_norm(key_states)

        cos, sin = position_embeddings
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
            self.config._attn_implementation, eager_attention_forward
        )
        attn_output, attn_weights = attention_interface(
            self,
            query_states,
            key_states,
            value_states,
            attention_mask,
            dropout=0.0 if not self.training else self.dropout,
            scaling=self.scaling,
            **kwargs,
        )

        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights


class Sapiens2LayerScale(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.lambda1 = nn.Parameter(config.layerscale_value * torch.ones(config.hidden_size))

    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
        return hidden_state * self.lambda1


class Sapiens2MLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.up_proj(x)))


class Sapiens2GatedMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return down_proj


class Sapiens2DropPath(nn.Module):
    """Stochastic depth (DropPath) per sample, for residual blocks.

    Identity when ``drop_prob`` is 0 or outside training. See `Deep Networks with Stochastic Depth
    <https://arxiv.org/abs/1603.09382>`_.
    """

    def __init__(self, drop_prob: float = 0.0) -> None:
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        if self.drop_prob == 0.0 or not self.training:
            return hidden_states
        keep_prob = 1 - self.drop_prob
        shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)
        random_tensor = torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)
        random_tensor = torch.floor(random_tensor + keep_prob)
        return hidden_states.div(keep_prob) * random_tensor

    def extra_repr(self) -> str:
        return f"p={self.drop_prob}"


class Sapiens2Layer(GradientCheckpointingLayer):
    """This corresponds to the Block class in the original implementation."""

    def __init__(self, config: Sapiens2Config, layer_idx: int):
        super().__init__()
        self.norm1 = Sapiens2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.attention = Sapiens2Attention(config, layer_idx=layer_idx)
        self.layer_scale1 = Sapiens2LayerScale(config)
        self.drop_path = Sapiens2DropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
        self.norm2 = Sapiens2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        if config.use_gated_mlp:
            self.mlp = Sapiens2GatedMLP(config)
        else:
            self.mlp = Sapiens2MLP(config)
        self.layer_scale2 = nn.Identity()

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> torch.Tensor:
        # Attention with residual connection
        residual = hidden_states
        hidden_states = self.norm1(hidden_states)
        hidden_states, _ = self.attention(
            hidden_states,
            attention_mask=attention_mask,
            position_embeddings=position_embeddings,
            **kwargs,
        )
        hidden_states = self.layer_scale1(hidden_states)
        hidden_states = self.drop_path(hidden_states) + residual

        # MLP with residual connection
        residual = hidden_states
        hidden_states = self.norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = self.layer_scale2(hidden_states)
        hidden_states = self.drop_path(hidden_states) + residual

        return hidden_states


class Sapiens2ConvLayer(nn.Module):
    """
    A basic wrapper for Convolution-BatchNorm-Activation, typically used for head components.
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int | tuple[int, int] = 1,
        stride: int = 1,
        padding: int | tuple[int, int] | str = 0,
        groups: int = 1,
        activation: str = "silu",
        bias: bool = True,
        convolution_transpose: bool = False,
        pixel_shuffle: bool = False,
        scale_factor: int = 2,
    ):
        super().__init__()
        if convolution_transpose:
            self.convolution = nn.ConvTranspose2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
            )
        else:
            self.convolution = nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                groups=groups,
                bias=bias,
            )
        self.norm = nn.InstanceNorm2d(out_channels)
        self.act_fn = ACT2FN[activation]
        if convolution_transpose:
            self.convolution = nn.ConvTranspose2d(
                in_channels,
                out_channels * scale_factor**2 if pixel_shuffle else out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias,
                groups=groups,
            )
        else:
            self.convolution = nn.Conv2d(
                in_channels,
                out_channels * scale_factor**2 if pixel_shuffle else out_channels,
                kernel_size=kernel_size,
                stride=stride,
                padding=padding,
                bias=bias,
                groups=groups,
            )
        self.pixel_shuffle = nn.PixelShuffle(scale_factor) if pixel_shuffle else nn.Identity()

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.convolution(hidden_states)
        hidden_states = self.pixel_shuffle(hidden_states)
        hidden_states = self.norm(hidden_states)
        hidden_states = self.act_fn(hidden_states)
        return hidden_states


class Sapiens2Head(nn.Module):
    def __init__(self, config: Sapiens2Config):
        super().__init__()
        self.input_conv = (
            Sapiens2ConvLayer(config.hidden_size, config.hidden_size, kernel_size=3, padding=1)
            if config.head_config.use_pixel_shuffle
            else nn.Identity()
        )
        upsample_in_channels = [config.hidden_size] + config.head_config.upsample_out_channels[:-1]
        self.upsample_layers = nn.ModuleList(
            Sapiens2ConvLayer(
                in_ch,
                out_ch,
                kernel_size=kernel_size,
                stride=1 if config.head_config.use_pixel_shuffle else 2,
                padding=(kernel_size - 1) // 2 if config.head_config.use_pixel_shuffle else 1,
                bias=bool(config.head_config.use_pixel_shuffle),
                pixel_shuffle=bool(config.head_config.use_pixel_shuffle),
                convolution_transpose=not config.head_config.use_pixel_shuffle,
            )
            for in_ch, out_ch, kernel_size in zip(
                upsample_in_channels,
                config.head_config.upsample_out_channels,
                config.head_config.upsample_kernel_sizes,
            )
        )
        conv_in_channels = [config.head_config.upsample_out_channels[-1]] + config.head_config.conv_out_channels[:-1]
        self.conv_layers = nn.ModuleList(
            Sapiens2ConvLayer(
                in_ch,
                out_ch,
                kernel_size=kernel_size,
                padding=(kernel_size - 1) // 2 if config.head_config.use_pixel_shuffle else 0,
            )
            for in_ch, out_ch, kernel_size in zip(
                conv_in_channels, config.head_config.conv_out_channels, config.head_config.conv_kernel_sizes
            )
        )
        predictor_in = (
            config.head_config.conv_out_channels[-1]
            if config.head_config.conv_out_channels
            else config.head_config.upsample_out_channels[-1]
            if config.head_config.upsample_out_channels
            else config.hidden_size
        )
        self.predictor = nn.Conv2d(predictor_in, config.num_labels, kernel_size=1)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.input_conv(hidden_states)
        for layer in self.upsample_layers:
            hidden_states = layer(hidden_states)
        for layer in self.conv_layers:
            hidden_states = layer(hidden_states)
        return self.predictor(hidden_states)


class Sapiens2PointmapFinalLayerBlock(nn.Module):
    def __init__(self, in_dim: int, out_dim: int, activation: nn.Module) -> None:
        super().__init__()
        self.layers = nn.ModuleList([nn.Linear(in_dim, out_dim), activation])

    def forward(self, input: Tensor) -> Tensor:
        hidden_state = input
        for layer in self.layers:
            hidden_state = layer(hidden_state)
        return hidden_state


class Sapiens2PointmapFinalLayer(nn.Module):
    def __init__(self, in_dim: int, hidden_sizes: tuple[int, int], out_dim: int = 1, activation: str = "silu"):
        super().__init__()
        self.flatten = nn.Flatten()
        self.block1 = Sapiens2PointmapFinalLayerBlock(
            in_dim=in_dim, out_dim=hidden_sizes[0], activation=ACT2FN[activation]
        )
        self.block2 = Sapiens2PointmapFinalLayerBlock(
            in_dim=hidden_sizes[0], out_dim=hidden_sizes[1], activation=ACT2FN[activation]
        )
        self.proj = nn.Linear(hidden_sizes[1], out_dim)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.flatten(hidden_states)
        hidden_states = self.block1(hidden_states)
        hidden_states = self.block2(hidden_states)
        return self.proj(hidden_states)


class Sapiens2PointmapScaleHead(nn.Module):
    def __init__(self, config: Sapiens2Config):
        super().__init__()
        self.conv_layers = nn.ModuleList()
        scale_in_channels = [config.hidden_size] + config.head_config.scale_conv_out_channels[:-1]
        for in_ch, out_ch, kernel_size in zip(
            scale_in_channels,
            config.head_config.scale_conv_out_channels,
            config.head_config.scale_conv_kernel_sizes,
        ):
            self.conv_layers.append(
                Sapiens2ConvLayer(in_ch, out_ch, kernel_size=kernel_size, stride=2, padding=(kernel_size - 1) // 2)
            )
        self.predictor = Sapiens2PointmapFinalLayer(
            config.head_config.scale_final_input_size,
            config.head_config.scale_final_hidden_sizes,
            activation=config.hidden_act,
        )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        for layer in self.conv_layers:
            hidden_states = layer(hidden_states)
        return self.predictor(hidden_states)


@auto_docstring
class Sapiens2PreTrainedModel(PreTrainedModel):
    config: Sapiens2Config
    base_model_prefix = "model"
    main_input_name = "pixel_values"
    input_modalities = ("image",)
    supports_gradient_checkpointing = True
    _no_split_modules = ["Sapiens2Layer"]
    _supports_sdpa = True
    _supports_flash_attn = True
    _supports_flex_attn = True
    _supports_attention_backend = True
    _can_record_outputs = {
        "hidden_states": Sapiens2Layer,
        "attentions": Sapiens2Attention,
    }

    # Ignore periods as we use inv_freq instead which is automatically calculated from the config.
    _keys_to_ignore_on_load_unexpected = [r"periods"]
    # mask_token is only used for masked image modeling pretraining and is absent in most checkpoints.
    _keys_to_ignore_on_load_missing = [r"mask_token"]

    @torch.no_grad()
    def _init_weights(self, module) -> None:
        """Initialize the weights"""
        super()._init_weights(module)
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            init.trunc_normal_(module.weight, mean=0.0, std=self.config.initializer_range)
        elif isinstance(module, nn.ConvTranspose2d):
            init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
        elif isinstance(module, Sapiens2Embeddings):
            init.trunc_normal_(module.cls_token, mean=0.0, std=self.config.initializer_range)
            if module.config.num_register_tokens > 0:
                init.trunc_normal_(module.register_tokens, mean=0.0, std=self.config.initializer_range)
            if module.config.use_mask_token:
                init.zeros_(module.mask_token)
        elif isinstance(module, Sapiens2LayerScale):
            init.constant_(module.lambda1, self.config.layerscale_value)
        elif isinstance(module, Sapiens2RopePositionEmbedding):
            inv_freq = 1 / module.base ** torch.arange(0, 1, 4 / module.head_dim, dtype=torch.float32)
            init.copy_(module.inv_freq, inv_freq)
        elif isinstance(module, (Sapiens2Head, Sapiens2PointmapScaleHead)):
            for head_module in module.modules():
                if isinstance(head_module, nn.Conv2d):
                    init.kaiming_normal_(head_module.weight, mode="fan_out", nonlinearity="relu")
                elif isinstance(head_module, nn.Linear):
                    init.kaiming_normal_(head_module.weight, mode="fan_in", nonlinearity="linear")


class Sapiens2Encoder(Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)
        self.layer = nn.ModuleList(
            [Sapiens2Layer(config, layer_idx=layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        # Initialize weights and apply final processing
        self.post_init()

    @merge_with_config_defaults
    @capture_outputs(tie_last_hidden_states=False)
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutput:
        for layer_module in self.layer:
            hidden_states = layer_module(hidden_states, position_embeddings=position_embeddings, **kwargs)

        return BaseModelOutput(last_hidden_state=hidden_states)


@auto_docstring
class Sapiens2Model(Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)
        self.embeddings = Sapiens2Embeddings(config)
        self.rope_embeddings = Sapiens2RopePositionEmbedding(config)
        self.model = Sapiens2Encoder(config)
        self.norm = Sapiens2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False
        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.patch_embeddings

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.Tensor,
        bool_masked_pos: torch.Tensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithPooling:
        r"""
        bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for
            pre-training.

        Example:

        ```python
        >>> from transformers import AutoImageProcessor, AutoModel
        >>> from transformers.image_utils import load_image
        >>> import torch

        >>> image = load_image("http://images.cocodataset.org/val2017/000000004016.jpg")
        >>> image_processor = AutoImageProcessor.from_pretrained("facebook/sapiens2-pretrain-0.4b")
        >>> model = AutoModel.from_pretrained("facebook/sapiens2-pretrain-0.4b")

        >>> inputs = image_processor(images=image, return_tensors="pt")
        >>> with torch.inference_mode():
        ...     outputs = model(**inputs)

        >>> cls_token = outputs.pooler_output
        >>> cls_token.shape
        torch.Size([1, 1024])
        ```
        """

        pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
        hidden_states = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
        position_embeddings = self.rope_embeddings(pixel_values)

        output = self.model(hidden_states, position_embeddings, **kwargs)
        sequence_output = self.norm(output.last_hidden_state)
        pooled_output = sequence_output[:, 0, :]

        return BaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=output.hidden_states,
            attentions=output.attentions,
        )


@auto_docstring
class Sapiens2Backbone(BackboneMixin, Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)

        self.embeddings = Sapiens2Embeddings(config)
        self.rope_embeddings = Sapiens2RopePositionEmbedding(config)
        self.model = Sapiens2Encoder(config)
        self.norm = Sapiens2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False

        self.num_features = [config.hidden_size for _ in range(config.num_hidden_layers + 1)]
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.patch_embeddings

    @can_return_tuple
    @filter_output_hidden_states
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.Tensor,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Sapiens2BackboneOutput:
        r"""
        Example:

        ```python
        >>> from transformers import AutoBackbone, AutoImageProcessor
        >>> from transformers.image_utils import load_image
        >>> import torch

        >>> image = load_image("http://images.cocodataset.org/val2017/000000004016.jpg")
        >>> image_processor = AutoImageProcessor.from_pretrained("facebook/sapiens2-pretrain-0.4b")
        >>> model = AutoBackbone.from_pretrained("facebook/sapiens2-pretrain-0.4b")

        >>> inputs = image_processor(images=image, return_tensors="pt")
        >>> with torch.inference_mode():
        ...     outputs = model(**inputs, return_class_token=True)

        >>> outputs.feature_maps[0].shape
        torch.Size([1, 1024, 64, 48])
        >>> outputs.cls_tokens[0].shape
        torch.Size([1, 1024])
        ```
        """
        pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
        hidden_states = self.embeddings(pixel_values)
        position_embeddings = self.rope_embeddings(pixel_values)

        kwargs["output_hidden_states"] = True  # required to extract layers for the stages
        output = self.model(hidden_states, position_embeddings, **kwargs)
        stage_hidden_states = output.hidden_states

        batch_size, _, image_height, image_width = pixel_values.shape
        patch_size = self.config.patch_size
        patch_size_h = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_size_w = patch_size if isinstance(patch_size, int) else patch_size[1]
        num_patches_height = image_height // patch_size_h
        num_patches_width = image_width // patch_size_w

        num_prefix = 1 + getattr(self.config, "num_register_tokens", 0)
        return_class_token = getattr(self.config, "return_class_token", False)

        feature_maps, cls_tokens = [], []
        for idx, (stage_name, hidden_state) in enumerate(zip(self.stage_names, stage_hidden_states)):
            if self.config.normalize_backbone_outputs:
                hidden_state = self.norm(hidden_state)

            if stage_name in self.out_features:
                if return_class_token:
                    cls_tokens.append(hidden_state[:, 0, :])
                patch_tokens = hidden_state[:, num_prefix:, :]
                if self.config.reshape_hidden_states:
                    feature_map = (
                        patch_tokens.reshape(batch_size, num_patches_height, num_patches_width, patch_tokens.shape[-1])
                        .permute(0, 3, 1, 2)
                        .contiguous()
                    )
                else:
                    feature_map = patch_tokens

                feature_maps.append(feature_map)

        return Sapiens2BackboneOutput(
            feature_maps=tuple(feature_maps),
            cls_tokens=tuple(cls_tokens) if return_class_token else None,
            hidden_states=output.hidden_states,
            attentions=output.attentions,
        )


@auto_docstring(checkpoint="facebook/sapiens2-seg-0.4b")
class Sapiens2ForSemanticSegmentation(Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = Sapiens2Model(config)
        self.decode_head = Sapiens2Head(config)
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        labels: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> SemanticSegmenterOutput:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, height, width)`, *optional*):
            Ground truth semantic segmentation maps for computing the loss.
            Indices should be in `[0, ..., config.num_labels - 1]`.
            If `config.num_labels > 1`, a classification loss is computed (Cross-Entropy).

        Example:

        ```python
        >>> from transformers import AutoImageProcessor, AutoModel
        >>> from transformers.image_utils import load_image
        >>> import torch

        >>> image = load_image("http://images.cocodataset.org/val2017/000000004016.jpg")
        >>> image_processor = AutoImageProcessor.from_pretrained("facebook/sapiens2-seg-0.4b")
        >>> model = AutoModel.from_pretrained("facebook/sapiens2-seg-0.4b")

        >>> inputs = image_processor(image, return_tensors="pt")
        >>> with torch.inference_mode():
        ...     outputs = model(**inputs)

        >>> outputs.logits.shape
        torch.Size([1, 29, 1024, 768])
        ```
        """
        if labels is not None and self.config.num_labels == 1:
            raise ValueError("The number of labels should be greater than one")

        outputs = self.model(pixel_values, **kwargs)

        batch_size, _, height, width = pixel_values.shape
        patch_size = self.config.patch_size
        patch_size_h = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_size_w = patch_size if isinstance(patch_size, int) else patch_size[1]
        patch_height = height // patch_size_h
        patch_width = width // patch_size_w

        patch_tokens = outputs.last_hidden_state[:, 1 + self.config.num_register_tokens :]
        feature_map = patch_tokens.transpose(1, 2).reshape(batch_size, -1, patch_height, patch_width)

        logits = self.decode_head(feature_map)

        loss = None
        if labels is not None:
            loss = self.loss_function(logits, labels, ignore_index=self.config.semantic_loss_ignore_index)

        return SemanticSegmenterOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


def flip_back(output_flipped, flip_pairs, target_type="gaussian-heatmap"):
    """Flip the flipped heatmaps back to the original form.

    Args:
        output_flipped (`torch.tensor` of shape `(batch_size, num_keypoints, height, width)`):
            The output heatmaps obtained from the flipped images.
        flip_pairs (`torch.Tensor` of shape `(num_keypoints, 2)`):
            Pairs of keypoints which are mirrored (for example, left ear -- right ear).
        target_type (`str`, *optional*, defaults to `"gaussian-heatmap"`):
            Target type to use. Can be gaussian-heatmap or combined-target.
            gaussian-heatmap: Classification target with gaussian distribution.
            combined-target: The combination of classification target (response map) and regression target (offset map).
            Paper ref: Huang et al. The Devil is in the Details: Delving into Unbiased Data Processing for Human Pose Estimation (CVPR 2020).

    Returns:
        torch.Tensor: heatmaps that flipped back to the original image
    """
    if target_type not in ["gaussian-heatmap", "combined-target"]:
        raise ValueError("target_type should be gaussian-heatmap or combined-target")

    if output_flipped.ndim != 4:
        raise ValueError("output_flipped should be [batch_size, num_keypoints, height, width]")

    batch_size, num_keypoints, height, width = output_flipped.shape
    channels = 1
    if target_type == "combined-target":
        channels = 3
        output_flipped = output_flipped.clone()  # clone to avoid mutation of output_flipped argument
        output_flipped[:, 1::3, ...] = -output_flipped[:, 1::3, ...]
    output_flipped = output_flipped.reshape(batch_size, -1, channels, height, width)
    output_flipped_back = output_flipped.clone()

    # Swap left-right parts
    left_indices, right_indices = flip_pairs.unbind(-1)
    output_flipped_back[:, left_indices, ...] = output_flipped[:, right_indices, ...]
    output_flipped_back[:, right_indices, ...] = output_flipped[:, left_indices, ...]
    output_flipped_back = output_flipped_back.reshape((batch_size, num_keypoints, height, width))
    # Flip horizontally
    output_flipped_back = output_flipped_back.flip(-1)
    return output_flipped_back


@auto_docstring(
    checkpoint="facebook/sapiens2-pose-0.4b",
    custom_intro="""
    The Sapiens2 model with a pose estimation head on top (a set of heatmap predictors on top of the hidden states output).
    """,
)
class Sapiens2ForPoseEstimation(Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = Sapiens2Model(config)
        self.decode_head = Sapiens2Head(config)
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        flip_pairs: torch.Tensor | None = None,
        labels: torch.FloatTensor | None = None,
        label_weights: torch.FloatTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Sapiens2PoseEstimatorOutput:
        r"""
        flip_pairs (`torch.Tensor` of shape `(num_pairs, 2)`, *optional*):
            Pairs of keypoints which are mirrored (for example, left ear -- right ear), used for
            test-time flip augmentation. When provided, the model assumes `pixel_values` contains
            horizontally-flipped images and calls `flip_back` on the output heatmaps to restore the
            original orientation.
        labels (`torch.FloatTensor` of shape `(batch_size, num_keypoints, height, width)`, *optional*):
            Heatmap ground truth for computing the loss.
        label_weights (`torch.FloatTensor` of shape `(batch_size, num_labels, 1, 1)` or `(batch_size, num_labels, height, width)`, *optional*):
            Visibility weights for each keypoint. Must be broadcastable to the shape of `labels`.

        Example:

        ```python
        >>> from transformers import AutoImageProcessor, AutoModel
        >>> from transformers.image_utils import load_image
        >>> import torch

        >>> image = load_image("http://images.cocodataset.org/val2017/000000004016.jpg")
        >>> image_processor = AutoImageProcessor.from_pretrained("facebook/sapiens2-pose-0.4b")
        >>> model = AutoModel.from_pretrained("facebook/sapiens2-pose-0.4b")

        >>> boxes = [[[270.8, 0.6, 294.1, 379.5]]]
        >>> inputs = image_processor(image, boxes=boxes, return_tensors="pt")
        >>> with torch.inference_mode():
        ...     outputs = model(**inputs)

        >>> outputs.heatmaps.shape
        torch.Size([1, 308, 256, 192])
        ```
        """
        outputs = self.model(pixel_values, **kwargs)

        batch_size, _, height, width = pixel_values.shape
        patch_size = self.config.patch_size
        patch_size_h = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_size_w = patch_size if isinstance(patch_size, int) else patch_size[1]
        patch_height = height // patch_size_h
        patch_width = width // patch_size_w

        patch_tokens = outputs.last_hidden_state[:, 1 + self.config.num_register_tokens :]
        feature_map = patch_tokens.transpose(1, 2).reshape(batch_size, -1, patch_height, patch_width)

        heatmaps = self.decode_head(feature_map)
        if flip_pairs is not None:
            heatmaps = flip_back(heatmaps, flip_pairs)

        loss = None
        if labels is not None:
            loss = F.mse_loss(heatmaps, labels, weight=label_weights)

        return Sapiens2PoseEstimatorOutput(
            loss=loss,
            heatmaps=heatmaps,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


@auto_docstring(
    checkpoint="facebook/sapiens2-normal-0.4b",
    custom_intro="""
    The Sapiens2 model with a normal estimation head on top (a PixelShuffle-based decoder that predicts surface normal maps).
    """,
)
class Sapiens2ForNormalEstimation(Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.model = Sapiens2Model(config)
        self.decode_head = Sapiens2Head(config)
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        labels: torch.FloatTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Sapiens2NormalEstimatorOutput:
        r"""
        labels (`torch.FloatTensor` of shape `(batch_size, num_labels, height, width)`, *optional*):
            Ground-truth surface normal maps for computing the loss.

        Example:

        ```python
        >>> from transformers import AutoImageProcessor, AutoModel
        >>> from transformers.image_utils import load_image
        >>> import torch

        >>> image = load_image("http://images.cocodataset.org/val2017/000000004016.jpg")
        >>> image_processor = AutoImageProcessor.from_pretrained("facebook/sapiens2-normal-0.4b")
        >>> model = AutoModel.from_pretrained("facebook/sapiens2-normal-0.4b")

        >>> inputs = image_processor(image, return_tensors="pt")
        >>> with torch.inference_mode():
        ...     outputs = model(**inputs)

        >>> outputs.normals.shape
        torch.Size([1, 3, 1024, 768])
        ```
        """
        outputs = self.model(pixel_values, **kwargs)

        batch_size, _, height, width = pixel_values.shape
        patch_size = self.config.patch_size
        patch_size_h = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_size_w = patch_size if isinstance(patch_size, int) else patch_size[1]
        patch_height = height // patch_size_h
        patch_width = width // patch_size_w

        patch_tokens = outputs.last_hidden_state[:, 1 + self.config.num_register_tokens :]
        feature_map = patch_tokens.transpose(1, 2).reshape(batch_size, -1, patch_height, patch_width)

        normals = self.decode_head(feature_map)

        loss = None
        if labels is not None:
            raise NotImplementedError("Training is not yet supported")

        return Sapiens2NormalEstimatorOutput(
            loss=loss,
            normals=normals,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


@auto_docstring(
    checkpoint="facebook/sapiens2-pointmap-0.4b",
    custom_intro="""
    The Sapiens2 model with a pointmap head on top (a PixelShuffle-based decoder that predicts per-pixel 3D XYZ
    coordinates, plus an optional scale branch for focal-length normalization).
    """,
)
class Sapiens2ForPointmapEstimation(Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)
        self.model = Sapiens2Model(config)
        self.decode_head = Sapiens2Head(config)
        self.scale_head = (
            Sapiens2PointmapScaleHead(config)
            if config.head_config is not None and config.head_config.scale_conv_out_channels is not None
            else nn.Identity()
        )
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        labels: torch.FloatTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Sapiens2PointmapEstimatorOutput:
        r"""
        labels (`torch.FloatTensor` of shape `(batch_size, 3, height, width)`, *optional*):
            Ground-truth pointmap for computing the loss.

        Example:

        ```python
        >>> from transformers import AutoImageProcessor, AutoModel
        >>> from transformers.image_utils import load_image
        >>> import torch

        >>> image = load_image("http://images.cocodataset.org/val2017/000000004016.jpg")
        >>> image_processor = AutoImageProcessor.from_pretrained("facebook/sapiens2-pointmap-0.4b")
        >>> model = AutoModel.from_pretrained("facebook/sapiens2-pointmap-0.4b")

        >>> inputs = image_processor(image, return_tensors="pt")
        >>> with torch.inference_mode():
        ...     outputs = model(**inputs)

        >>> outputs.pointmaps.shape
        torch.Size([1, 3, 1024, 768])
        ```
        """
        outputs = self.model(pixel_values, **kwargs)

        batch_size, _, height, width = pixel_values.shape
        patch_size = self.config.patch_size
        patch_size_h = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_size_w = patch_size if isinstance(patch_size, int) else patch_size[1]
        patch_height = height // patch_size_h
        patch_width = width // patch_size_w

        patch_tokens = outputs.last_hidden_state[:, 1 + self.config.num_register_tokens :]
        feature_map = patch_tokens.transpose(1, 2).reshape(batch_size, -1, patch_height, patch_width)

        pointmaps = self.decode_head(feature_map)
        scales = None if isinstance(self.scale_head, nn.Identity) else self.scale_head(feature_map)

        loss = None
        if labels is not None:
            raise NotImplementedError("Training is not yet supported")

        return Sapiens2PointmapEstimatorOutput(
            loss=loss,
            pointmaps=pointmaps,
            scales=scales,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


@auto_docstring(
    checkpoint="facebook/sapiens2-matting-1b",
    custom_intro="""
    The Sapiens2 model with a matting head on top (a PixelShuffle-based decoder that predicts a
    pre-multiplied RGB foreground and an alpha matte).
    """,
)
class Sapiens2ForImageMatting(Sapiens2PreTrainedModel):
    def __init__(self, config: Sapiens2Config):
        super().__init__(config)
        self.model = Sapiens2Model(config)
        self.decode_head = Sapiens2Head(config)  # config.num_labels = 4
        self.post_init()

    @can_return_tuple
    @auto_docstring
    def forward(
        self,
        pixel_values: torch.FloatTensor,
        labels: torch.FloatTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> Sapiens2ImageMattingOutput:
        r"""
        labels (`torch.FloatTensor` of shape `(batch_size, 4, height, width)`, *optional*):
            Ground-truth matting targets for computing the loss.

        Example:

        ```python
        >>> from transformers import AutoImageProcessor, AutoModel
        >>> from transformers.image_utils import load_image
        >>> import torch

        >>> image = load_image("http://images.cocodataset.org/val2017/000000004016.jpg")
        >>> image_processor = AutoImageProcessor.from_pretrained("facebook/sapiens2-matting-1b")
        >>> model = AutoModel.from_pretrained("facebook/sapiens2-matting-1b")

        >>> inputs = image_processor(image, return_tensors="pt")
        >>> with torch.inference_mode():
        ...     outputs = model(**inputs)

        >>> outputs.alphas.shape
        torch.Size([1, 1, 1024, 768])
        >>> outputs.foregrounds.shape
        torch.Size([1, 3, 1024, 768])
        ```
        """
        outputs = self.model(pixel_values, **kwargs)

        batch_size, _, height, width = pixel_values.shape
        patch_size = self.config.patch_size
        patch_size_h = patch_size if isinstance(patch_size, int) else patch_size[0]
        patch_size_w = patch_size if isinstance(patch_size, int) else patch_size[1]
        patch_height = height // patch_size_h
        patch_width = width // patch_size_w

        patch_tokens = outputs.last_hidden_state[:, 1 + self.config.num_register_tokens :]
        feature_map = patch_tokens.transpose(1, 2).reshape(batch_size, -1, patch_height, patch_width)

        matting = self.decode_head(feature_map).sigmoid()  # (B, 4, H, W)
        foregrounds = matting[:, :3]  # (B, 3, H, W)
        alphas = matting[:, 3:]  # (B, 1, H, W)

        loss = None
        if labels is not None:
            raise NotImplementedError("Training is not yet supported")

        return Sapiens2ImageMattingOutput(
            loss=loss,
            alphas=alphas,
            foregrounds=foregrounds,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


__all__ = [
    "Sapiens2ForSemanticSegmentation",
    "Sapiens2ForPoseEstimation",
    "Sapiens2ForNormalEstimation",
    "Sapiens2ForPointmapEstimation",
    "Sapiens2ForImageMatting",
    "Sapiens2Model",
    "Sapiens2PreTrainedModel",
    "Sapiens2Backbone",
]