ffgs 0.1.0.dev0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (63) hide show
  1. ffgs/__init__.py +39 -0
  2. ffgs/dynamic.py +239 -0
  3. ffgs/geometry.py +418 -0
  4. ffgs/image.py +204 -0
  5. ffgs/models/__init__.py +2 -0
  6. ffgs/models/anysplat/LICENSE +1120 -0
  7. ffgs/models/anysplat/__init__.py +49 -0
  8. ffgs/models/anysplat/encoder.py +372 -0
  9. ffgs/models/anysplat/evaluation.py +98 -0
  10. ffgs/models/anysplat/gaussian_adapter.py +96 -0
  11. ffgs/models/anysplat/gs_head.py +166 -0
  12. ffgs/models/anysplat/modeling.py +95 -0
  13. ffgs/models/anysplat/processing.py +249 -0
  14. ffgs/models/anysplat/vggt/__init__.py +0 -0
  15. ffgs/models/anysplat/vggt/heads/__init__.py +0 -0
  16. ffgs/models/anysplat/vggt/heads/camera_head.py +181 -0
  17. ffgs/models/anysplat/vggt/heads/dpt_head.py +612 -0
  18. ffgs/models/anysplat/vggt/heads/head_act.py +134 -0
  19. ffgs/models/anysplat/vggt/heads/utils.py +126 -0
  20. ffgs/models/anysplat/vggt/layers/__init__.py +28 -0
  21. ffgs/models/anysplat/vggt/layers/attention.py +92 -0
  22. ffgs/models/anysplat/vggt/layers/block.py +157 -0
  23. ffgs/models/anysplat/vggt/layers/drop_path.py +42 -0
  24. ffgs/models/anysplat/vggt/layers/layer_scale.py +32 -0
  25. ffgs/models/anysplat/vggt/layers/mlp.py +45 -0
  26. ffgs/models/anysplat/vggt/layers/patch_embed.py +106 -0
  27. ffgs/models/anysplat/vggt/layers/rope.py +211 -0
  28. ffgs/models/anysplat/vggt/layers/swiglu_ffn.py +77 -0
  29. ffgs/models/anysplat/vggt/layers/vision_transformer.py +463 -0
  30. ffgs/models/anysplat/vggt/models/__init__.py +0 -0
  31. ffgs/models/anysplat/vggt/models/aggregator.py +398 -0
  32. ffgs/models/anysplat/vggt/utils/__init__.py +0 -0
  33. ffgs/models/anysplat/vggt/utils/geometry.py +230 -0
  34. ffgs/models/anysplat/vggt/utils/pose_enc.py +144 -0
  35. ffgs/models/anysplat/vggt/utils/rotation.py +148 -0
  36. ffgs/models/tokengs/__init__.py +43 -0
  37. ffgs/models/tokengs/activations.py +135 -0
  38. ffgs/models/tokengs/attention.py +189 -0
  39. ffgs/models/tokengs/enc_dec.py +323 -0
  40. ffgs/models/tokengs/modeling.py +173 -0
  41. ffgs/models/tokengs/processing.py +339 -0
  42. ffgs/pipeline.py +341 -0
  43. ffgs/processor.py +236 -0
  44. ffgs/registry.py +132 -0
  45. ffgs/render.py +111 -0
  46. ffgs/types.py +291 -0
  47. ffgs/zoo/__init__.py +234 -0
  48. ffgs/zoo/anysplat/default.json +34 -0
  49. ffgs/zoo/tokengs/dl3dv-2v.json +29 -0
  50. ffgs/zoo/tokengs/dl3dv-4v.json +29 -0
  51. ffgs/zoo/tokengs/dl3dv-6v.json +29 -0
  52. ffgs/zoo/tokengs/dl3dv-base.json +29 -0
  53. ffgs/zoo/tokengs/dl3dv-latent-2v-lpips.json +48 -0
  54. ffgs/zoo/tokengs/dl3dv-latent-2v-ssim.json +48 -0
  55. ffgs/zoo/tokengs/dl3dv-latent-4v-lpips.json +48 -0
  56. ffgs/zoo/tokengs/dl3dv-latent-4v-ssim.json +48 -0
  57. ffgs/zoo/tokengs/dl3dv-latent-6v-lpips.json +48 -0
  58. ffgs/zoo/tokengs/dl3dv-latent-6v-ssim.json +48 -0
  59. ffgs-0.1.0.dev0.dist-info/METADATA +320 -0
  60. ffgs-0.1.0.dev0.dist-info/RECORD +63 -0
  61. ffgs-0.1.0.dev0.dist-info/WHEEL +4 -0
  62. ffgs-0.1.0.dev0.dist-info/licenses/LICENSE +237 -0
  63. ffgs-0.1.0.dev0.dist-info/licenses/src/ffgs/models/anysplat/LICENSE +1120 -0
@@ -0,0 +1,134 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+ #
7
+ # Modified by the libgs contributors, 2026, from
8
+ # src/model/encoder/vggt/heads/head_act.py of https://github.com/InternRobotics/AnySplat
9
+ # (5f5e208): formatted with black, unused imports removed.
10
+ # "The license found in the LICENSE file" above is the one VGGT's code
11
+ # had when AnySplat vendored it (2025-06-30): CC BY-NC 4.0, non-commercial
12
+ # (facebookresearch/vggt c6bf698); see LICENSE (this directory).
13
+
14
+
15
+ import torch
16
+ import torch.nn.functional as F
17
+
18
+
19
+ def activate_pose(
20
+ pred_pose_enc, trans_act="linear", quat_act="linear", fl_act="linear"
21
+ ):
22
+ """
23
+ Activate pose parameters with specified activation functions.
24
+
25
+ Args:
26
+ pred_pose_enc: Tensor containing encoded pose parameters [translation, quaternion, focal length]
27
+ trans_act: Activation type for translation component
28
+ quat_act: Activation type for quaternion component
29
+ fl_act: Activation type for focal length component
30
+
31
+ Returns:
32
+ Activated pose parameters tensor
33
+ """
34
+ T = pred_pose_enc[..., :3]
35
+ quat = pred_pose_enc[..., 3:7]
36
+ fl = pred_pose_enc[..., 7:] # or fov
37
+
38
+ T = base_pose_act(T, trans_act)
39
+ quat = base_pose_act(quat, quat_act)
40
+ fl = base_pose_act(fl, fl_act) # or fov
41
+
42
+ pred_pose_enc = torch.cat([T, quat, fl], dim=-1)
43
+
44
+ return pred_pose_enc
45
+
46
+
47
+ def base_pose_act(pose_enc, act_type="linear"):
48
+ """
49
+ Apply basic activation function to pose parameters.
50
+
51
+ Args:
52
+ pose_enc: Tensor containing encoded pose parameters
53
+ act_type: Activation type ("linear", "inv_log", "exp", "relu")
54
+
55
+ Returns:
56
+ Activated pose parameters
57
+ """
58
+ if act_type == "linear":
59
+ return pose_enc
60
+ elif act_type == "inv_log":
61
+ return inverse_log_transform(pose_enc)
62
+ elif act_type == "exp":
63
+ return torch.exp(pose_enc)
64
+ elif act_type == "relu":
65
+ return F.relu(pose_enc)
66
+ else:
67
+ raise ValueError(f"Unknown act_type: {act_type}")
68
+
69
+
70
+ def activate_head(out, activation="norm_exp", conf_activation="expp1"):
71
+ """
72
+ Process network output to extract 3D points and confidence values.
73
+
74
+ Args:
75
+ out: Network output tensor (B, C, H, W)
76
+ activation: Activation type for 3D points
77
+ conf_activation: Activation type for confidence values
78
+
79
+ Returns:
80
+ Tuple of (3D points tensor, confidence tensor)
81
+ """
82
+ # Move channels from last dim to the 4th dimension => (B, H, W, C)
83
+ fmap = out.permute(0, 2, 3, 1) # B,H,W,C expected
84
+
85
+ # Split into xyz (first C-1 channels) and confidence (last channel)
86
+ xyz = fmap[:, :, :, :-1]
87
+ conf = fmap[:, :, :, -1]
88
+
89
+ if activation == "norm_exp":
90
+ d = xyz.norm(dim=-1, keepdim=True).clamp(min=1e-8)
91
+ xyz_normed = xyz / d
92
+ pts3d = xyz_normed * torch.expm1(d)
93
+ elif activation == "norm":
94
+ pts3d = xyz / xyz.norm(dim=-1, keepdim=True)
95
+ elif activation == "exp":
96
+ pts3d = torch.exp(xyz)
97
+ elif activation == "relu":
98
+ pts3d = F.relu(xyz)
99
+ elif activation == "inv_log":
100
+ pts3d = inverse_log_transform(xyz)
101
+ elif activation == "xy_inv_log":
102
+ xy, z = xyz.split([2, 1], dim=-1)
103
+ z = inverse_log_transform(z)
104
+ pts3d = torch.cat([xy * z, z], dim=-1)
105
+ elif activation == "sigmoid":
106
+ pts3d = torch.sigmoid(xyz)
107
+ elif activation == "linear":
108
+ pts3d = xyz
109
+ else:
110
+ raise ValueError(f"Unknown activation: {activation}")
111
+
112
+ if conf_activation == "expp1":
113
+ conf_out = 1 + conf.exp()
114
+ elif conf_activation == "expp0":
115
+ conf_out = conf.exp()
116
+ elif conf_activation == "sigmoid":
117
+ conf_out = torch.sigmoid(conf)
118
+ else:
119
+ raise ValueError(f"Unknown conf_activation: {conf_activation}")
120
+
121
+ return pts3d, conf_out
122
+
123
+
124
+ def inverse_log_transform(y):
125
+ """
126
+ Apply inverse log transform: sign(y) * (exp(|y|) - 1)
127
+
128
+ Args:
129
+ y: Input tensor
130
+
131
+ Returns:
132
+ Transformed tensor
133
+ """
134
+ return torch.sign(y) * (torch.expm1(torch.abs(y)))
@@ -0,0 +1,126 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+ #
7
+ # Modified by the libgs contributors, 2026, from src/model/encoder/vggt/heads/utils.py
8
+ # of https://github.com/InternRobotics/AnySplat (5f5e208): formatted with black, unused
9
+ # imports removed.
10
+ # "The license found in the LICENSE file" above is the one VGGT's code
11
+ # had when AnySplat vendored it (2025-06-30): CC BY-NC 4.0, non-commercial
12
+ # (facebookresearch/vggt c6bf698); see LICENSE (this directory).
13
+
14
+ import torch
15
+
16
+
17
+ def position_grid_to_embed(
18
+ pos_grid: torch.Tensor, embed_dim: int, omega_0: float = 100
19
+ ) -> torch.Tensor:
20
+ """
21
+ Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC)
22
+
23
+ Args:
24
+ pos_grid: Tensor of shape (H, W, 2) containing 2D coordinates
25
+ embed_dim: Output channel dimension for embeddings
26
+
27
+ Returns:
28
+ Tensor of shape (H, W, embed_dim) with positional embeddings
29
+ """
30
+ H, W, grid_dim = pos_grid.shape
31
+ assert grid_dim == 2
32
+ pos_flat = pos_grid.reshape(-1, grid_dim) # Flatten to (H*W, 2)
33
+
34
+ # Process x and y coordinates separately
35
+ emb_x = make_sincos_pos_embed(
36
+ embed_dim // 2, pos_flat[:, 0], omega_0=omega_0
37
+ ) # [1, H*W, D/2]
38
+ emb_y = make_sincos_pos_embed(
39
+ embed_dim // 2, pos_flat[:, 1], omega_0=omega_0
40
+ ) # [1, H*W, D/2]
41
+
42
+ # Combine and reshape
43
+ emb = torch.cat([emb_x, emb_y], dim=-1) # [1, H*W, D]
44
+
45
+ return emb.view(H, W, embed_dim) # [H, W, D]
46
+
47
+
48
+ def make_sincos_pos_embed(
49
+ embed_dim: int, pos: torch.Tensor, omega_0: float = 100
50
+ ) -> torch.Tensor:
51
+ """
52
+ This function generates a 1D positional embedding from a given grid using sine and cosine functions.
53
+
54
+ Args:
55
+ - embed_dim: The embedding dimension.
56
+ - pos: The position to generate the embedding from.
57
+
58
+ Returns:
59
+ - emb: The generated 1D positional embedding.
60
+ """
61
+ assert embed_dim % 2 == 0
62
+ omega = torch.arange(embed_dim // 2, dtype=torch.double, device=pos.device)
63
+ omega /= embed_dim / 2.0
64
+ omega = 1.0 / omega_0**omega # (D/2,)
65
+
66
+ pos = pos.reshape(-1) # (M,)
67
+ out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product
68
+
69
+ emb_sin = torch.sin(out) # (M, D/2)
70
+ emb_cos = torch.cos(out) # (M, D/2)
71
+
72
+ emb = torch.cat([emb_sin, emb_cos], dim=1) # (M, D)
73
+ return emb.float()
74
+
75
+
76
+ # Inspired by https://github.com/microsoft/moge
77
+
78
+
79
+ def create_uv_grid(
80
+ width: int,
81
+ height: int,
82
+ aspect_ratio: float = None,
83
+ dtype: torch.dtype = None,
84
+ device: torch.device = None,
85
+ ) -> torch.Tensor:
86
+ """
87
+ Create a normalized UV grid of shape (width, height, 2).
88
+
89
+ The grid spans horizontally and vertically according to an aspect ratio,
90
+ ensuring the top-left corner is at (-x_span, -y_span) and the bottom-right
91
+ corner is at (x_span, y_span), normalized by the diagonal of the plane.
92
+
93
+ Args:
94
+ width (int): Number of points horizontally.
95
+ height (int): Number of points vertically.
96
+ aspect_ratio (float, optional): Width-to-height ratio. Defaults to width/height.
97
+ dtype (torch.dtype, optional): Data type of the resulting tensor.
98
+ device (torch.device, optional): Device on which the tensor is created.
99
+
100
+ Returns:
101
+ torch.Tensor: A (width, height, 2) tensor of UV coordinates.
102
+ """
103
+ # Derive aspect ratio if not explicitly provided
104
+ if aspect_ratio is None:
105
+ aspect_ratio = float(width) / float(height)
106
+
107
+ # Compute normalized spans for X and Y
108
+ diag_factor = (aspect_ratio**2 + 1.0) ** 0.5
109
+ span_x = aspect_ratio / diag_factor
110
+ span_y = 1.0 / diag_factor
111
+
112
+ # Establish the linspace boundaries
113
+ left_x = -span_x * (width - 1) / width
114
+ right_x = span_x * (width - 1) / width
115
+ top_y = -span_y * (height - 1) / height
116
+ bottom_y = span_y * (height - 1) / height
117
+
118
+ # Generate 1D coordinates
119
+ x_coords = torch.linspace(left_x, right_x, steps=width, dtype=dtype, device=device)
120
+ y_coords = torch.linspace(top_y, bottom_y, steps=height, dtype=dtype, device=device)
121
+
122
+ # Create 2D meshgrid (width x height) and stack into UV
123
+ uu, vv = torch.meshgrid(x_coords, y_coords, indexing="xy")
124
+ uv_grid = torch.stack((uu, vv), dim=-1)
125
+
126
+ return uv_grid
@@ -0,0 +1,28 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ # All rights reserved.
3
+ #
4
+ # This source code is licensed under the license found in the
5
+ # LICENSE file in the root directory of this source tree.
6
+ #
7
+ # Modified by the libgs contributors, 2026, from
8
+ # src/model/encoder/vggt/layers/__init__.py of
9
+ # https://github.com/InternRobotics/AnySplat (5f5e208): `__all__` added (re-exports, not
10
+ # unused imports), imports sorted.
11
+ # "The license found in the LICENSE file" above is the one VGGT's code
12
+ # had when AnySplat vendored it (2025-06-30): CC BY-NC 4.0, non-commercial
13
+ # (facebookresearch/vggt c6bf698); see LICENSE (this directory).
14
+
15
+ from .attention import MemEffAttention
16
+ from .block import NestedTensorBlock
17
+ from .mlp import Mlp
18
+ from .patch_embed import PatchEmbed
19
+ from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
20
+
21
+ __all__ = [
22
+ "MemEffAttention",
23
+ "Mlp",
24
+ "NestedTensorBlock",
25
+ "PatchEmbed",
26
+ "SwiGLUFFN",
27
+ "SwiGLUFFNFused",
28
+ ]
@@ -0,0 +1,92 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ #
3
+ # This source code is licensed under the Apache License, Version 2.0
4
+ # found in the LICENSE file in the root directory of this source tree.
5
+ #
6
+ # Modified by the libgs contributors, 2026, from
7
+ # src/model/encoder/vggt/layers/attention.py of
8
+ # https://github.com/InternRobotics/AnySplat (5f5e208): the xFormers branch of
9
+ # `MemEffAttention` removed (upstream sets XFORMERS_AVAILABLE = False, so it never
10
+ # runs), formatted with black, unused imports removed.
11
+ # Apache License 2.0: see the Apache-2.0 text in LICENSE (this directory).
12
+
13
+ # References:
14
+ # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
15
+ # https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
16
+
17
+
18
+ import torch.nn.functional as F
19
+ from torch import Tensor, nn
20
+
21
+ XFORMERS_AVAILABLE = False
22
+
23
+
24
+ class Attention(nn.Module):
25
+ def __init__(
26
+ self,
27
+ dim: int,
28
+ num_heads: int = 8,
29
+ qkv_bias: bool = True,
30
+ proj_bias: bool = True,
31
+ attn_drop: float = 0.0,
32
+ proj_drop: float = 0.0,
33
+ norm_layer: nn.Module = nn.LayerNorm,
34
+ qk_norm: bool = False,
35
+ fused_attn: bool = True, # use F.scaled_dot_product_attention or not
36
+ rope=None,
37
+ ) -> None:
38
+ super().__init__()
39
+ assert dim % num_heads == 0, "dim should be divisible by num_heads"
40
+ self.num_heads = num_heads
41
+ self.head_dim = dim // num_heads
42
+ self.scale = self.head_dim**-0.5
43
+ self.fused_attn = fused_attn
44
+
45
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
46
+ self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
47
+ self.k_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
48
+ self.attn_drop = nn.Dropout(attn_drop)
49
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
50
+ self.proj_drop = nn.Dropout(proj_drop)
51
+ self.rope = rope
52
+
53
+ def forward(self, x: Tensor, pos=None) -> Tensor:
54
+ B, N, C = x.shape
55
+ qkv = (
56
+ self.qkv(x)
57
+ .reshape(B, N, 3, self.num_heads, self.head_dim)
58
+ .permute(2, 0, 3, 1, 4)
59
+ )
60
+ q, k, v = qkv.unbind(0)
61
+ q, k = self.q_norm(q), self.k_norm(k)
62
+
63
+ if self.rope is not None:
64
+ q = self.rope(q, pos)
65
+ k = self.rope(k, pos)
66
+
67
+ if self.fused_attn:
68
+ x = F.scaled_dot_product_attention(
69
+ q,
70
+ k,
71
+ v,
72
+ dropout_p=self.attn_drop.p if self.training else 0.0,
73
+ )
74
+ else:
75
+ q = q * self.scale
76
+ attn = q @ k.transpose(-2, -1)
77
+ attn = attn.softmax(dim=-1)
78
+ attn = self.attn_drop(attn)
79
+ x = attn @ v
80
+
81
+ x = x.transpose(1, 2).reshape(B, N, C)
82
+ x = self.proj(x)
83
+ x = self.proj_drop(x)
84
+ return x
85
+
86
+
87
+ class MemEffAttention(Attention):
88
+ def forward(self, x: Tensor, attn_bias=None, pos=None) -> Tensor:
89
+ assert pos is None
90
+ if attn_bias is not None:
91
+ raise AssertionError("xFormers is required for using nested tensors")
92
+ return super().forward(x)
@@ -0,0 +1,157 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ #
3
+ # This source code is licensed under the Apache License, Version 2.0
4
+ # found in the LICENSE file in the root directory of this source tree.
5
+ #
6
+ # Modified by the libgs contributors, 2026, from src/model/encoder/vggt/layers/block.py
7
+ # of https://github.com/InternRobotics/AnySplat (5f5e208): the xFormers nested-tensor
8
+ # code removed (`get_branges_scales`, `add_residual`, `get_attn_bias_and_cat`,
9
+ # `drop_add_residual_stochastic_depth_list`, `NestedTensorBlock.forward_nested`;
10
+ # unreachable with upstream's XFORMERS_AVAILABLE = False), formatted with black, unused
11
+ # imports removed.
12
+ # Apache License 2.0: see the Apache-2.0 text in LICENSE (this directory).
13
+
14
+ # References:
15
+ # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
16
+ # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
17
+
18
+ from typing import Callable
19
+
20
+ import torch
21
+ from torch import Tensor, nn
22
+
23
+ from .attention import Attention
24
+ from .drop_path import DropPath
25
+ from .layer_scale import LayerScale
26
+ from .mlp import Mlp
27
+
28
+ XFORMERS_AVAILABLE = False
29
+
30
+
31
+ class Block(nn.Module):
32
+ def __init__(
33
+ self,
34
+ dim: int,
35
+ num_heads: int,
36
+ mlp_ratio: float = 4.0,
37
+ qkv_bias: bool = True,
38
+ proj_bias: bool = True,
39
+ ffn_bias: bool = True,
40
+ drop: float = 0.0,
41
+ attn_drop: float = 0.0,
42
+ init_values=None,
43
+ drop_path: float = 0.0,
44
+ act_layer: Callable[..., nn.Module] = nn.GELU,
45
+ norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
46
+ attn_class: Callable[..., nn.Module] = Attention,
47
+ ffn_layer: Callable[..., nn.Module] = Mlp,
48
+ qk_norm: bool = False,
49
+ fused_attn: bool = True, # use F.scaled_dot_product_attention or not
50
+ rope=None,
51
+ ) -> None:
52
+ super().__init__()
53
+
54
+ self.norm1 = norm_layer(dim)
55
+
56
+ self.attn = attn_class(
57
+ dim,
58
+ num_heads=num_heads,
59
+ qkv_bias=qkv_bias,
60
+ proj_bias=proj_bias,
61
+ attn_drop=attn_drop,
62
+ proj_drop=drop,
63
+ qk_norm=qk_norm,
64
+ fused_attn=fused_attn,
65
+ rope=rope,
66
+ )
67
+
68
+ self.ls1 = (
69
+ LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
70
+ )
71
+ self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
72
+
73
+ self.norm2 = norm_layer(dim)
74
+ mlp_hidden_dim = int(dim * mlp_ratio)
75
+ self.mlp = ffn_layer(
76
+ in_features=dim,
77
+ hidden_features=mlp_hidden_dim,
78
+ act_layer=act_layer,
79
+ drop=drop,
80
+ bias=ffn_bias,
81
+ )
82
+ self.ls2 = (
83
+ LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
84
+ )
85
+ self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
86
+
87
+ self.sample_drop_ratio = drop_path
88
+
89
+ def forward(self, x: Tensor, pos=None) -> Tensor:
90
+ def attn_residual_func(x: Tensor, pos=None) -> Tensor:
91
+ return self.ls1(self.attn(self.norm1(x), pos=pos))
92
+
93
+ def ffn_residual_func(x: Tensor) -> Tensor:
94
+ return self.ls2(self.mlp(self.norm2(x)))
95
+
96
+ if self.training and self.sample_drop_ratio > 0.1:
97
+ # the overhead is compensated only for a drop path rate larger than 0.1
98
+ x = drop_add_residual_stochastic_depth(
99
+ x,
100
+ pos=pos,
101
+ residual_func=attn_residual_func,
102
+ sample_drop_ratio=self.sample_drop_ratio,
103
+ )
104
+ x = drop_add_residual_stochastic_depth(
105
+ x,
106
+ residual_func=ffn_residual_func,
107
+ sample_drop_ratio=self.sample_drop_ratio,
108
+ )
109
+ elif self.training and self.sample_drop_ratio > 0.0:
110
+ x = x + self.drop_path1(attn_residual_func(x, pos=pos))
111
+ x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
112
+ else:
113
+ x = x + attn_residual_func(x, pos=pos)
114
+ x = x + ffn_residual_func(x)
115
+ return x
116
+
117
+
118
+ def drop_add_residual_stochastic_depth(
119
+ x: Tensor,
120
+ residual_func: Callable[[Tensor], Tensor],
121
+ sample_drop_ratio: float = 0.0,
122
+ pos=None,
123
+ ) -> Tensor:
124
+ # 1) extract subset using permutation
125
+ b, n, d = x.shape
126
+ sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
127
+ brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
128
+ x_subset = x[brange]
129
+
130
+ # 2) apply residual_func to get residual
131
+ if pos is not None:
132
+ # if necessary, apply rope to the subset
133
+ pos = pos[brange]
134
+ residual = residual_func(x_subset, pos=pos)
135
+ else:
136
+ residual = residual_func(x_subset)
137
+
138
+ x_flat = x.flatten(1)
139
+ residual = residual.flatten(1)
140
+
141
+ residual_scale_factor = b / sample_subset_size
142
+
143
+ # 3) add the residual
144
+ x_plus_residual = torch.index_add(
145
+ x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor
146
+ )
147
+ return x_plus_residual.view_as(x)
148
+
149
+
150
+ class NestedTensorBlock(Block):
151
+ def forward(self, x_or_x_list):
152
+ if isinstance(x_or_x_list, Tensor):
153
+ return super().forward(x_or_x_list)
154
+ elif isinstance(x_or_x_list, list):
155
+ raise AssertionError("xFormers is required for using nested tensors")
156
+ else:
157
+ raise AssertionError
@@ -0,0 +1,42 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ #
3
+ # This source code is licensed under the Apache License, Version 2.0
4
+ # found in the LICENSE file in the root directory of this source tree.
5
+ #
6
+ # Modified by the libgs contributors, 2026, from
7
+ # src/model/encoder/vggt/layers/drop_path.py of
8
+ # https://github.com/InternRobotics/AnySplat (5f5e208): formatted with black, unused
9
+ # imports removed.
10
+ # Apache License 2.0: see the Apache-2.0 text in LICENSE (this directory).
11
+
12
+ # References:
13
+ # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
14
+ # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
15
+
16
+
17
+ from torch import nn
18
+
19
+
20
+ def drop_path(x, drop_prob: float = 0.0, training: bool = False):
21
+ if drop_prob == 0.0 or not training:
22
+ return x
23
+ keep_prob = 1 - drop_prob
24
+ shape = (x.shape[0],) + (1,) * (
25
+ x.ndim - 1
26
+ ) # work with diff dim tensors, not just 2D ConvNets
27
+ random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
28
+ if keep_prob > 0.0:
29
+ random_tensor.div_(keep_prob)
30
+ output = x * random_tensor
31
+ return output
32
+
33
+
34
+ class DropPath(nn.Module):
35
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
36
+
37
+ def __init__(self, drop_prob=None):
38
+ super(DropPath, self).__init__()
39
+ self.drop_prob = drop_prob
40
+
41
+ def forward(self, x):
42
+ return drop_path(x, self.drop_prob, self.training)
@@ -0,0 +1,32 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ #
3
+ # This source code is licensed under the Apache License, Version 2.0
4
+ # found in the LICENSE file in the root directory of this source tree.
5
+ #
6
+ # Modified by the libgs contributors, 2026, from
7
+ # src/model/encoder/vggt/layers/layer_scale.py of
8
+ # https://github.com/InternRobotics/AnySplat (5f5e208): formatted with black, unused
9
+ # imports removed.
10
+ # Apache License 2.0: see the Apache-2.0 text in LICENSE (this directory).
11
+
12
+ # Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
13
+
14
+ from typing import Union
15
+
16
+ import torch
17
+ from torch import Tensor, nn
18
+
19
+
20
+ class LayerScale(nn.Module):
21
+ def __init__(
22
+ self,
23
+ dim: int,
24
+ init_values: Union[float, Tensor] = 1e-5,
25
+ inplace: bool = False,
26
+ ) -> None:
27
+ super().__init__()
28
+ self.inplace = inplace
29
+ self.gamma = nn.Parameter(init_values * torch.ones(dim))
30
+
31
+ def forward(self, x: Tensor) -> Tensor:
32
+ return x.mul_(self.gamma) if self.inplace else x * self.gamma
@@ -0,0 +1,45 @@
1
+ # Copyright (c) Meta Platforms, Inc. and affiliates.
2
+ #
3
+ # This source code is licensed under the Apache License, Version 2.0
4
+ # found in the LICENSE file in the root directory of this source tree.
5
+ #
6
+ # Modified by the libgs contributors, 2026, from src/model/encoder/vggt/layers/mlp.py of
7
+ # https://github.com/InternRobotics/AnySplat (5f5e208): formatted with black, unused
8
+ # imports removed.
9
+ # Apache License 2.0: see the Apache-2.0 text in LICENSE (this directory).
10
+
11
+ # References:
12
+ # https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
13
+ # https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
14
+
15
+
16
+ from typing import Callable, Optional
17
+
18
+ from torch import Tensor, nn
19
+
20
+
21
+ class Mlp(nn.Module):
22
+ def __init__(
23
+ self,
24
+ in_features: int,
25
+ hidden_features: Optional[int] = None,
26
+ out_features: Optional[int] = None,
27
+ act_layer: Callable[..., nn.Module] = nn.GELU,
28
+ drop: float = 0.0,
29
+ bias: bool = True,
30
+ ) -> None:
31
+ super().__init__()
32
+ out_features = out_features or in_features
33
+ hidden_features = hidden_features or in_features
34
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
35
+ self.act = act_layer()
36
+ self.fc2 = nn.Linear(hidden_features, out_features, bias=bias)
37
+ self.drop = nn.Dropout(drop)
38
+
39
+ def forward(self, x: Tensor) -> Tensor:
40
+ x = self.fc1(x)
41
+ x = self.act(x)
42
+ x = self.drop(x)
43
+ x = self.fc2(x)
44
+ x = self.drop(x)
45
+ return x