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.
- ffgs/__init__.py +39 -0
- ffgs/dynamic.py +239 -0
- ffgs/geometry.py +418 -0
- ffgs/image.py +204 -0
- ffgs/models/__init__.py +2 -0
- ffgs/models/anysplat/LICENSE +1120 -0
- ffgs/models/anysplat/__init__.py +49 -0
- ffgs/models/anysplat/encoder.py +372 -0
- ffgs/models/anysplat/evaluation.py +98 -0
- ffgs/models/anysplat/gaussian_adapter.py +96 -0
- ffgs/models/anysplat/gs_head.py +166 -0
- ffgs/models/anysplat/modeling.py +95 -0
- ffgs/models/anysplat/processing.py +249 -0
- ffgs/models/anysplat/vggt/__init__.py +0 -0
- ffgs/models/anysplat/vggt/heads/__init__.py +0 -0
- ffgs/models/anysplat/vggt/heads/camera_head.py +181 -0
- ffgs/models/anysplat/vggt/heads/dpt_head.py +612 -0
- ffgs/models/anysplat/vggt/heads/head_act.py +134 -0
- ffgs/models/anysplat/vggt/heads/utils.py +126 -0
- ffgs/models/anysplat/vggt/layers/__init__.py +28 -0
- ffgs/models/anysplat/vggt/layers/attention.py +92 -0
- ffgs/models/anysplat/vggt/layers/block.py +157 -0
- ffgs/models/anysplat/vggt/layers/drop_path.py +42 -0
- ffgs/models/anysplat/vggt/layers/layer_scale.py +32 -0
- ffgs/models/anysplat/vggt/layers/mlp.py +45 -0
- ffgs/models/anysplat/vggt/layers/patch_embed.py +106 -0
- ffgs/models/anysplat/vggt/layers/rope.py +211 -0
- ffgs/models/anysplat/vggt/layers/swiglu_ffn.py +77 -0
- ffgs/models/anysplat/vggt/layers/vision_transformer.py +463 -0
- ffgs/models/anysplat/vggt/models/__init__.py +0 -0
- ffgs/models/anysplat/vggt/models/aggregator.py +398 -0
- ffgs/models/anysplat/vggt/utils/__init__.py +0 -0
- ffgs/models/anysplat/vggt/utils/geometry.py +230 -0
- ffgs/models/anysplat/vggt/utils/pose_enc.py +144 -0
- ffgs/models/anysplat/vggt/utils/rotation.py +148 -0
- ffgs/models/tokengs/__init__.py +43 -0
- ffgs/models/tokengs/activations.py +135 -0
- ffgs/models/tokengs/attention.py +189 -0
- ffgs/models/tokengs/enc_dec.py +323 -0
- ffgs/models/tokengs/modeling.py +173 -0
- ffgs/models/tokengs/processing.py +339 -0
- ffgs/pipeline.py +341 -0
- ffgs/processor.py +236 -0
- ffgs/registry.py +132 -0
- ffgs/render.py +111 -0
- ffgs/types.py +291 -0
- ffgs/zoo/__init__.py +234 -0
- ffgs/zoo/anysplat/default.json +34 -0
- ffgs/zoo/tokengs/dl3dv-2v.json +29 -0
- ffgs/zoo/tokengs/dl3dv-4v.json +29 -0
- ffgs/zoo/tokengs/dl3dv-6v.json +29 -0
- ffgs/zoo/tokengs/dl3dv-base.json +29 -0
- ffgs/zoo/tokengs/dl3dv-latent-2v-lpips.json +48 -0
- ffgs/zoo/tokengs/dl3dv-latent-2v-ssim.json +48 -0
- ffgs/zoo/tokengs/dl3dv-latent-4v-lpips.json +48 -0
- ffgs/zoo/tokengs/dl3dv-latent-4v-ssim.json +48 -0
- ffgs/zoo/tokengs/dl3dv-latent-6v-lpips.json +48 -0
- ffgs/zoo/tokengs/dl3dv-latent-6v-ssim.json +48 -0
- ffgs-0.1.0.dev0.dist-info/METADATA +320 -0
- ffgs-0.1.0.dev0.dist-info/RECORD +63 -0
- ffgs-0.1.0.dev0.dist-info/WHEEL +4 -0
- ffgs-0.1.0.dev0.dist-info/licenses/LICENSE +237 -0
- 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
|