ffgs 0.1.0.dev0__py3-none-any.whl

ffgs/models/anysplat/vggt/utils/pose_enc.py

@@ -0,0 +1,144 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Modified by the libgs contributors, 2026, from
+# src/model/encoder/vggt/utils/pose_enc.py of https://github.com/InternRobotics/AnySplat
+# (5f5e208): formatted with black, unused imports removed.
+# "The license found in the LICENSE file" above is the one VGGT's code
+# had when AnySplat vendored it (2025-06-30): CC BY-NC 4.0, non-commercial
+# (facebookresearch/vggt c6bf698); see LICENSE (this directory).
+
+import torch
+
+from .rotation import mat_to_quat, quat_to_mat
+
+
+def extri_intri_to_pose_encoding(
+    extrinsics,
+    intrinsics,
+    image_size_hw=None,  # e.g., (256, 512)
+    pose_encoding_type="absT_quaR_FoV",
+):
+    """Convert camera extrinsics and intrinsics to a compact pose encoding.
+
+    This function transforms camera parameters into a unified pose encoding format,
+    which can be used for various downstream tasks like pose prediction or representation.
+
+    Args:
+        extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4,
+            where B is batch size and S is sequence length.
+            In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world transformation.
+            The format is [R|t] where R is a 3x3 rotation matrix and t is a 3x1 translation vector.
+        intrinsics (torch.Tensor): Camera intrinsic parameters with shape BxSx3x3.
+            Defined in pixels, with format:
+            [[fx, 0, cx],
+             [0, fy, cy],
+             [0,  0,  1]]
+            where fx, fy are focal lengths and (cx, cy) is the principal point
+        image_size_hw (tuple): Tuple of (height, width) of the image in pixels.
+            Required for computing field of view values. For example: (256, 512).
+        pose_encoding_type (str): Type of pose encoding to use. Currently only
+            supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view).
+
+    Returns:
+        torch.Tensor: Encoded camera pose parameters with shape BxSx9.
+            For "absT_quaR_FoV" type, the 9 dimensions are:
+            - [:3] = absolute translation vector T (3D)
+            - [3:7] = rotation as quaternion quat (4D)
+            - [7:] = field of view (2D)
+    """
+
+    # extrinsics: BxSx3x4
+    # intrinsics: BxSx3x3
+
+    if pose_encoding_type == "absT_quaR_FoV":
+        R = extrinsics[:, :, :3, :3]  # BxSx3x3
+        T = extrinsics[:, :, :3, 3]  # BxSx3
+
+        quat = mat_to_quat(R)
+        # Note the order of h and w here
+        H, W = image_size_hw
+        fov_h = 2 * torch.atan((H / 2) / intrinsics[..., 1, 1])
+        fov_w = 2 * torch.atan((W / 2) / intrinsics[..., 0, 0])
+        pose_encoding = torch.cat(
+            [T, quat, fov_h[..., None], fov_w[..., None]], dim=-1
+        ).float()
+    else:
+        raise NotImplementedError
+
+    return pose_encoding
+
+
+def pose_encoding_to_extri_intri(
+    pose_encoding,
+    image_size_hw=None,  # e.g., (256, 512)
+    pose_encoding_type="absT_quaR_FoV",
+    build_intrinsics=True,
+):
+    """Convert a pose encoding back to camera extrinsics and intrinsics.
+
+    This function performs the inverse operation of extri_intri_to_pose_encoding,
+    reconstructing the full camera parameters from the compact encoding.
+
+    Args:
+        pose_encoding (torch.Tensor): Encoded camera pose parameters with shape BxSx9,
+            where B is batch size and S is sequence length.
+            For "absT_quaR_FoV" type, the 9 dimensions are:
+            - [:3] = absolute translation vector T (3D)
+            - [3:7] = rotation as quaternion quat (4D)
+            - [7:] = field of view (2D)
+        image_size_hw (tuple): Tuple of (height, width) of the image in pixels.
+            Required for reconstructing intrinsics from field of view values.
+            For example: (256, 512).
+        pose_encoding_type (str): Type of pose encoding used. Currently only
+            supports "absT_quaR_FoV" (absolute translation, quaternion rotation, field of view).
+        build_intrinsics (bool): Whether to reconstruct the intrinsics matrix.
+            If False, only extrinsics are returned and intrinsics will be None.
+
+    Returns:
+        tuple: (extrinsics, intrinsics)
+            - extrinsics (torch.Tensor): Camera extrinsic parameters with shape BxSx3x4.
+              In OpenCV coordinate system (x-right, y-down, z-forward), representing camera from world
+              transformation. The format is [R|t] where R is a 3x3 rotation matrix and t is
+              a 3x1 translation vector.
+            - intrinsics (torch.Tensor or None): Camera intrinsic parameters with shape BxSx3x3,
+              or None if build_intrinsics is False. Defined in pixels, with format:
+              [[fx, 0, cx],
+               [0, fy, cy],
+               [0,  0,  1]]
+              where fx, fy are focal lengths and (cx, cy) is the principal point,
+              assumed to be at the center of the image (W/2, H/2).
+    """
+
+    intrinsics = None
+
+    if pose_encoding_type == "absT_quaR_FoV":
+        T = pose_encoding[..., :3]
+        quat = pose_encoding[..., 3:7]
+        fov_h = pose_encoding[..., 7]
+        fov_w = pose_encoding[..., 8]
+
+        R = quat_to_mat(quat)
+        extrinsics = torch.cat([R, T[..., None]], dim=-1)
+
+        if build_intrinsics:
+            H, W = image_size_hw
+            fy = (H / 2.0) / (torch.tan(fov_h / 2.0) + 1e-3)
+            fx = (W / 2.0) / (torch.tan(fov_w / 2.0) + 1e-3)
+            intrinsics = torch.zeros(
+                pose_encoding.shape[:2] + (3, 3),
+                device=pose_encoding.device,
+                dtype=pose_encoding.dtype,
+            )
+            intrinsics[..., 0, 0] = fx
+            intrinsics[..., 1, 1] = fy
+            intrinsics[..., 0, 2] = W / 2
+            intrinsics[..., 1, 2] = H / 2
+            intrinsics[..., 2, 2] = 1.0  # Set the homogeneous coordinate to 1
+    else:
+        raise NotImplementedError
+
+    return extrinsics, intrinsics

ffgs/models/anysplat/vggt/utils/rotation.py

@@ -0,0 +1,148 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+# Modified by the libgs contributors, 2026, from
+# src/model/encoder/vggt/utils/rotation.py of https://github.com/InternRobotics/AnySplat
+# (5f5e208): formatted with black, unused imports removed.
+# "The license found in the LICENSE file" above is the one VGGT's code
+# had when AnySplat vendored it (2025-06-30): CC BY-NC 4.0, non-commercial
+# (facebookresearch/vggt c6bf698); see LICENSE (this directory).
+
+# Modified from PyTorch3D, https://github.com/facebookresearch/pytorch3d
+
+import torch
+import torch.nn.functional as F
+
+
+def quat_to_mat(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Quaternion Order: XYZW or say ijkr, scalar-last
+
+    Convert rotations given as quaternions to rotation matrices.
+    Args:
+        quaternions: quaternions with real part last,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    i, j, k, r = torch.unbind(quaternions, -1)
+    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def mat_to_quat(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part last, as tensor of shape (..., 4).
+        Quaternion Order: XYZW or say ijkr, scalar-last
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
+        matrix.reshape(batch_dim + (9,)), dim=-1
+    )
+
+    q_abs = _sqrt_positive_part(
+        torch.stack(
+            [
+                1.0 + m00 + m11 + m22,
+                1.0 + m00 - m11 - m22,
+                1.0 - m00 + m11 - m22,
+                1.0 - m00 - m11 + m22,
+            ],
+            dim=-1,
+        )
+    )
+
+    # we produce the desired quaternion multiplied by each of r, i, j, k
+    quat_by_rijk = torch.stack(
+        [
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], dim=-1),
+            # pyre-fixme[58]: `**` is not supported for operand types `Tensor` and
+            #  `int`.
+            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], dim=-1),
+        ],
+        dim=-2,
+    )
+
+    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+    # the candidate won't be picked.
+    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+    # forall i; we pick the best-conditioned one (with the largest denominator)
+    out = quat_candidates[
+        F.one_hot(q_abs.argmax(dim=-1), num_classes=4) > 0.5, :
+    ].reshape(batch_dim + (4,))
+
+    # Convert from rijk to ijkr
+    out = out[..., [1, 2, 3, 0]]
+
+    out = standardize_quaternion(out)
+
+    return out
+
+
+def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    if torch.is_grad_enabled():
+        ret[positive_mask] = torch.sqrt(x[positive_mask])
+    else:
+        ret = torch.where(positive_mask, torch.sqrt(x), ret)
+    return ret
+
+
+def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part last,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 3:4] < 0, -quaternions, quaternions)

ffgs/models/tokengs/__init__.py

@@ -0,0 +1,43 @@
+"""TokenGS (Ren, Tyszkiewicz, Huang, Gojcic; CVPR 2026), vendored from the official
+Apache-2.0 code at https://github.com/nv-tlabs/TokenGS (b16269c).
+
+`model_type` "tokengs". Imported on first use: `import ffgs` does not import it.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from ...registry import ModelSpec, register_model
+from .modeling import TokenGS
+from .processing import TokenGSProcessor, TokenGSProcessorConfig
+
+MODEL_TYPE = "tokengs"
+
+
+def convert_state_dict(checkpoint: dict[str, Any]) -> dict[str, Any]:
+    """An official checkpoint (the upstream module's state dict) -> `TokenGS`.
+
+    The parameter names are the upstream ones; only the LPIPS training loss that
+    upstream checkpoints may carry is dropped, as the upstream evaluation does.
+    """
+    return {k: v for k, v in checkpoint.items() if "lpips_loss" not in k}
+
+
+SPEC = register_model(
+    ModelSpec(
+        model_type=MODEL_TYPE,
+        model_cls=TokenGS,
+        processor_cls=TokenGSProcessor,
+        convert_state_dict=convert_state_dict,
+    )
+)
+
+__all__ = [
+    "MODEL_TYPE",
+    "SPEC",
+    "TokenGS",
+    "TokenGSProcessor",
+    "TokenGSProcessorConfig",
+    "convert_state_dict",
+]

ffgs/models/tokengs/activations.py

@@ -0,0 +1,135 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+#
+# Modified by the libgs contributors, 2026, from tokengs/models/activations.py of
+# https://github.com/nv-tlabs/TokenGS (b16269c): the clip head only, taking a
+# `TokenGSConfig`; einops replaced by the equivalent reshape / permute; the unused
+# ray arguments and weight-decay flags removed. Parameter names are unchanged.
+
+"""Gaussian tokens -> activated Gaussian parameters."""
+
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.init import trunc_normal_
+
+from .enc_dec import TokenGSConfig
+
+
+class ClipActivationHead(nn.Module):
+    """
+    Clip-based activation head that uses exponential activations with hard
+    clipping. Uses a single deconv (linear) layer for all Gaussian components.
+    """
+
+    def __init__(self, opt: TokenGSConfig):
+        super().__init__()
+        self.opt = opt
+
+        self.num_gaussians_per_token = self.opt.dec_patch_size**2
+
+        scale_cap = self.opt.gaussian_scale_cap
+        self.scale_shift = 1 - math.log(scale_cap)
+        self.scale_cap = scale_cap
+
+        self.output_dims = 3 + 1 + 3 + 4 + 3  # x, y, z, opacity, scale, rotation, rgb
+
+        # Single linear layer (deconv) for all components
+        self.deconv = nn.Linear(
+            self.opt.enc_embed_dim,
+            self.output_dims * self.opt.dec_patch_size * self.opt.dec_patch_size,
+            bias=True,
+        )
+
+        self._init_weights(self.deconv)
+        if self.opt.clip_head_z_init is not None:
+            bias_val = math.log(1.0 + self.opt.clip_head_z_init)
+            with torch.no_grad():
+                self.deconv.bias.view(
+                    self.opt.dec_patch_size * self.opt.dec_patch_size,
+                    self.output_dims,
+                )[:, 2].fill_(bias_val)
+
+    def _init_weights(self, m):
+        """Initialize deconv weights"""
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=self.opt.clip_head_readout_std)
+            if m.bias is not None:
+                m.bias.data.zero_()
+
+    def pos_act(self, x: torch.Tensor) -> torch.Tensor:
+        """Position activation without scaling factor (for xyz prediction)"""
+        pos = torch.sign(x) * (torch.expm1(torch.abs(x)))  # inverse log transform
+        return pos
+
+    def scale_act(self, x: torch.Tensor) -> torch.Tensor:
+        """Scale activation with exponential and hard clipping"""
+        x = x - self.scale_shift
+        return torch.minimum(
+            torch.exp(x),
+            torch.tensor(self.scale_cap, device=x.device, dtype=x.dtype),
+        )
+
+    def opacity_act(self, x: torch.Tensor) -> torch.Tensor:
+        """Opacity activation with hard clipping"""
+        return torch.sigmoid(x - self.opt.opacity_bias)
+
+    def rot_act(self, x: torch.Tensor) -> torch.Tensor:
+        """Rotation normalization"""
+        return F.normalize(x, dim=-1)
+
+    def rgb_act(self, x: torch.Tensor) -> torch.Tensor:
+        """RGB activation (always tanh for clip-based head)"""
+        return 0.5 * torch.tanh(x) + 0.5
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Convert tokens to Gaussians using a single deconv layer.
+
+        Args:
+            x: Input tokens [B, N, C]
+
+        Returns:
+            Gaussians [B, N * num_gaussians_per_token, 14]: xyz, opacity, scale,
+            rotation (wxyz), rgb
+        """
+        B = x.shape[0]
+
+        x = self.deconv(x)  # [B, N, output_dims * P * P]
+        # "b n (p c) -> b c (n p)"
+        n = x.shape[1]
+        p = self.opt.dec_patch_size**2
+        x = x.reshape(B, n, p, -1).permute(0, 3, 1, 2).reshape(B, -1, n * p)
+
+        x = x.reshape(B, self.output_dims, -1)  # B, output_dims, N * P * P
+        x = x.permute(0, 2, 1).contiguous()  # B, N * P * P, output_dims
+
+        pos, rgb, scaling, rotation, opacity = x.split([3, 3, 3, 4, 1], dim=-1)
+        pos = self.pos_act(pos)
+
+        opacity = self.opacity_act(opacity)
+        scale = self.scale_act(scaling)
+        rotation = self.rot_act(rotation)
+        rgbs = self.rgb_act(rgb)
+
+        gaussians = torch.cat(
+            [pos, opacity, scale, rotation, rgbs], dim=-1
+        )  # [B, N, 14]
+
+        return gaussians

ffgs/models/tokengs/attention.py

@@ -0,0 +1,189 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+#
+# Modified by the libgs contributors, 2026, from tokengs/models/attention.py of
+# https://github.com/nv-tlabs/TokenGS (b16269c): inference-only subset. Removed
+# FlexAttention (and the global torch._dynamo settings made on import), RoPE and
+# stochastic depth; parameter names are unchanged.
+
+"""Transformer building blocks of TokenGS."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: int | None = None,
+        out_features: int | None = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        bias: bool = True,
+    ) -> None:
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = True,
+        proj_bias: bool = True,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        qk_norm: bool = False,
+    ) -> None:
+        super().__init__()
+        assert dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
+        self.proj = nn.Linear(dim, dim, bias=proj_bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, self.head_dim)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv.unbind(0)
+        q, k = self.q_norm(q), self.k_norm(k)
+
+        x = F.scaled_dot_product_attention(q, k, v)
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        return x
+
+
+class LayerScale(nn.Module):
+    def __init__(self, dim: int, init_values: float | Tensor = 1e-5) -> None:
+        super().__init__()
+        self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+    def forward(self, x: Tensor) -> Tensor:
+        return x * self.gamma
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = True,
+        proj_bias: bool = True,
+        ffn_bias: bool = True,
+        init_values: float | Tensor | None = None,
+        act_layer: Callable[..., nn.Module] = nn.GELU,
+        norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+        qk_norm: bool = False,
+    ) -> None:
+        super().__init__()
+
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            proj_bias=proj_bias,
+            qk_norm=qk_norm,
+        )
+        self.ls1 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+
+        self.norm2 = norm_layer(dim)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=int(dim * mlp_ratio),
+            act_layer=act_layer,
+            bias=ffn_bias,
+        )
+        self.ls2 = (
+            LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.ls1(self.attn(self.norm1(x)))
+        x = x + self.ls2(self.mlp(self.norm2(x)))
+        return x
+
+
+def make_2tuple(x: int | tuple[int, int]) -> tuple[int, int]:
+    if isinstance(x, tuple):
+        assert len(x) == 2
+        return x
+    assert isinstance(x, int)
+    return (x, x)
+
+
+class PatchEmbed(nn.Module):
+    """2D image to patch embedding: (B, C, H, W) -> (B, N, D)."""
+
+    def __init__(
+        self,
+        patch_size: int | tuple[int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        norm_layer: Callable[..., nn.Module] | None = None,
+    ) -> None:
+        super().__init__()
+        patch_HW = make_2tuple(patch_size)
+        self.patch_size = patch_HW
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(
+            in_chans, embed_dim, kernel_size=patch_HW, stride=patch_HW
+        )
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        patch_H, patch_W = self.patch_size
+
+        assert (
+            H % patch_H == 0
+        ), f"Input image height {H} is not a multiple of patch height {patch_H}"
+        assert (
+            W % patch_W == 0
+        ), f"Input image width {W} is not a multiple of patch width: {patch_W}"
+
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)
+        x = self.norm(x)
+        return x