diff --git a/Stable3DGen/__init__.py b/Stable3DGen/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/Stable3DGen/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..22902f5043a5587b1150dd12a94b6caa65e05a63
Binary files /dev/null and b/Stable3DGen/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/__init__.py b/Stable3DGen/hi3dgen/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..8fafab7776f0c7099302dd635ca2e673f0d35019
--- /dev/null
+++ b/Stable3DGen/hi3dgen/__init__.py
@@ -0,0 +1,36 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+#
+# Original file was released under MIT, with the full license text
+# available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+#
+# This modified file is released under the same license.
+from . import models
+from . import modules
+from . import pipelines
+from . import representations
diff --git a/Stable3DGen/hi3dgen/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1720249ecb5cf727be049abbedce04976e283cbc
Binary files /dev/null and b/Stable3DGen/hi3dgen/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/__init__.py b/Stable3DGen/hi3dgen/models/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..dadd77eb4dca3c4f23ab1814dd065fa5d7753ce5
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/__init__.py
@@ -0,0 +1,98 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+import importlib
+
+__attributes = {
+ 'SparseStructureEncoder': 'sparse_structure_vae',
+ 'SparseStructureDecoder': 'sparse_structure_vae',
+ 'SparseStructureFlowModel': 'sparse_structure_flow',
+ 'SLatEncoder': 'structured_latent_vae',
+ 'SLatGaussianDecoder': 'structured_latent_vae',
+ 'SLatRadianceFieldDecoder': 'structured_latent_vae',
+ 'SLatMeshDecoder': 'structured_latent_vae',
+ 'SLatFlowModel': 'structured_latent_flow',
+}
+
+__submodules = []
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+ if name not in globals():
+ if name in __attributes:
+ module_name = __attributes[name]
+ module = importlib.import_module(f".{module_name}", __name__)
+ globals()[name] = getattr(module, name)
+ elif name in __submodules:
+ module = importlib.import_module(f".{name}", __name__)
+ globals()[name] = module
+ else:
+ raise AttributeError(f"module {__name__} has no attribute {name}")
+ return globals()[name]
+
+
+def from_pretrained(path: str, **kwargs):
+ """
+ Load a model from a pretrained checkpoint.
+
+ Args:
+ path: The path to the checkpoint. Can be either local path or a Hugging Face model name.
+ NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively.
+ **kwargs: Additional arguments for the model constructor.
+ """
+ import os
+ import json
+ from safetensors.torch import load_file
+ is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors")
+
+ if is_local:
+ config_file = f"{path}.json"
+ model_file = f"{path}.safetensors"
+ else:
+ from huggingface_hub import hf_hub_download
+ path_parts = path.split('/')
+ repo_id = f'{path_parts[0]}/{path_parts[1]}'
+ model_name = '/'.join(path_parts[2:])
+ config_file = hf_hub_download(repo_id, f"{model_name}.json")
+ model_file = hf_hub_download(repo_id, f"{model_name}.safetensors")
+
+ with open(config_file, 'r') as f:
+ config = json.load(f)
+ model = __getattr__(config['name'])(**config['args'], **kwargs)
+ model.load_state_dict(load_file(model_file))
+
+ return model
+
+
+# For Pylance
+if __name__ == '__main__':
+ from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder
+ from .sparse_structure_flow import SparseStructureFlowModel
+ from .structured_latent_vae import SLatEncoder, SLatGaussianDecoder, SLatRadianceFieldDecoder, SLatMeshDecoder
+ from .structured_latent_flow import SLatFlowModel
diff --git a/Stable3DGen/hi3dgen/models/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/models/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a332823037be1fad7ed1a740ac43df04d107b2c1
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/__pycache__/sparse_structure_flow.cpython-310.pyc b/Stable3DGen/hi3dgen/models/__pycache__/sparse_structure_flow.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..98ec496bc9a319c5ca3e6f36bc775423258989d9
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/__pycache__/sparse_structure_flow.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/__pycache__/sparse_structure_vae.cpython-310.pyc b/Stable3DGen/hi3dgen/models/__pycache__/sparse_structure_vae.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4d3d842eda15cf51ea74c6127a6e9bd636ceba86
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/__pycache__/sparse_structure_vae.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/__pycache__/structured_latent_flow.cpython-310.pyc b/Stable3DGen/hi3dgen/models/__pycache__/structured_latent_flow.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b06e9681be1a090da3188cdd1e96528dec0dcf79
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/__pycache__/structured_latent_flow.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/sparse_structure_flow.py b/Stable3DGen/hi3dgen/models/sparse_structure_flow.py
new file mode 100755
index 0000000000000000000000000000000000000000..858fc4011c95b2dba5c2c392ce53ae1fdca5f3d9
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/sparse_structure_flow.py
@@ -0,0 +1,228 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import convert_module_to_f16, convert_module_to_f32
+from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock
+from ..modules.spatial import patchify, unpatchify
+
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+
+ Args:
+ t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ dim: the dimension of the output.
+ max_period: controls the minimum frequency of the embeddings.
+
+ Returns:
+ an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class SparseStructureFlowModel(nn.Module):
+ def __init__(
+ self,
+ resolution: int,
+ in_channels: int,
+ model_channels: int,
+ cond_channels: int,
+ out_channels: int,
+ num_blocks: int,
+ num_heads: Optional[int] = None,
+ num_head_channels: Optional[int] = 64,
+ mlp_ratio: float = 4,
+ patch_size: int = 2,
+ pe_mode: Literal["ape", "rope"] = "ape",
+ use_fp16: bool = False,
+ use_checkpoint: bool = False,
+ share_mod: bool = False,
+ qk_rms_norm: bool = False,
+ qk_rms_norm_cross: bool = False,
+ ):
+ super().__init__()
+ self.resolution = resolution
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.cond_channels = cond_channels
+ self.out_channels = out_channels
+ self.num_blocks = num_blocks
+ self.num_heads = num_heads or model_channels // num_head_channels
+ self.mlp_ratio = mlp_ratio
+ self.patch_size = patch_size
+ self.pe_mode = pe_mode
+ self.use_fp16 = use_fp16
+ self.use_checkpoint = use_checkpoint
+ self.share_mod = share_mod
+ self.qk_rms_norm = qk_rms_norm
+ self.qk_rms_norm_cross = qk_rms_norm_cross
+ self.dtype = torch.float16 if use_fp16 else torch.float32
+
+ self.t_embedder = TimestepEmbedder(model_channels)
+ if share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(model_channels, 6 * model_channels, bias=True)
+ )
+
+ if pe_mode == "ape":
+ pos_embedder = AbsolutePositionEmbedder(model_channels, 3)
+ coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution // patch_size] * 3], indexing='ij')
+ coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+ pos_emb = pos_embedder(coords)
+ self.register_buffer("pos_emb", pos_emb)
+
+ self.input_layer = nn.Linear(in_channels * patch_size**3, model_channels)
+
+ self.blocks = nn.ModuleList([
+ ModulatedTransformerCrossBlock(
+ model_channels,
+ cond_channels,
+ num_heads=self.num_heads,
+ mlp_ratio=self.mlp_ratio,
+ attn_mode='full',
+ use_checkpoint=self.use_checkpoint,
+ use_rope=(pe_mode == "rope"),
+ share_mod=share_mod,
+ qk_rms_norm=self.qk_rms_norm,
+ qk_rms_norm_cross=self.qk_rms_norm_cross,
+ )
+ for _ in range(num_blocks)
+ ])
+
+ self.out_layer = nn.Linear(model_channels, out_channels * patch_size**3)
+
+ self.initialize_weights()
+ if use_fp16:
+ self.convert_to_fp16()
+
+ @property
+ def device(self) -> torch.device:
+ """
+ Return the device of the model.
+ """
+ return next(self.parameters()).device
+
+ def convert_to_fp16(self) -> None:
+ """
+ Convert the torso of the model to float16.
+ """
+ self.blocks.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self) -> None:
+ """
+ Convert the torso of the model to float32.
+ """
+ self.blocks.apply(convert_module_to_f32)
+
+ def initialize_weights(self) -> None:
+ # Initialize transformer layers:
+ def _basic_init(module):
+ if isinstance(module, nn.Linear):
+ torch.nn.init.xavier_uniform_(module.weight)
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+ self.apply(_basic_init)
+
+ # Initialize timestep embedding MLP:
+ nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+ nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+ # Zero-out adaLN modulation layers in DiT blocks:
+ if self.share_mod:
+ nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+ nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+ else:
+ for block in self.blocks:
+ nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+ nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+ # Zero-out output layers:
+ nn.init.constant_(self.out_layer.weight, 0)
+ nn.init.constant_(self.out_layer.bias, 0)
+
+ def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+ assert [*x.shape] == [x.shape[0], self.in_channels, *[self.resolution] * 3], \
+ f"Input shape mismatch, got {x.shape}, expected {[x.shape[0], self.in_channels, *[self.resolution] * 3]}"
+
+ h = patchify(x, self.patch_size)
+ h = h.view(*h.shape[:2], -1).permute(0, 2, 1).contiguous()
+
+ h = self.input_layer(h)
+ h = h + self.pos_emb[None]
+ t_emb = self.t_embedder(t)
+ if self.share_mod:
+ t_emb = self.adaLN_modulation(t_emb)
+ t_emb = t_emb.type(self.dtype)
+ h = h.type(self.dtype)
+ cond = cond.type(self.dtype)
+ for block in self.blocks:
+ h = block(h, t_emb, cond)
+ h = h.type(x.dtype)
+ h = F.layer_norm(h, h.shape[-1:])
+ h = self.out_layer(h)
+
+ h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution // self.patch_size] * 3)
+ h = unpatchify(h, self.patch_size).contiguous()
+
+ return h
diff --git a/Stable3DGen/hi3dgen/models/sparse_structure_vae.py b/Stable3DGen/hi3dgen/models/sparse_structure_vae.py
new file mode 100755
index 0000000000000000000000000000000000000000..750d8234bef1b211a51439773642b18c1c990ec0
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/sparse_structure_vae.py
@@ -0,0 +1,334 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..modules.norm import GroupNorm32, ChannelLayerNorm32
+from ..modules.spatial import pixel_shuffle_3d
+from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
+
+
+def norm_layer(norm_type: str, *args, **kwargs) -> nn.Module:
+ """
+ Return a normalization layer.
+ """
+ if norm_type == "group":
+ return GroupNorm32(32, *args, **kwargs)
+ elif norm_type == "layer":
+ return ChannelLayerNorm32(*args, **kwargs)
+ else:
+ raise ValueError(f"Invalid norm type {norm_type}")
+
+
+class ResBlock3d(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ out_channels: Optional[int] = None,
+ norm_type: Literal["group", "layer"] = "layer",
+ ):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+
+ self.norm1 = norm_layer(norm_type, channels)
+ self.norm2 = norm_layer(norm_type, self.out_channels)
+ self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1)
+ self.conv2 = zero_module(nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1))
+ self.skip_connection = nn.Conv3d(channels, self.out_channels, 1) if channels != self.out_channels else nn.Identity()
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ h = self.norm1(x)
+ h = F.silu(h)
+ h = self.conv1(h)
+ h = self.norm2(h)
+ h = F.silu(h)
+ h = self.conv2(h)
+ h = h + self.skip_connection(x)
+ return h
+
+
+class DownsampleBlock3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ mode: Literal["conv", "avgpool"] = "conv",
+ ):
+ assert mode in ["conv", "avgpool"], f"Invalid mode {mode}"
+
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ if mode == "conv":
+ self.conv = nn.Conv3d(in_channels, out_channels, 2, stride=2)
+ elif mode == "avgpool":
+ assert in_channels == out_channels, "Pooling mode requires in_channels to be equal to out_channels"
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if hasattr(self, "conv"):
+ return self.conv(x)
+ else:
+ return F.avg_pool3d(x, 2)
+
+
+class UpsampleBlock3d(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ mode: Literal["conv", "nearest"] = "conv",
+ ):
+ assert mode in ["conv", "nearest"], f"Invalid mode {mode}"
+
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ if mode == "conv":
+ self.conv = nn.Conv3d(in_channels, out_channels*8, 3, padding=1)
+ elif mode == "nearest":
+ assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels"
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if hasattr(self, "conv"):
+ x = self.conv(x)
+ return pixel_shuffle_3d(x, 2)
+ else:
+ return F.interpolate(x, scale_factor=2, mode="nearest")
+
+
+class SparseStructureEncoder(nn.Module):
+ """
+ Encoder for Sparse Structure (\mathcal{E}_S in the paper Sec. 3.3).
+
+ Args:
+ in_channels (int): Channels of the input.
+ latent_channels (int): Channels of the latent representation.
+ num_res_blocks (int): Number of residual blocks at each resolution.
+ channels (List[int]): Channels of the encoder blocks.
+ num_res_blocks_middle (int): Number of residual blocks in the middle.
+ norm_type (Literal["group", "layer"]): Type of normalization layer.
+ use_fp16 (bool): Whether to use FP16.
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ latent_channels: int,
+ num_res_blocks: int,
+ channels: List[int],
+ num_res_blocks_middle: int = 2,
+ norm_type: Literal["group", "layer"] = "layer",
+ use_fp16: bool = False,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.latent_channels = latent_channels
+ self.num_res_blocks = num_res_blocks
+ self.channels = channels
+ self.num_res_blocks_middle = num_res_blocks_middle
+ self.norm_type = norm_type
+ self.use_fp16 = use_fp16
+ self.dtype = torch.float16 if use_fp16 else torch.float32
+
+ self.input_layer = nn.Conv3d(in_channels, channels[0], 3, padding=1)
+
+ self.blocks = nn.ModuleList([])
+ for i, ch in enumerate(channels):
+ self.blocks.extend([
+ ResBlock3d(ch, ch)
+ for _ in range(num_res_blocks)
+ ])
+ if i < len(channels) - 1:
+ self.blocks.append(
+ DownsampleBlock3d(ch, channels[i+1])
+ )
+
+ self.middle_block = nn.Sequential(*[
+ ResBlock3d(channels[-1], channels[-1])
+ for _ in range(num_res_blocks_middle)
+ ])
+
+ self.out_layer = nn.Sequential(
+ norm_layer(norm_type, channels[-1]),
+ nn.SiLU(),
+ nn.Conv3d(channels[-1], latent_channels*2, 3, padding=1)
+ )
+
+ if use_fp16:
+ self.convert_to_fp16()
+
+ @property
+ def device(self) -> torch.device:
+ """
+ Return the device of the model.
+ """
+ return next(self.parameters()).device
+
+ def convert_to_fp16(self) -> None:
+ """
+ Convert the torso of the model to float16.
+ """
+ self.use_fp16 = True
+ self.dtype = torch.float16
+ self.blocks.apply(convert_module_to_f16)
+ self.middle_block.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self) -> None:
+ """
+ Convert the torso of the model to float32.
+ """
+ self.use_fp16 = False
+ self.dtype = torch.float32
+ self.blocks.apply(convert_module_to_f32)
+ self.middle_block.apply(convert_module_to_f32)
+
+ def forward(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False) -> torch.Tensor:
+ h = self.input_layer(x)
+ h = h.type(self.dtype)
+
+ for block in self.blocks:
+ h = block(h)
+ h = self.middle_block(h)
+
+ h = h.type(x.dtype)
+ h = self.out_layer(h)
+
+ mean, logvar = h.chunk(2, dim=1)
+
+ if sample_posterior:
+ std = torch.exp(0.5 * logvar)
+ z = mean + std * torch.randn_like(std)
+ else:
+ z = mean
+
+ if return_raw:
+ return z, mean, logvar
+ return z
+
+
+class SparseStructureDecoder(nn.Module):
+ """
+ Decoder for Sparse Structure (\mathcal{D}_S in the paper Sec. 3.3).
+
+ Args:
+ out_channels (int): Channels of the output.
+ latent_channels (int): Channels of the latent representation.
+ num_res_blocks (int): Number of residual blocks at each resolution.
+ channels (List[int]): Channels of the decoder blocks.
+ num_res_blocks_middle (int): Number of residual blocks in the middle.
+ norm_type (Literal["group", "layer"]): Type of normalization layer.
+ use_fp16 (bool): Whether to use FP16.
+ """
+ def __init__(
+ self,
+ out_channels: int,
+ latent_channels: int,
+ num_res_blocks: int,
+ channels: List[int],
+ num_res_blocks_middle: int = 2,
+ norm_type: Literal["group", "layer"] = "layer",
+ use_fp16: bool = False,
+ ):
+ super().__init__()
+ self.out_channels = out_channels
+ self.latent_channels = latent_channels
+ self.num_res_blocks = num_res_blocks
+ self.channels = channels
+ self.num_res_blocks_middle = num_res_blocks_middle
+ self.norm_type = norm_type
+ self.use_fp16 = use_fp16
+ self.dtype = torch.float16 if use_fp16 else torch.float32
+
+ self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1)
+
+ self.middle_block = nn.Sequential(*[
+ ResBlock3d(channels[0], channels[0])
+ for _ in range(num_res_blocks_middle)
+ ])
+
+ self.blocks = nn.ModuleList([])
+ for i, ch in enumerate(channels):
+ self.blocks.extend([
+ ResBlock3d(ch, ch)
+ for _ in range(num_res_blocks)
+ ])
+ if i < len(channels) - 1:
+ self.blocks.append(
+ UpsampleBlock3d(ch, channels[i+1])
+ )
+
+ self.out_layer = nn.Sequential(
+ norm_layer(norm_type, channels[-1]),
+ nn.SiLU(),
+ nn.Conv3d(channels[-1], out_channels, 3, padding=1)
+ )
+
+ if use_fp16:
+ self.convert_to_fp16()
+
+ @property
+ def device(self) -> torch.device:
+ """
+ Return the device of the model.
+ """
+ return next(self.parameters()).device
+
+ def convert_to_fp16(self) -> None:
+ """
+ Convert the torso of the model to float16.
+ """
+ self.use_fp16 = True
+ self.dtype = torch.float16
+ self.blocks.apply(convert_module_to_f16)
+ self.middle_block.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self) -> None:
+ """
+ Convert the torso of the model to float32.
+ """
+ self.use_fp16 = False
+ self.dtype = torch.float32
+ self.blocks.apply(convert_module_to_f32)
+ self.middle_block.apply(convert_module_to_f32)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ h = self.input_layer(x)
+
+ h = h.type(self.dtype)
+
+ h = self.middle_block(h)
+ for block in self.blocks:
+ h = block(h)
+
+ h = h.type(x.dtype)
+ h = self.out_layer(h)
+ return h
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_flow.py b/Stable3DGen/hi3dgen/models/structured_latent_flow.py
new file mode 100755
index 0000000000000000000000000000000000000000..a1d9d8e3ca75a34242d03aa7f91d5f692a7845fb
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/structured_latent_flow.py
@@ -0,0 +1,290 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
+from ..modules.transformer import AbsolutePositionEmbedder
+from ..modules.norm import LayerNorm32
+from ..modules import sparse as sp
+from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock
+from .sparse_structure_flow import TimestepEmbedder
+
+
+class SparseResBlock3d(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ emb_channels: int,
+ out_channels: Optional[int] = None,
+ downsample: bool = False,
+ upsample: bool = False,
+ ):
+ super().__init__()
+ self.channels = channels
+ self.emb_channels = emb_channels
+ self.out_channels = out_channels or channels
+ self.downsample = downsample
+ self.upsample = upsample
+
+ assert not (downsample and upsample), "Cannot downsample and upsample at the same time"
+
+ self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+ self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+ self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
+ self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+ self.emb_layers = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(emb_channels, 2 * self.out_channels, bias=True),
+ )
+ self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+ self.updown = None
+ if self.downsample:
+ self.updown = sp.SparseDownsample(2)
+ elif self.upsample:
+ self.updown = sp.SparseUpsample(2)
+
+ def _updown(self, x: sp.SparseTensor) -> sp.SparseTensor:
+ if self.updown is not None:
+ x = self.updown(x)
+ return x
+
+ def forward(self, x: sp.SparseTensor, emb: torch.Tensor) -> sp.SparseTensor:
+ emb_out = self.emb_layers(emb).type(x.dtype)
+ scale, shift = torch.chunk(emb_out, 2, dim=1)
+
+ x = self._updown(x)
+ h = x.replace(self.norm1(x.feats))
+ h = h.replace(F.silu(h.feats))
+ h = self.conv1(h)
+ h = h.replace(self.norm2(h.feats)) * (1 + scale) + shift
+ h = h.replace(F.silu(h.feats))
+ h = self.conv2(h)
+ h = h + self.skip_connection(x)
+
+ return h
+
+
+class SLatFlowModel(nn.Module):
+ def __init__(
+ self,
+ resolution: int,
+ in_channels: int,
+ model_channels: int,
+ cond_channels: int,
+ out_channels: int,
+ num_blocks: int,
+ num_heads: Optional[int] = None,
+ num_head_channels: Optional[int] = 64,
+ mlp_ratio: float = 4,
+ patch_size: int = 2,
+ num_io_res_blocks: int = 2,
+ io_block_channels: List[int] = None,
+ pe_mode: Literal["ape", "rope"] = "ape",
+ use_fp16: bool = False,
+ use_checkpoint: bool = False,
+ use_skip_connection: bool = True,
+ share_mod: bool = False,
+ qk_rms_norm: bool = False,
+ qk_rms_norm_cross: bool = False,
+ ):
+ super().__init__()
+ self.resolution = resolution
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.cond_channels = cond_channels
+ self.out_channels = out_channels
+ self.num_blocks = num_blocks
+ self.num_heads = num_heads or model_channels // num_head_channels
+ self.mlp_ratio = mlp_ratio
+ self.patch_size = patch_size
+ self.num_io_res_blocks = num_io_res_blocks
+ self.io_block_channels = io_block_channels
+ self.pe_mode = pe_mode
+ self.use_fp16 = use_fp16
+ self.use_checkpoint = use_checkpoint
+ self.use_skip_connection = use_skip_connection
+ self.share_mod = share_mod
+ self.qk_rms_norm = qk_rms_norm
+ self.qk_rms_norm_cross = qk_rms_norm_cross
+ self.dtype = torch.float16 if use_fp16 else torch.float32
+
+ assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2"
+ assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages"
+
+ self.t_embedder = TimestepEmbedder(model_channels)
+ if share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(model_channels, 6 * model_channels, bias=True)
+ )
+
+ if pe_mode == "ape":
+ self.pos_embedder = AbsolutePositionEmbedder(model_channels)
+
+ self.input_layer = sp.SparseLinear(in_channels, io_block_channels[0])
+ self.input_blocks = nn.ModuleList([])
+ for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]):
+ self.input_blocks.extend([
+ SparseResBlock3d(
+ chs,
+ model_channels,
+ out_channels=chs,
+ )
+ for _ in range(num_io_res_blocks-1)
+ ])
+ self.input_blocks.append(
+ SparseResBlock3d(
+ chs,
+ model_channels,
+ out_channels=next_chs,
+ downsample=True,
+ )
+ )
+
+ self.blocks = nn.ModuleList([
+ ModulatedSparseTransformerCrossBlock(
+ model_channels,
+ cond_channels,
+ num_heads=self.num_heads,
+ mlp_ratio=self.mlp_ratio,
+ attn_mode='full',
+ use_checkpoint=self.use_checkpoint,
+ use_rope=(pe_mode == "rope"),
+ share_mod=self.share_mod,
+ qk_rms_norm=self.qk_rms_norm,
+ qk_rms_norm_cross=self.qk_rms_norm_cross,
+ )
+ for _ in range(num_blocks)
+ ])
+
+ self.out_blocks = nn.ModuleList([])
+ for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))):
+ self.out_blocks.append(
+ SparseResBlock3d(
+ prev_chs * 2 if self.use_skip_connection else prev_chs,
+ model_channels,
+ out_channels=chs,
+ upsample=True,
+ )
+ )
+ self.out_blocks.extend([
+ SparseResBlock3d(
+ chs * 2 if self.use_skip_connection else chs,
+ model_channels,
+ out_channels=chs,
+ )
+ for _ in range(num_io_res_blocks-1)
+ ])
+ self.out_layer = sp.SparseLinear(io_block_channels[0], out_channels)
+
+ self.initialize_weights()
+ if use_fp16:
+ self.convert_to_fp16()
+
+ @property
+ def device(self) -> torch.device:
+ """
+ Return the device of the model.
+ """
+ return next(self.parameters()).device
+
+ def convert_to_fp16(self) -> None:
+ """
+ Convert the torso of the model to float16.
+ """
+ self.input_blocks.apply(convert_module_to_f16)
+ self.blocks.apply(convert_module_to_f16)
+ self.out_blocks.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self) -> None:
+ """
+ Convert the torso of the model to float32.
+ """
+ self.input_blocks.apply(convert_module_to_f32)
+ self.blocks.apply(convert_module_to_f32)
+ self.out_blocks.apply(convert_module_to_f32)
+
+ def initialize_weights(self) -> None:
+ # Initialize transformer layers:
+ def _basic_init(module):
+ if isinstance(module, nn.Linear):
+ torch.nn.init.xavier_uniform_(module.weight)
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+ self.apply(_basic_init)
+
+ # Initialize timestep embedding MLP:
+ nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+ nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+ # Zero-out adaLN modulation layers in DiT blocks:
+ if self.share_mod:
+ nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+ nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+ else:
+ for block in self.blocks:
+ nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+ nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+ # Zero-out output layers:
+ nn.init.constant_(self.out_layer.weight, 0)
+ nn.init.constant_(self.out_layer.bias, 0)
+
+ def forward(self, x: sp.SparseTensor, t: torch.Tensor, cond: torch.Tensor) -> sp.SparseTensor:
+ h = self.input_layer(x).type(self.dtype)
+ t_emb = self.t_embedder(t)
+ if self.share_mod:
+ t_emb = self.adaLN_modulation(t_emb)
+ t_emb = t_emb.type(self.dtype)
+ cond = cond.type(self.dtype)
+
+ skips = []
+ # pack with input blocks
+ for block in self.input_blocks:
+ h = block(h, t_emb)
+ skips.append(h.feats)
+
+ if self.pe_mode == "ape":
+ h = h + self.pos_embedder(h.coords[:, 1:]).type(self.dtype)
+ for block in self.blocks:
+ h = block(h, t_emb, cond)
+
+ # unpack with output blocks
+ for block, skip in zip(self.out_blocks, reversed(skips)):
+ if self.use_skip_connection:
+ h = block(h.replace(torch.cat([h.feats, skip], dim=1)), t_emb)
+ else:
+ h = block(h, t_emb)
+
+ h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+ h = self.out_layer(h.type(x.dtype))
+ return h
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/__init__.py b/Stable3DGen/hi3dgen/models/structured_latent_vae/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..275fa4cfdc1d650b8476276a4df7492cbfb4dbde
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/structured_latent_vae/__init__.py
@@ -0,0 +1,30 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+from .encoder import SLatEncoder
+from .decoder_mesh import SLatMeshDecoder
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0403bee668efc6f51f3b0559c1c77961fcc15514
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/base.cpython-310.pyc b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/base.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..eba64705dfa72ec6781286d44567fcae79b4edb7
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/base.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/decoder_mesh.cpython-310.pyc b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/decoder_mesh.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..273bdd2f78fc2dcef567e209244d3524bd9d23b1
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/decoder_mesh.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/encoder.cpython-310.pyc b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/encoder.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..664c6bc3243b13affd4d3ada9fcc15844ee6b287
Binary files /dev/null and b/Stable3DGen/hi3dgen/models/structured_latent_vae/__pycache__/encoder.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/base.py b/Stable3DGen/hi3dgen/models/structured_latent_vae/base.py
new file mode 100755
index 0000000000000000000000000000000000000000..6403441603c829be1efb604a13ab452dfad432c9
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/structured_latent_vae/base.py
@@ -0,0 +1,145 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+from typing import *
+import torch
+import torch.nn as nn
+from ...modules.utils import convert_module_to_f16, convert_module_to_f32
+from ...modules import sparse as sp
+from ...modules.transformer import AbsolutePositionEmbedder
+from ...modules.sparse.transformer import SparseTransformerBlock
+
+
+def block_attn_config(self):
+ """
+ Return the attention configuration of the model.
+ """
+ for i in range(self.num_blocks):
+ if self.attn_mode == "shift_window":
+ yield "serialized", self.window_size, 0, (16 * (i % 2),) * 3, sp.SerializeMode.Z_ORDER
+ elif self.attn_mode == "shift_sequence":
+ yield "serialized", self.window_size, self.window_size // 2 * (i % 2), (0, 0, 0), sp.SerializeMode.Z_ORDER
+ elif self.attn_mode == "shift_order":
+ yield "serialized", self.window_size, 0, (0, 0, 0), sp.SerializeModes[i % 4]
+ elif self.attn_mode == "full":
+ yield "full", None, None, None, None
+ elif self.attn_mode == "swin":
+ yield "windowed", self.window_size, None, self.window_size // 2 * (i % 2), None
+
+
+class SparseTransformerBase(nn.Module):
+ """
+ Sparse Transformer without output layers.
+ Serve as the base class for encoder and decoder.
+ """
+ def __init__(
+ self,
+ in_channels: int,
+ model_channels: int,
+ num_blocks: int,
+ num_heads: Optional[int] = None,
+ num_head_channels: Optional[int] = 64,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
+ window_size: Optional[int] = None,
+ pe_mode: Literal["ape", "rope"] = "ape",
+ use_fp16: bool = False,
+ use_checkpoint: bool = False,
+ qk_rms_norm: bool = False,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.num_blocks = num_blocks
+ self.window_size = window_size
+ self.num_heads = num_heads or model_channels // num_head_channels
+ self.mlp_ratio = mlp_ratio
+ self.attn_mode = attn_mode
+ self.pe_mode = pe_mode
+ self.use_fp16 = use_fp16
+ self.use_checkpoint = use_checkpoint
+ self.qk_rms_norm = qk_rms_norm
+ self.dtype = torch.float16 if use_fp16 else torch.float32
+
+ if pe_mode == "ape":
+ self.pos_embedder = AbsolutePositionEmbedder(model_channels)
+
+ self.input_layer = sp.SparseLinear(in_channels, model_channels)
+ self.blocks = nn.ModuleList([
+ SparseTransformerBlock(
+ model_channels,
+ num_heads=self.num_heads,
+ mlp_ratio=self.mlp_ratio,
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_sequence=shift_sequence,
+ shift_window=shift_window,
+ serialize_mode=serialize_mode,
+ use_checkpoint=self.use_checkpoint,
+ use_rope=(pe_mode == "rope"),
+ qk_rms_norm=self.qk_rms_norm,
+ )
+ for attn_mode, window_size, shift_sequence, shift_window, serialize_mode in block_attn_config(self)
+ ])
+
+ @property
+ def device(self) -> torch.device:
+ """
+ Return the device of the model.
+ """
+ return next(self.parameters()).device
+
+ def convert_to_fp16(self) -> None:
+ """
+ Convert the torso of the model to float16.
+ """
+ self.blocks.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self) -> None:
+ """
+ Convert the torso of the model to float32.
+ """
+ self.blocks.apply(convert_module_to_f32)
+
+ def initialize_weights(self) -> None:
+ # Initialize transformer layers:
+ def _basic_init(module):
+ if isinstance(module, nn.Linear):
+ torch.nn.init.xavier_uniform_(module.weight)
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+ self.apply(_basic_init)
+
+ def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+ h = self.input_layer(x)
+ if self.pe_mode == "ape":
+ h = h + self.pos_embedder(x.coords[:, 1:])
+ h = h.type(self.dtype)
+ for block in self.blocks:
+ h = block(h)
+ return h
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/decoder_mesh.py b/Stable3DGen/hi3dgen/models/structured_latent_vae/decoder_mesh.py
new file mode 100755
index 0000000000000000000000000000000000000000..9f595771547e05221b7e7b10010dddcc0023e26f
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/structured_latent_vae/decoder_mesh.py
@@ -0,0 +1,195 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ...modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
+from ...modules import sparse as sp
+from .base import SparseTransformerBase
+from ...representations import MeshExtractResult
+from ...representations.mesh import SparseFeatures2Mesh
+
+
+class SparseSubdivideBlock3d(nn.Module):
+ """
+ A 3D subdivide block that can subdivide the sparse tensor.
+
+ Args:
+ channels: channels in the inputs and outputs.
+ out_channels: if specified, the number of output channels.
+ num_groups: the number of groups for the group norm.
+ """
+ def __init__(
+ self,
+ channels: int,
+ resolution: int,
+ out_channels: Optional[int] = None,
+ num_groups: int = 32
+ ):
+ super().__init__()
+ self.channels = channels
+ self.resolution = resolution
+ self.out_resolution = resolution * 2
+ self.out_channels = out_channels or channels
+
+ self.act_layers = nn.Sequential(
+ sp.SparseGroupNorm32(num_groups, channels),
+ sp.SparseSiLU()
+ )
+
+ self.sub = sp.SparseSubdivide()
+
+ self.out_layers = nn.Sequential(
+ sp.SparseConv3d(channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}"),
+ sp.SparseGroupNorm32(num_groups, self.out_channels),
+ sp.SparseSiLU(),
+ zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}")),
+ )
+
+ if self.out_channels == channels:
+ self.skip_connection = nn.Identity()
+ else:
+ self.skip_connection = sp.SparseConv3d(channels, self.out_channels, 1, indice_key=f"res_{self.out_resolution}")
+
+ def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+ """
+ Apply the block to a Tensor, conditioned on a timestep embedding.
+
+ Args:
+ x: an [N x C x ...] Tensor of features.
+ Returns:
+ an [N x C x ...] Tensor of outputs.
+ """
+ h = self.act_layers(x)
+ h = self.sub(h)
+ x = self.sub(x)
+ h = self.out_layers(h)
+ h = h + self.skip_connection(x)
+ return h
+
+
+class SLatMeshDecoder(SparseTransformerBase):
+ def __init__(
+ self,
+ resolution: int,
+ model_channels: int,
+ latent_channels: int,
+ num_blocks: int,
+ num_heads: Optional[int] = None,
+ num_head_channels: Optional[int] = 64,
+ mlp_ratio: float = 4,
+ attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
+ window_size: int = 8,
+ pe_mode: Literal["ape", "rope"] = "ape",
+ use_fp16: bool = False,
+ use_checkpoint: bool = False,
+ qk_rms_norm: bool = False,
+ representation_config: dict = None,
+ ):
+ super().__init__(
+ in_channels=latent_channels,
+ model_channels=model_channels,
+ num_blocks=num_blocks,
+ num_heads=num_heads,
+ num_head_channels=num_head_channels,
+ mlp_ratio=mlp_ratio,
+ attn_mode=attn_mode,
+ window_size=window_size,
+ pe_mode=pe_mode,
+ use_fp16=use_fp16,
+ use_checkpoint=use_checkpoint,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.resolution = resolution
+ self.rep_config = representation_config
+ self.mesh_extractor = SparseFeatures2Mesh(res=self.resolution*4, use_color=self.rep_config.get('use_color', False))
+ self.out_channels = self.mesh_extractor.feats_channels
+ self.upsample = nn.ModuleList([
+ SparseSubdivideBlock3d(
+ channels=model_channels,
+ resolution=resolution,
+ out_channels=model_channels // 4
+ ),
+ SparseSubdivideBlock3d(
+ channels=model_channels // 4,
+ resolution=resolution * 2,
+ out_channels=model_channels // 8
+ )
+ ])
+ self.out_layer = sp.SparseLinear(model_channels // 8, self.out_channels)
+
+ self.initialize_weights()
+ if use_fp16:
+ self.convert_to_fp16()
+
+ def initialize_weights(self) -> None:
+ super().initialize_weights()
+ # Zero-out output layers:
+ nn.init.constant_(self.out_layer.weight, 0)
+ nn.init.constant_(self.out_layer.bias, 0)
+
+ def convert_to_fp16(self) -> None:
+ """
+ Convert the torso of the model to float16.
+ """
+ super().convert_to_fp16()
+ self.upsample.apply(convert_module_to_f16)
+
+ def convert_to_fp32(self) -> None:
+ """
+ Convert the torso of the model to float32.
+ """
+ super().convert_to_fp32()
+ self.upsample.apply(convert_module_to_f32)
+
+ def to_representation(self, x: sp.SparseTensor) -> List[MeshExtractResult]:
+ """
+ Convert a batch of network outputs to 3D representations.
+
+ Args:
+ x: The [N x * x C] sparse tensor output by the network.
+
+ Returns:
+ list of representations
+ """
+ ret = []
+ for i in range(x.shape[0]):
+ mesh = self.mesh_extractor(x[i], training=self.training)
+ ret.append(mesh)
+ return ret
+
+ def forward(self, x: sp.SparseTensor) -> List[MeshExtractResult]:
+ h = super().forward(x)
+ for block in self.upsample:
+ h = block(h)
+ h = h.type(x.dtype)
+ h = self.out_layer(h)
+ return self.to_representation(h)
diff --git a/Stable3DGen/hi3dgen/models/structured_latent_vae/encoder.py b/Stable3DGen/hi3dgen/models/structured_latent_vae/encoder.py
new file mode 100755
index 0000000000000000000000000000000000000000..f86c88dcaa9a3f6b47d9293df15c0f40f24f0694
--- /dev/null
+++ b/Stable3DGen/hi3dgen/models/structured_latent_vae/encoder.py
@@ -0,0 +1,100 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ...modules import sparse as sp
+from .base import SparseTransformerBase
+
+
+class SLatEncoder(SparseTransformerBase):
+ def __init__(
+ self,
+ resolution: int,
+ in_channels: int,
+ model_channels: int,
+ latent_channels: int,
+ num_blocks: int,
+ num_heads: Optional[int] = None,
+ num_head_channels: Optional[int] = 64,
+ mlp_ratio: float = 4,
+ attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
+ window_size: int = 8,
+ pe_mode: Literal["ape", "rope"] = "ape",
+ use_fp16: bool = False,
+ use_checkpoint: bool = False,
+ qk_rms_norm: bool = False,
+ ):
+ super().__init__(
+ in_channels=in_channels,
+ model_channels=model_channels,
+ num_blocks=num_blocks,
+ num_heads=num_heads,
+ num_head_channels=num_head_channels,
+ mlp_ratio=mlp_ratio,
+ attn_mode=attn_mode,
+ window_size=window_size,
+ pe_mode=pe_mode,
+ use_fp16=use_fp16,
+ use_checkpoint=use_checkpoint,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.resolution = resolution
+ self.out_layer = sp.SparseLinear(model_channels, 2 * latent_channels)
+
+ self.initialize_weights()
+ if use_fp16:
+ self.convert_to_fp16()
+
+ def initialize_weights(self) -> None:
+ super().initialize_weights()
+ # Zero-out output layers:
+ nn.init.constant_(self.out_layer.weight, 0)
+ nn.init.constant_(self.out_layer.bias, 0)
+
+ def forward(self, x: sp.SparseTensor, sample_posterior=True, return_raw=False):
+ h = super().forward(x)
+ h = h.type(x.dtype)
+ h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+ h = self.out_layer(h)
+
+ # Sample from the posterior distribution
+ mean, logvar = h.feats.chunk(2, dim=-1)
+ if sample_posterior:
+ std = torch.exp(0.5 * logvar)
+ z = mean + std * torch.randn_like(std)
+ else:
+ z = mean
+ z = h.replace(z)
+
+ if return_raw:
+ return z, mean, logvar
+ else:
+ return z
diff --git a/Stable3DGen/hi3dgen/modules/__pycache__/norm.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/__pycache__/norm.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9d4734766a74369055622ec8cc57c29f51bf317c
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/__pycache__/norm.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/__pycache__/spatial.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/__pycache__/spatial.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..516594a63f77d930ad164d110b4d0da77b075a27
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/__pycache__/spatial.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/__pycache__/utils.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c8945460e548240fb3361f46690cb65afb8856ed
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/__pycache__/utils.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/attention/__init__.py b/Stable3DGen/hi3dgen/modules/attention/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..25af9a695abbdc86e669c09983dbecfff05e174b
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/attention/__init__.py
@@ -0,0 +1,60 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+
+BACKEND = 'xformers'
+DEBUG = False
+
+def __from_env():
+ import os
+
+ global BACKEND
+ global DEBUG
+
+ env_attn_backend = os.environ.get('ATTN_BACKEND')
+ env_sttn_debug = os.environ.get('ATTN_DEBUG')
+
+ if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'sdpa', 'naive']:
+ BACKEND = env_attn_backend
+ if env_sttn_debug is not None:
+ DEBUG = env_sttn_debug == '1'
+
+ print(f"[ATTENTION] Using backend: {BACKEND}")
+
+
+__from_env()
+
+
+def set_backend(backend: Literal['xformers', 'flash_attn']):
+ global BACKEND
+ BACKEND = backend
+
+def set_debug(debug: bool):
+ global DEBUG
+ DEBUG = debug
+
+
+from .full_attn import *
+from .modules import *
diff --git a/Stable3DGen/hi3dgen/modules/attention/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/attention/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2b04c1be1be23fc1807cb16777fd1b9d1c2d61c1
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/attention/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/attention/__pycache__/full_attn.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/attention/__pycache__/full_attn.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..de151097651661a8e38da9ffb76974602a3a9068
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/attention/__pycache__/full_attn.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/attention/__pycache__/modules.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/attention/__pycache__/modules.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f07b9b75a1d99edff623197390368f6054a6c4e2
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/attention/__pycache__/modules.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/attention/full_attn.py b/Stable3DGen/hi3dgen/modules/attention/full_attn.py
new file mode 100755
index 0000000000000000000000000000000000000000..3a622983b85a3c5c047cf4e996282cabcbc7dd8b
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/attention/full_attn.py
@@ -0,0 +1,164 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import math
+from . import DEBUG, BACKEND
+
+if BACKEND == 'xformers':
+ import xformers.ops as xops
+elif BACKEND == 'flash_attn':
+ import flash_attn
+elif BACKEND == 'sdpa':
+ from torch.nn.functional import scaled_dot_product_attention as sdpa
+elif BACKEND == 'naive':
+ pass
+else:
+ raise ValueError(f"Unknown attention backend: {BACKEND}")
+
+
+__all__ = [
+ 'scaled_dot_product_attention',
+]
+
+
+def _naive_sdpa(q, k, v):
+ """
+ Naive implementation of scaled dot product attention.
+ """
+ q = q.permute(0, 2, 1, 3) # [N, H, L, C]
+ k = k.permute(0, 2, 1, 3) # [N, H, L, C]
+ v = v.permute(0, 2, 1, 3) # [N, H, L, C]
+ scale_factor = 1 / math.sqrt(q.size(-1))
+ attn_weight = q @ k.transpose(-2, -1) * scale_factor
+ attn_weight = torch.softmax(attn_weight, dim=-1)
+ out = attn_weight @ v
+ out = out.permute(0, 2, 1, 3) # [N, L, H, C]
+ return out
+
+
+@overload
+def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention.
+
+ Args:
+ qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs.
+ """
+ ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention.
+
+ Args:
+ q (torch.Tensor): A [N, L, H, C] tensor containing Qs.
+ kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs.
+ """
+ ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention.
+
+ Args:
+ q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs.
+ k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks.
+ v (torch.Tensor): A [N, L, H, Co] tensor containing Vs.
+
+ Note:
+ k and v are assumed to have the same coordinate map.
+ """
+ ...
+
+def scaled_dot_product_attention(*args, **kwargs):
+ arg_names_dict = {
+ 1: ['qkv'],
+ 2: ['q', 'kv'],
+ 3: ['q', 'k', 'v']
+ }
+ num_all_args = len(args) + len(kwargs)
+ assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+ for key in arg_names_dict[num_all_args][len(args):]:
+ assert key in kwargs, f"Missing argument {key}"
+
+ if num_all_args == 1:
+ qkv = args[0] if len(args) > 0 else kwargs['qkv']
+ assert len(qkv.shape) == 5 and qkv.shape[2] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]"
+ device = qkv.device
+
+ elif num_all_args == 2:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ kv = args[1] if len(args) > 1 else kwargs['kv']
+ assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+ assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+ device = q.device
+
+ elif num_all_args == 3:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ k = args[1] if len(args) > 1 else kwargs['k']
+ v = args[2] if len(args) > 2 else kwargs['v']
+ assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+ assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+ assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+ device = q.device
+
+ if BACKEND == 'xformers':
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=2)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=2)
+ out = xops.memory_efficient_attention(q, k, v)
+ elif BACKEND == 'flash_attn':
+ if num_all_args == 1:
+ out = flash_attn.flash_attn_qkvpacked_func(qkv)
+ elif num_all_args == 2:
+ out = flash_attn.flash_attn_kvpacked_func(q, kv)
+ elif num_all_args == 3:
+ out = flash_attn.flash_attn_func(q, k, v)
+ elif BACKEND == 'sdpa':
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=2)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=2)
+ q = q.permute(0, 2, 1, 3) # [N, H, L, C]
+ k = k.permute(0, 2, 1, 3) # [N, H, L, C]
+ v = v.permute(0, 2, 1, 3) # [N, H, L, C]
+ out = sdpa(q, k, v) # [N, H, L, C]
+ out = out.permute(0, 2, 1, 3) # [N, L, H, C]
+ elif BACKEND == 'naive':
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=2)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=2)
+ out = _naive_sdpa(q, k, v)
+ else:
+ raise ValueError(f"Unknown attention module: {BACKEND}")
+
+ return out
diff --git a/Stable3DGen/hi3dgen/modules/attention/modules.py b/Stable3DGen/hi3dgen/modules/attention/modules.py
new file mode 100755
index 0000000000000000000000000000000000000000..a6ce75da6db996f6d63cd3fedcf837ee8ee48fff
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/attention/modules.py
@@ -0,0 +1,170 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .full_attn import scaled_dot_product_attention
+
+
+class MultiHeadRMSNorm(nn.Module):
+ def __init__(self, dim: int, heads: int):
+ super().__init__()
+ self.scale = dim ** 0.5
+ self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
+
+
+class RotaryPositionEmbedder(nn.Module):
+ def __init__(self, hidden_size: int, in_channels: int = 3):
+ super().__init__()
+ assert hidden_size % 2 == 0, "Hidden size must be divisible by 2"
+ self.hidden_size = hidden_size
+ self.in_channels = in_channels
+ self.freq_dim = hidden_size // in_channels // 2
+ self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+ self.freqs = 1.0 / (10000 ** self.freqs)
+
+ def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+ self.freqs = self.freqs.to(indices.device)
+ phases = torch.outer(indices, self.freqs)
+ phases = torch.polar(torch.ones_like(phases), phases)
+ return phases
+
+ def _rotary_embedding(self, x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+ x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+ x_rotated = x_complex * phases
+ x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+ return x_embed
+
+ def forward(self, q: torch.Tensor, k: torch.Tensor, indices: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+ """
+ Args:
+ q (sp.SparseTensor): [..., N, D] tensor of queries
+ k (sp.SparseTensor): [..., N, D] tensor of keys
+ indices (torch.Tensor): [..., N, C] tensor of spatial positions
+ """
+ if indices is None:
+ indices = torch.arange(q.shape[-2], device=q.device)
+ if len(q.shape) > 2:
+ indices = indices.unsqueeze(0).expand(q.shape[:-2] + (-1,))
+
+ phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
+ if phases.shape[1] < self.hidden_size // 2:
+ phases = torch.cat([phases, torch.polar(
+ torch.ones(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device),
+ torch.zeros(*phases.shape[:-1], self.hidden_size // 2 - phases.shape[1], device=phases.device)
+ )], dim=-1)
+ q_embed = self._rotary_embedding(q, phases)
+ k_embed = self._rotary_embedding(k, phases)
+ return q_embed, k_embed
+
+
+class MultiHeadAttention(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ ctx_channels: Optional[int]=None,
+ type: Literal["self", "cross"] = "self",
+ attn_mode: Literal["full", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ qkv_bias: bool = True,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ ):
+ super().__init__()
+ assert channels % num_heads == 0
+ assert type in ["self", "cross"], f"Invalid attention type: {type}"
+ assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}"
+ assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+
+ if attn_mode == "windowed":
+ raise NotImplementedError("Windowed attention is not yet implemented")
+
+ self.channels = channels
+ self.head_dim = channels // num_heads
+ self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+ self.num_heads = num_heads
+ self._type = type
+ self.attn_mode = attn_mode
+ self.window_size = window_size
+ self.shift_window = shift_window
+ self.use_rope = use_rope
+ self.qk_rms_norm = qk_rms_norm
+
+ if self._type == "self":
+ self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+ else:
+ self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+ self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+
+ if self.qk_rms_norm:
+ self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+ self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+
+ self.to_out = nn.Linear(channels, channels)
+
+ if use_rope:
+ self.rope = RotaryPositionEmbedder(channels)
+
+ def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, indices: Optional[torch.Tensor] = None) -> torch.Tensor:
+ B, L, C = x.shape
+ if self._type == "self":
+ qkv = self.to_qkv(x)
+ qkv = qkv.reshape(B, L, 3, self.num_heads, -1)
+ if self.use_rope:
+ q, k, v = qkv.unbind(dim=2)
+ q, k = self.rope(q, k, indices)
+ qkv = torch.stack([q, k, v], dim=2)
+ if self.attn_mode == "full":
+ if self.qk_rms_norm:
+ q, k, v = qkv.unbind(dim=2)
+ q = self.q_rms_norm(q)
+ k = self.k_rms_norm(k)
+ h = scaled_dot_product_attention(q, k, v)
+ else:
+ h = scaled_dot_product_attention(qkv)
+ elif self.attn_mode == "windowed":
+ raise NotImplementedError("Windowed attention is not yet implemented")
+ else:
+ Lkv = context.shape[1]
+ q = self.to_q(x)
+ kv = self.to_kv(context)
+ q = q.reshape(B, L, self.num_heads, -1)
+ kv = kv.reshape(B, Lkv, 2, self.num_heads, -1)
+ if self.qk_rms_norm:
+ q = self.q_rms_norm(q)
+ k, v = kv.unbind(dim=2)
+ k = self.k_rms_norm(k)
+ h = scaled_dot_product_attention(q, k, v)
+ else:
+ h = scaled_dot_product_attention(q, kv)
+ h = h.reshape(B, L, -1)
+ h = self.to_out(h)
+ return h
diff --git a/Stable3DGen/hi3dgen/modules/norm.py b/Stable3DGen/hi3dgen/modules/norm.py
new file mode 100755
index 0000000000000000000000000000000000000000..d52f676604d869692ce35a39c72d1c3504238c72
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/norm.py
@@ -0,0 +1,49 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch
+import torch.nn as nn
+
+
+class LayerNorm32(nn.LayerNorm):
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return super().forward(x.float()).type(x.dtype)
+
+
+class GroupNorm32(nn.GroupNorm):
+ """
+ A GroupNorm layer that converts to float32 before the forward pass.
+ """
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return super().forward(x.float()).type(x.dtype)
+
+
+class ChannelLayerNorm32(LayerNorm32):
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ DIM = x.dim()
+ x = x.permute(0, *range(2, DIM), 1).contiguous()
+ x = super().forward(x)
+ x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous()
+ return x
+
\ No newline at end of file
diff --git a/Stable3DGen/hi3dgen/modules/sparse/__init__.py b/Stable3DGen/hi3dgen/modules/sparse/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..7c01a7e08a8d0feb48d73444f8c5a224a46dda20
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/__init__.py
@@ -0,0 +1,126 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+
+BACKEND = 'spconv'
+DEBUG = False
+ATTN = 'xformers'
+
+def __from_env():
+ import os
+
+ global BACKEND
+ global DEBUG
+ global ATTN
+
+ env_sparse_backend = os.environ.get('SPARSE_BACKEND')
+ env_sparse_debug = os.environ.get('SPARSE_DEBUG')
+ env_sparse_attn = os.environ.get('SPARSE_ATTN_BACKEND')
+ if env_sparse_attn is None:
+ env_sparse_attn = os.environ.get('ATTN_BACKEND')
+
+ if env_sparse_backend is not None and env_sparse_backend in ['spconv', 'torchsparse']:
+ BACKEND = env_sparse_backend
+ if env_sparse_debug is not None:
+ DEBUG = env_sparse_debug == '1'
+ if env_sparse_attn is not None and env_sparse_attn in ['xformers', 'flash_attn']:
+ ATTN = env_sparse_attn
+
+ print(f"[SPARSE] Backend: {BACKEND}, Attention: {ATTN}")
+
+
+__from_env()
+
+
+def set_backend(backend: Literal['spconv', 'torchsparse']):
+ global BACKEND
+ BACKEND = backend
+
+def set_debug(debug: bool):
+ global DEBUG
+ DEBUG = debug
+
+def set_attn(attn: Literal['xformers', 'flash_attn']):
+ global ATTN
+ ATTN = attn
+
+
+import importlib
+
+__attributes = {
+ 'SparseTensor': 'basic',
+ 'sparse_batch_broadcast': 'basic',
+ 'sparse_batch_op': 'basic',
+ 'sparse_cat': 'basic',
+ 'sparse_unbind': 'basic',
+ 'SparseGroupNorm': 'norm',
+ 'SparseLayerNorm': 'norm',
+ 'SparseGroupNorm32': 'norm',
+ 'SparseLayerNorm32': 'norm',
+ 'SparseReLU': 'nonlinearity',
+ 'SparseSiLU': 'nonlinearity',
+ 'SparseGELU': 'nonlinearity',
+ 'SparseActivation': 'nonlinearity',
+ 'SparseLinear': 'linear',
+ 'sparse_scaled_dot_product_attention': 'attention',
+ 'SerializeMode': 'attention',
+ 'sparse_serialized_scaled_dot_product_self_attention': 'attention',
+ 'sparse_windowed_scaled_dot_product_self_attention': 'attention',
+ 'SparseMultiHeadAttention': 'attention',
+ 'SparseConv3d': 'conv',
+ 'SparseInverseConv3d': 'conv',
+ 'SparseDownsample': 'spatial',
+ 'SparseUpsample': 'spatial',
+ 'SparseSubdivide' : 'spatial'
+}
+
+__submodules = ['transformer']
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+ if name not in globals():
+ if name in __attributes:
+ module_name = __attributes[name]
+ module = importlib.import_module(f".{module_name}", __name__)
+ globals()[name] = getattr(module, name)
+ elif name in __submodules:
+ module = importlib.import_module(f".{name}", __name__)
+ globals()[name] = module
+ else:
+ raise AttributeError(f"module {__name__} has no attribute {name}")
+ return globals()[name]
+
+
+# For Pylance
+if __name__ == '__main__':
+ from .basic import *
+ from .norm import *
+ from .nonlinearity import *
+ from .linear import *
+ from .attention import *
+ from .conv import *
+ from .spatial import *
+ import transformer
diff --git a/Stable3DGen/hi3dgen/modules/sparse/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..012cb94ca416199c779e2bed9264a62a3eb4b10d
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/__pycache__/basic.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/basic.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..23ad72c076d3a08aafbbb9cb6324bfc53b12444f
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/basic.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/__pycache__/linear.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/linear.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bc9eb38f72dfc06ef788ea6dc621cc9dfcd3dcf8
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/linear.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/__pycache__/nonlinearity.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/nonlinearity.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..37ebebd4ffbf4d213bba4177e435c5d04cc858db
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/nonlinearity.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/__pycache__/norm.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/norm.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b9779287b077ba11ef223df8002bbbc331bcf35c
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/norm.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/__pycache__/spatial.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/spatial.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0ca2cb6deee3d7797fac5030286ab6b7cbfa60f9
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/__pycache__/spatial.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/__init__.py b/Stable3DGen/hi3dgen/modules/sparse/attention/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..67dabbf1c706c4bf33b98805b644c5844cec394d
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/attention/__init__.py
@@ -0,0 +1,28 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from .full_attn import *
+from .serialized_attn import *
+from .windowed_attn import *
+from .modules import *
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..200fa8b8e85454c1de3b1140a5e0ef79a64e7463
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/full_attn.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/full_attn.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..19c6bfadd4494fa0cb1f8b3d690bf2fea4bde189
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/full_attn.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/modules.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/modules.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f4d8a3f1afbeb7f2445e81d7922a4a32ae6e91d0
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/modules.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/serialized_attn.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/serialized_attn.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cd03010a3b581cbc637b62fd2db15921421ad204
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/serialized_attn.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/windowed_attn.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/windowed_attn.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9cca0595c16bfe9fcab9f4116ccc15b58293e692
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/attention/__pycache__/windowed_attn.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/full_attn.py b/Stable3DGen/hi3dgen/modules/sparse/attention/full_attn.py
new file mode 100755
index 0000000000000000000000000000000000000000..79085079fe15e0ee6d247e86e9872306270f7383
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/attention/full_attn.py
@@ -0,0 +1,239 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+from .. import SparseTensor
+from .. import DEBUG, ATTN
+
+if ATTN == 'xformers':
+ import xformers.ops as xops
+elif ATTN == 'flash_attn':
+ import flash_attn
+else:
+ raise ValueError(f"Unknown attention module: {ATTN}")
+
+
+__all__ = [
+ 'sparse_scaled_dot_product_attention',
+]
+
+
+@overload
+def sparse_scaled_dot_product_attention(qkv: SparseTensor) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ qkv (SparseTensor): A [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, kv: Union[SparseTensor, torch.Tensor]) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, *, H, C] sparse tensor containing Qs.
+ kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor or a [N, L, 2, H, C] dense tensor containing Ks and Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, kv: SparseTensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, L, H, C] dense tensor containing Qs.
+ kv (SparseTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor containing Ks and Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, k: SparseTensor, v: SparseTensor) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+ k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+ v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
+
+ Note:
+ k and v are assumed to have the same coordinate map.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: SparseTensor, k: torch.Tensor, v: torch.Tensor) -> SparseTensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (SparseTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+ k (torch.Tensor): A [N, L, H, Ci] dense tensor containing Ks.
+ v (torch.Tensor): A [N, L, H, Co] dense tensor containing Vs.
+ """
+ ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, k: SparseTensor, v: SparseTensor) -> torch.Tensor:
+ """
+ Apply scaled dot product attention to a sparse tensor.
+
+ Args:
+ q (torch.Tensor): A [N, L, H, Ci] dense tensor containing Qs.
+ k (SparseTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+ v (SparseTensor): A [N, *, H, Co] sparse tensor containing Vs.
+ """
+ ...
+
+def sparse_scaled_dot_product_attention(*args, **kwargs):
+ arg_names_dict = {
+ 1: ['qkv'],
+ 2: ['q', 'kv'],
+ 3: ['q', 'k', 'v']
+ }
+ num_all_args = len(args) + len(kwargs)
+ assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+ for key in arg_names_dict[num_all_args][len(args):]:
+ assert key in kwargs, f"Missing argument {key}"
+
+ if num_all_args == 1:
+ qkv = args[0] if len(args) > 0 else kwargs['qkv']
+ assert isinstance(qkv, SparseTensor), f"qkv must be a SparseTensor, got {type(qkv)}"
+ assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+ device = qkv.device
+
+ s = qkv
+ q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
+ kv_seqlen = q_seqlen
+ qkv = qkv.feats # [T, 3, H, C]
+
+ elif num_all_args == 2:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ kv = args[1] if len(args) > 1 else kwargs['kv']
+ assert isinstance(q, SparseTensor) and isinstance(kv, (SparseTensor, torch.Tensor)) or \
+ isinstance(q, torch.Tensor) and isinstance(kv, SparseTensor), \
+ f"Invalid types, got {type(q)} and {type(kv)}"
+ assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+ device = q.device
+
+ if isinstance(q, SparseTensor):
+ assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]"
+ s = q
+ q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+ q = q.feats # [T_Q, H, C]
+ else:
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+ s = None
+ N, L, H, C = q.shape
+ q_seqlen = [L] * N
+ q = q.reshape(N * L, H, C) # [T_Q, H, C]
+
+ if isinstance(kv, SparseTensor):
+ assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]"
+ kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
+ kv = kv.feats # [T_KV, 2, H, C]
+ else:
+ assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+ N, L, _, H, C = kv.shape
+ kv_seqlen = [L] * N
+ kv = kv.reshape(N * L, 2, H, C) # [T_KV, 2, H, C]
+
+ elif num_all_args == 3:
+ q = args[0] if len(args) > 0 else kwargs['q']
+ k = args[1] if len(args) > 1 else kwargs['k']
+ v = args[2] if len(args) > 2 else kwargs['v']
+ assert isinstance(q, SparseTensor) and isinstance(k, (SparseTensor, torch.Tensor)) and type(k) == type(v) or \
+ isinstance(q, torch.Tensor) and isinstance(k, SparseTensor) and isinstance(v, SparseTensor), \
+ f"Invalid types, got {type(q)}, {type(k)}, and {type(v)}"
+ assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+ device = q.device
+
+ if isinstance(q, SparseTensor):
+ assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, Ci]"
+ s = q
+ q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+ q = q.feats # [T_Q, H, Ci]
+ else:
+ assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+ s = None
+ N, L, H, CI = q.shape
+ q_seqlen = [L] * N
+ q = q.reshape(N * L, H, CI) # [T_Q, H, Ci]
+
+ if isinstance(k, SparseTensor):
+ assert len(k.shape) == 3, f"Invalid shape for k, got {k.shape}, expected [N, *, H, Ci]"
+ assert len(v.shape) == 3, f"Invalid shape for v, got {v.shape}, expected [N, *, H, Co]"
+ kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
+ k = k.feats # [T_KV, H, Ci]
+ v = v.feats # [T_KV, H, Co]
+ else:
+ assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+ assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+ N, L, H, CI, CO = *k.shape, v.shape[-1]
+ kv_seqlen = [L] * N
+ k = k.reshape(N * L, H, CI) # [T_KV, H, Ci]
+ v = v.reshape(N * L, H, CO) # [T_KV, H, Co]
+
+ if DEBUG:
+ if s is not None:
+ for i in range(s.shape[0]):
+ assert (s.coords[s.layout[i]] == i).all(), f"SparseScaledDotProductSelfAttention: batch index mismatch"
+ if num_all_args in [2, 3]:
+ assert q.shape[:2] == [1, sum(q_seqlen)], f"SparseScaledDotProductSelfAttention: q shape mismatch"
+ if num_all_args == 3:
+ assert k.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: k shape mismatch"
+ assert v.shape[:2] == [1, sum(kv_seqlen)], f"SparseScaledDotProductSelfAttention: v shape mismatch"
+
+ if ATTN == 'xformers':
+ if num_all_args == 1:
+ q, k, v = qkv.unbind(dim=1)
+ elif num_all_args == 2:
+ k, v = kv.unbind(dim=1)
+ q = q.unsqueeze(0)
+ k = k.unsqueeze(0)
+ v = v.unsqueeze(0)
+ mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
+ out = xops.memory_efficient_attention(q, k, v, mask)[0]
+ elif ATTN == 'flash_attn':
+ cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+ if num_all_args in [2, 3]:
+ cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+ if num_all_args == 1:
+ out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
+ elif num_all_args == 2:
+ out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+ elif num_all_args == 3:
+ out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN}")
+
+ if s is not None:
+ return s.replace(out)
+ else:
+ return out.reshape(N, L, H, -1)
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/modules.py b/Stable3DGen/hi3dgen/modules/sparse/attention/modules.py
new file mode 100755
index 0000000000000000000000000000000000000000..9a54356b4045b8283703d97b19a07891e45faf2b
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/attention/modules.py
@@ -0,0 +1,163 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .. import SparseTensor
+from .full_attn import sparse_scaled_dot_product_attention
+from .serialized_attn import SerializeMode, sparse_serialized_scaled_dot_product_self_attention
+from .windowed_attn import sparse_windowed_scaled_dot_product_self_attention
+from ...attention import RotaryPositionEmbedder
+
+
+class SparseMultiHeadRMSNorm(nn.Module):
+ def __init__(self, dim: int, heads: int):
+ super().__init__()
+ self.scale = dim ** 0.5
+ self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+ def forward(self, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]:
+ x_type = x.dtype
+ x = x.float()
+ if isinstance(x, SparseTensor):
+ x = x.replace(F.normalize(x.feats, dim=-1))
+ else:
+ x = F.normalize(x, dim=-1)
+ return (x * self.gamma * self.scale).to(x_type)
+
+
+class SparseMultiHeadAttention(nn.Module):
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ ctx_channels: Optional[int] = None,
+ type: Literal["self", "cross"] = "self",
+ attn_mode: Literal["full", "serialized", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_sequence: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ serialize_mode: Optional[SerializeMode] = None,
+ qkv_bias: bool = True,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ ):
+ super().__init__()
+ assert channels % num_heads == 0
+ assert type in ["self", "cross"], f"Invalid attention type: {type}"
+ assert attn_mode in ["full", "serialized", "windowed"], f"Invalid attention mode: {attn_mode}"
+ assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+ assert type == "self" or use_rope is False, "Rotary position embeddings only supported for self-attention"
+ self.channels = channels
+ self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+ self.num_heads = num_heads
+ self._type = type
+ self.attn_mode = attn_mode
+ self.window_size = window_size
+ self.shift_sequence = shift_sequence
+ self.shift_window = shift_window
+ self.serialize_mode = serialize_mode
+ self.use_rope = use_rope
+ self.qk_rms_norm = qk_rms_norm
+
+ if self._type == "self":
+ self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+ else:
+ self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+ self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+
+ if self.qk_rms_norm:
+ self.q_rms_norm = SparseMultiHeadRMSNorm(channels // num_heads, num_heads)
+ self.k_rms_norm = SparseMultiHeadRMSNorm(channels // num_heads, num_heads)
+
+ self.to_out = nn.Linear(channels, channels)
+
+ if use_rope:
+ self.rope = RotaryPositionEmbedder(channels)
+
+ @staticmethod
+ def _linear(module: nn.Linear, x: Union[SparseTensor, torch.Tensor]) -> Union[SparseTensor, torch.Tensor]:
+ if isinstance(x, SparseTensor):
+ return x.replace(module(x.feats))
+ else:
+ return module(x)
+
+ @staticmethod
+ def _reshape_chs(x: Union[SparseTensor, torch.Tensor], shape: Tuple[int, ...]) -> Union[SparseTensor, torch.Tensor]:
+ if isinstance(x, SparseTensor):
+ return x.reshape(*shape)
+ else:
+ return x.reshape(*x.shape[:2], *shape)
+
+ def _fused_pre(self, x: Union[SparseTensor, torch.Tensor], num_fused: int) -> Union[SparseTensor, torch.Tensor]:
+ if isinstance(x, SparseTensor):
+ x_feats = x.feats.unsqueeze(0)
+ else:
+ x_feats = x
+ x_feats = x_feats.reshape(*x_feats.shape[:2], num_fused, self.num_heads, -1)
+ return x.replace(x_feats.squeeze(0)) if isinstance(x, SparseTensor) else x_feats
+
+ def _rope(self, qkv: SparseTensor) -> SparseTensor:
+ q, k, v = qkv.feats.unbind(dim=1) # [T, H, C]
+ q, k = self.rope(q, k, qkv.coords[:, 1:])
+ qkv = qkv.replace(torch.stack([q, k, v], dim=1))
+ return qkv
+
+ def forward(self, x: Union[SparseTensor, torch.Tensor], context: Optional[Union[SparseTensor, torch.Tensor]] = None) -> Union[SparseTensor, torch.Tensor]:
+ if self._type == "self":
+ qkv = self._linear(self.to_qkv, x)
+ qkv = self._fused_pre(qkv, num_fused=3)
+ if self.use_rope:
+ qkv = self._rope(qkv)
+ if self.qk_rms_norm:
+ q, k, v = qkv.unbind(dim=1)
+ q = self.q_rms_norm(q)
+ k = self.k_rms_norm(k)
+ qkv = qkv.replace(torch.stack([q.feats, k.feats, v.feats], dim=1))
+ if self.attn_mode == "full":
+ h = sparse_scaled_dot_product_attention(qkv)
+ elif self.attn_mode == "serialized":
+ h = sparse_serialized_scaled_dot_product_self_attention(
+ qkv, self.window_size, serialize_mode=self.serialize_mode, shift_sequence=self.shift_sequence, shift_window=self.shift_window
+ )
+ elif self.attn_mode == "windowed":
+ h = sparse_windowed_scaled_dot_product_self_attention(
+ qkv, self.window_size, shift_window=self.shift_window
+ )
+ else:
+ q = self._linear(self.to_q, x)
+ q = self._reshape_chs(q, (self.num_heads, -1))
+ kv = self._linear(self.to_kv, context)
+ kv = self._fused_pre(kv, num_fused=2)
+ if self.qk_rms_norm:
+ q = self.q_rms_norm(q)
+ k, v = kv.unbind(dim=1)
+ k = self.k_rms_norm(k)
+ kv = kv.replace(torch.stack([k.feats, v.feats], dim=1))
+ h = sparse_scaled_dot_product_attention(q, kv)
+ h = self._reshape_chs(h, (-1,))
+ h = self._linear(self.to_out, h)
+ return h
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/serialized_attn.py b/Stable3DGen/hi3dgen/modules/sparse/attention/serialized_attn.py
new file mode 100755
index 0000000000000000000000000000000000000000..62b973e2bfa0f82b17f67183165fca359e1e0d5d
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/attention/serialized_attn.py
@@ -0,0 +1,217 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+from enum import Enum
+import torch
+import math
+from .. import SparseTensor
+from .. import DEBUG, ATTN
+
+if ATTN == 'xformers':
+ import xformers.ops as xops
+elif ATTN == 'flash_attn':
+ import flash_attn
+else:
+ raise ValueError(f"Unknown attention module: {ATTN}")
+
+
+__all__ = [
+ 'sparse_serialized_scaled_dot_product_self_attention',
+]
+
+
+class SerializeMode(Enum):
+ Z_ORDER = 0
+ Z_ORDER_TRANSPOSED = 1
+ HILBERT = 2
+ HILBERT_TRANSPOSED = 3
+
+
+SerializeModes = [
+ SerializeMode.Z_ORDER,
+ SerializeMode.Z_ORDER_TRANSPOSED,
+ SerializeMode.HILBERT,
+ SerializeMode.HILBERT_TRANSPOSED
+]
+
+
+def calc_serialization(
+ tensor: SparseTensor,
+ window_size: int,
+ serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
+ shift_sequence: int = 0,
+ shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> Tuple[torch.Tensor, torch.Tensor, List[int]]:
+ """
+ Calculate serialization and partitioning for a set of coordinates.
+
+ Args:
+ tensor (SparseTensor): The input tensor.
+ window_size (int): The window size to use.
+ serialize_mode (SerializeMode): The serialization mode to use.
+ shift_sequence (int): The shift of serialized sequence.
+ shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
+
+ Returns:
+ (torch.Tensor, torch.Tensor): Forwards and backwards indices.
+ """
+ fwd_indices = []
+ bwd_indices = []
+ seq_lens = []
+ seq_batch_indices = []
+ offsets = [0]
+
+ if 'vox2seq' not in globals():
+ import vox2seq
+
+ # Serialize the input
+ serialize_coords = tensor.coords[:, 1:].clone()
+ serialize_coords += torch.tensor(shift_window, dtype=torch.int32, device=tensor.device).reshape(1, 3)
+ if serialize_mode == SerializeMode.Z_ORDER:
+ code = vox2seq.encode(serialize_coords, mode='z_order', permute=[0, 1, 2])
+ elif serialize_mode == SerializeMode.Z_ORDER_TRANSPOSED:
+ code = vox2seq.encode(serialize_coords, mode='z_order', permute=[1, 0, 2])
+ elif serialize_mode == SerializeMode.HILBERT:
+ code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[0, 1, 2])
+ elif serialize_mode == SerializeMode.HILBERT_TRANSPOSED:
+ code = vox2seq.encode(serialize_coords, mode='hilbert', permute=[1, 0, 2])
+ else:
+ raise ValueError(f"Unknown serialize mode: {serialize_mode}")
+
+ for bi, s in enumerate(tensor.layout):
+ num_points = s.stop - s.start
+ num_windows = (num_points + window_size - 1) // window_size
+ valid_window_size = num_points / num_windows
+ to_ordered = torch.argsort(code[s.start:s.stop])
+ if num_windows == 1:
+ fwd_indices.append(to_ordered)
+ bwd_indices.append(torch.zeros_like(to_ordered).scatter_(0, to_ordered, torch.arange(num_points, device=tensor.device)))
+ fwd_indices[-1] += s.start
+ bwd_indices[-1] += offsets[-1]
+ seq_lens.append(num_points)
+ seq_batch_indices.append(bi)
+ offsets.append(offsets[-1] + seq_lens[-1])
+ else:
+ # Partition the input
+ offset = 0
+ mids = [(i + 0.5) * valid_window_size + shift_sequence for i in range(num_windows)]
+ split = [math.floor(i * valid_window_size + shift_sequence) for i in range(num_windows + 1)]
+ bwd_index = torch.zeros((num_points,), dtype=torch.int64, device=tensor.device)
+ for i in range(num_windows):
+ mid = mids[i]
+ valid_start = split[i]
+ valid_end = split[i + 1]
+ padded_start = math.floor(mid - 0.5 * window_size)
+ padded_end = padded_start + window_size
+ fwd_indices.append(to_ordered[torch.arange(padded_start, padded_end, device=tensor.device) % num_points])
+ offset += valid_start - padded_start
+ bwd_index.scatter_(0, fwd_indices[-1][valid_start-padded_start:valid_end-padded_start], torch.arange(offset, offset + valid_end - valid_start, device=tensor.device))
+ offset += padded_end - valid_start
+ fwd_indices[-1] += s.start
+ seq_lens.extend([window_size] * num_windows)
+ seq_batch_indices.extend([bi] * num_windows)
+ bwd_indices.append(bwd_index + offsets[-1])
+ offsets.append(offsets[-1] + num_windows * window_size)
+
+ fwd_indices = torch.cat(fwd_indices)
+ bwd_indices = torch.cat(bwd_indices)
+
+ return fwd_indices, bwd_indices, seq_lens, seq_batch_indices
+
+
+def sparse_serialized_scaled_dot_product_self_attention(
+ qkv: SparseTensor,
+ window_size: int,
+ serialize_mode: SerializeMode = SerializeMode.Z_ORDER,
+ shift_sequence: int = 0,
+ shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> SparseTensor:
+ """
+ Apply serialized scaled dot product self attention to a sparse tensor.
+
+ Args:
+ qkv (SparseTensor): [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+ window_size (int): The window size to use.
+ serialize_mode (SerializeMode): The serialization mode to use.
+ shift_sequence (int): The shift of serialized sequence.
+ shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
+ shift (int): The shift to use.
+ """
+ assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+
+ serialization_spatial_cache_name = f'serialization_{serialize_mode}_{window_size}_{shift_sequence}_{shift_window}'
+ serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name)
+ if serialization_spatial_cache is None:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = calc_serialization(qkv, window_size, serialize_mode, shift_sequence, shift_window)
+ qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, seq_batch_indices))
+ else:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = serialization_spatial_cache
+
+ M = fwd_indices.shape[0]
+ T = qkv.feats.shape[0]
+ H = qkv.feats.shape[2]
+ C = qkv.feats.shape[3]
+
+ qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C]
+
+ if DEBUG:
+ start = 0
+ qkv_coords = qkv.coords[fwd_indices]
+ for i in range(len(seq_lens)):
+ assert (qkv_coords[start:start+seq_lens[i], 0] == seq_batch_indices[i]).all(), f"SparseWindowedScaledDotProductSelfAttention: batch index mismatch"
+ start += seq_lens[i]
+
+ if all([seq_len == window_size for seq_len in seq_lens]):
+ B = len(seq_lens)
+ N = window_size
+ qkv_feats = qkv_feats.reshape(B, N, 3, H, C)
+ if ATTN == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=2) # [B, N, H, C]
+ out = xops.memory_efficient_attention(q, k, v) # [B, N, H, C]
+ elif ATTN == 'flash_attn':
+ out = flash_attn.flash_attn_qkvpacked_func(qkv_feats) # [B, N, H, C]
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN}")
+ out = out.reshape(B * N, H, C) # [M, H, C]
+ else:
+ if ATTN == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=1) # [M, H, C]
+ q = q.unsqueeze(0) # [1, M, H, C]
+ k = k.unsqueeze(0) # [1, M, H, C]
+ v = v.unsqueeze(0) # [1, M, H, C]
+ mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+ out = xops.memory_efficient_attention(q, k, v, mask)[0] # [M, H, C]
+ elif ATTN == 'flash_attn':
+ cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
+ .to(qkv.device).int()
+ out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
+
+ out = out[bwd_indices] # [T, H, C]
+
+ if DEBUG:
+ qkv_coords = qkv_coords[bwd_indices]
+ assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch"
+
+ return qkv.replace(out)
diff --git a/Stable3DGen/hi3dgen/modules/sparse/attention/windowed_attn.py b/Stable3DGen/hi3dgen/modules/sparse/attention/windowed_attn.py
new file mode 100755
index 0000000000000000000000000000000000000000..642a2e282671080d8841296d3760919278ed994b
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/attention/windowed_attn.py
@@ -0,0 +1,159 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import math
+from .. import SparseTensor
+from .. import DEBUG, ATTN
+
+if ATTN == 'xformers':
+ import xformers.ops as xops
+elif ATTN == 'flash_attn':
+ import flash_attn
+else:
+ raise ValueError(f"Unknown attention module: {ATTN}")
+
+
+__all__ = [
+ 'sparse_windowed_scaled_dot_product_self_attention',
+]
+
+
+def calc_window_partition(
+ tensor: SparseTensor,
+ window_size: Union[int, Tuple[int, ...]],
+ shift_window: Union[int, Tuple[int, ...]] = 0
+) -> Tuple[torch.Tensor, torch.Tensor, List[int], List[int]]:
+ """
+ Calculate serialization and partitioning for a set of coordinates.
+
+ Args:
+ tensor (SparseTensor): The input tensor.
+ window_size (int): The window size to use.
+ shift_window (Tuple[int, ...]): The shift of serialized coordinates.
+
+ Returns:
+ (torch.Tensor): Forwards indices.
+ (torch.Tensor): Backwards indices.
+ (List[int]): Sequence lengths.
+ (List[int]): Sequence batch indices.
+ """
+ DIM = tensor.coords.shape[1] - 1
+ shift_window = (shift_window,) * DIM if isinstance(shift_window, int) else shift_window
+ window_size = (window_size,) * DIM if isinstance(window_size, int) else window_size
+ shifted_coords = tensor.coords.clone().detach()
+ shifted_coords[:, 1:] += torch.tensor(shift_window, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+
+ MAX_COORDS = shifted_coords[:, 1:].max(dim=0).values.tolist()
+ NUM_WINDOWS = [math.ceil((mc + 1) / ws) for mc, ws in zip(MAX_COORDS, window_size)]
+ OFFSET = torch.cumprod(torch.tensor([1] + NUM_WINDOWS[::-1]), dim=0).tolist()[::-1]
+
+ shifted_coords[:, 1:] //= torch.tensor(window_size, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+ shifted_indices = (shifted_coords * torch.tensor(OFFSET, device=tensor.device, dtype=torch.int32).unsqueeze(0)).sum(dim=1)
+ fwd_indices = torch.argsort(shifted_indices)
+ bwd_indices = torch.empty_like(fwd_indices)
+ bwd_indices[fwd_indices] = torch.arange(fwd_indices.shape[0], device=tensor.device)
+ seq_lens = torch.bincount(shifted_indices)
+ seq_batch_indices = torch.arange(seq_lens.shape[0], device=tensor.device, dtype=torch.int32) // OFFSET[0]
+ mask = seq_lens != 0
+ seq_lens = seq_lens[mask].tolist()
+ seq_batch_indices = seq_batch_indices[mask].tolist()
+
+ return fwd_indices, bwd_indices, seq_lens, seq_batch_indices
+
+
+def sparse_windowed_scaled_dot_product_self_attention(
+ qkv: SparseTensor,
+ window_size: int,
+ shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> SparseTensor:
+ """
+ Apply windowed scaled dot product self attention to a sparse tensor.
+
+ Args:
+ qkv (SparseTensor): [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+ window_size (int): The window size to use.
+ shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
+ shift (int): The shift to use.
+ """
+ assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+
+ serialization_spatial_cache_name = f'window_partition_{window_size}_{shift_window}'
+ serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name)
+ if serialization_spatial_cache is None:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = calc_window_partition(qkv, window_size, shift_window)
+ qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, seq_batch_indices))
+ else:
+ fwd_indices, bwd_indices, seq_lens, seq_batch_indices = serialization_spatial_cache
+
+ M = fwd_indices.shape[0]
+ T = qkv.feats.shape[0]
+ H = qkv.feats.shape[2]
+ C = qkv.feats.shape[3]
+
+ qkv_feats = qkv.feats[fwd_indices] # [M, 3, H, C]
+
+ if DEBUG:
+ start = 0
+ qkv_coords = qkv.coords[fwd_indices]
+ for i in range(len(seq_lens)):
+ seq_coords = qkv_coords[start:start+seq_lens[i]]
+ assert (seq_coords[:, 0] == seq_batch_indices[i]).all(), f"SparseWindowedScaledDotProductSelfAttention: batch index mismatch"
+ assert (seq_coords[:, 1:].max(dim=0).values - seq_coords[:, 1:].min(dim=0).values < window_size).all(), \
+ f"SparseWindowedScaledDotProductSelfAttention: window size exceeded"
+ start += seq_lens[i]
+
+ if all([seq_len == window_size for seq_len in seq_lens]):
+ B = len(seq_lens)
+ N = window_size
+ qkv_feats = qkv_feats.reshape(B, N, 3, H, C)
+ if ATTN == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=2) # [B, N, H, C]
+ out = xops.memory_efficient_attention(q, k, v) # [B, N, H, C]
+ elif ATTN == 'flash_attn':
+ out = flash_attn.flash_attn_qkvpacked_func(qkv_feats) # [B, N, H, C]
+ else:
+ raise ValueError(f"Unknown attention module: {ATTN}")
+ out = out.reshape(B * N, H, C) # [M, H, C]
+ else:
+ if ATTN == 'xformers':
+ q, k, v = qkv_feats.unbind(dim=1) # [M, H, C]
+ q = q.unsqueeze(0) # [1, M, H, C]
+ k = k.unsqueeze(0) # [1, M, H, C]
+ v = v.unsqueeze(0) # [1, M, H, C]
+ mask = xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+ out = xops.memory_efficient_attention(q, k, v, mask)[0] # [M, H, C]
+ elif ATTN == 'flash_attn':
+ cu_seqlens = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(seq_lens), dim=0)], dim=0) \
+ .to(qkv.device).int()
+ out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, cu_seqlens, max(seq_lens)) # [M, H, C]
+
+ out = out[bwd_indices] # [T, H, C]
+
+ if DEBUG:
+ qkv_coords = qkv_coords[bwd_indices]
+ assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch"
+
+ return qkv.replace(out)
diff --git a/Stable3DGen/hi3dgen/modules/sparse/basic.py b/Stable3DGen/hi3dgen/modules/sparse/basic.py
new file mode 100755
index 0000000000000000000000000000000000000000..3058f5c579852d8eb7cc79446f3db205d2d060df
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/basic.py
@@ -0,0 +1,483 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+from . import BACKEND, DEBUG
+SparseTensorData = None # Lazy import
+
+
+__all__ = [
+ 'SparseTensor',
+ 'sparse_batch_broadcast',
+ 'sparse_batch_op',
+ 'sparse_cat',
+ 'sparse_unbind',
+]
+
+
+class SparseTensor:
+ """
+ Sparse tensor with support for both torchsparse and spconv backends.
+
+ Parameters:
+ - feats (torch.Tensor): Features of the sparse tensor.
+ - coords (torch.Tensor): Coordinates of the sparse tensor.
+ - shape (torch.Size): Shape of the sparse tensor.
+ - layout (List[slice]): Layout of the sparse tensor for each batch
+ - data (SparseTensorData): Sparse tensor data used for convolusion
+
+ NOTE:
+ - Data corresponding to a same batch should be contiguous.
+ - Coords should be in [0, 1023]
+ """
+ @overload
+ def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ...
+
+ @overload
+ def __init__(self, data, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ...
+
+ def __init__(self, *args, **kwargs):
+ # Lazy import of sparse tensor backend
+ global SparseTensorData
+ if SparseTensorData is None:
+ import importlib
+ if BACKEND == 'torchsparse':
+ SparseTensorData = importlib.import_module('torchsparse').SparseTensor
+ elif BACKEND == 'spconv':
+ SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor
+
+ method_id = 0
+ if len(args) != 0:
+ method_id = 0 if isinstance(args[0], torch.Tensor) else 1
+ else:
+ method_id = 1 if 'data' in kwargs else 0
+
+ if method_id == 0:
+ feats, coords, shape, layout = args + (None,) * (4 - len(args))
+ if 'feats' in kwargs:
+ feats = kwargs['feats']
+ del kwargs['feats']
+ if 'coords' in kwargs:
+ coords = kwargs['coords']
+ del kwargs['coords']
+ if 'shape' in kwargs:
+ shape = kwargs['shape']
+ del kwargs['shape']
+ if 'layout' in kwargs:
+ layout = kwargs['layout']
+ del kwargs['layout']
+
+ if shape is None:
+ shape = self.__cal_shape(feats, coords)
+ if layout is None:
+ layout = self.__cal_layout(coords, shape[0])
+ if BACKEND == 'torchsparse':
+ self.data = SparseTensorData(feats, coords, **kwargs)
+ elif BACKEND == 'spconv':
+ spatial_shape = list(coords.max(0)[0] + 1)[1:]
+ self.data = SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape, shape[0], **kwargs)
+ self.data._features = feats
+ elif method_id == 1:
+ data, shape, layout = args + (None,) * (3 - len(args))
+ if 'data' in kwargs:
+ data = kwargs['data']
+ del kwargs['data']
+ if 'shape' in kwargs:
+ shape = kwargs['shape']
+ del kwargs['shape']
+ if 'layout' in kwargs:
+ layout = kwargs['layout']
+ del kwargs['layout']
+
+ self.data = data
+ if shape is None:
+ shape = self.__cal_shape(self.feats, self.coords)
+ if layout is None:
+ layout = self.__cal_layout(self.coords, shape[0])
+
+ self._shape = shape
+ self._layout = layout
+ self._scale = kwargs.get('scale', (1, 1, 1))
+ self._spatial_cache = kwargs.get('spatial_cache', {})
+
+ if DEBUG:
+ try:
+ assert self.feats.shape[0] == self.coords.shape[0], f"Invalid feats shape: {self.feats.shape}, coords shape: {self.coords.shape}"
+ assert self.shape == self.__cal_shape(self.feats, self.coords), f"Invalid shape: {self.shape}"
+ assert self.layout == self.__cal_layout(self.coords, self.shape[0]), f"Invalid layout: {self.layout}"
+ for i in range(self.shape[0]):
+ assert torch.all(self.coords[self.layout[i], 0] == i), f"The data of batch {i} is not contiguous"
+ except Exception as e:
+ print('Debugging information:')
+ print(f"- Shape: {self.shape}")
+ print(f"- Layout: {self.layout}")
+ print(f"- Scale: {self._scale}")
+ print(f"- Coords: {self.coords}")
+ raise e
+
+ def __cal_shape(self, feats, coords):
+ shape = []
+ shape.append(coords[:, 0].max().item() + 1)
+ shape.extend([*feats.shape[1:]])
+ return torch.Size(shape)
+
+ def __cal_layout(self, coords, batch_size):
+ seq_len = torch.bincount(coords[:, 0], minlength=batch_size)
+ offset = torch.cumsum(seq_len, dim=0)
+ layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)]
+ return layout
+
+ @property
+ def shape(self) -> torch.Size:
+ return self._shape
+
+ def dim(self) -> int:
+ return len(self.shape)
+
+ @property
+ def layout(self) -> List[slice]:
+ return self._layout
+
+ @property
+ def feats(self) -> torch.Tensor:
+ if BACKEND == 'torchsparse':
+ return self.data.F
+ elif BACKEND == 'spconv':
+ return self.data.features
+
+ @feats.setter
+ def feats(self, value: torch.Tensor):
+ if BACKEND == 'torchsparse':
+ self.data.F = value
+ elif BACKEND == 'spconv':
+ self.data.features = value
+
+ @property
+ def coords(self) -> torch.Tensor:
+ if BACKEND == 'torchsparse':
+ return self.data.C
+ elif BACKEND == 'spconv':
+ return self.data.indices
+
+ @coords.setter
+ def coords(self, value: torch.Tensor):
+ if BACKEND == 'torchsparse':
+ self.data.C = value
+ elif BACKEND == 'spconv':
+ self.data.indices = value
+
+ @property
+ def dtype(self):
+ return self.feats.dtype
+
+ @property
+ def device(self):
+ return self.feats.device
+
+ @overload
+ def to(self, dtype: torch.dtype) -> 'SparseTensor': ...
+
+ @overload
+ def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None) -> 'SparseTensor': ...
+
+ def to(self, *args, **kwargs) -> 'SparseTensor':
+ device = None
+ dtype = None
+ if len(args) == 2:
+ device, dtype = args
+ elif len(args) == 1:
+ if isinstance(args[0], torch.dtype):
+ dtype = args[0]
+ else:
+ device = args[0]
+ if 'dtype' in kwargs:
+ assert dtype is None, "to() received multiple values for argument 'dtype'"
+ dtype = kwargs['dtype']
+ if 'device' in kwargs:
+ assert device is None, "to() received multiple values for argument 'device'"
+ device = kwargs['device']
+
+ new_feats = self.feats.to(device=device, dtype=dtype)
+ new_coords = self.coords.to(device=device)
+ return self.replace(new_feats, new_coords)
+
+ def type(self, dtype):
+ new_feats = self.feats.type(dtype)
+ return self.replace(new_feats)
+
+ def cpu(self) -> 'SparseTensor':
+ new_feats = self.feats.cpu()
+ new_coords = self.coords.cpu()
+ return self.replace(new_feats, new_coords)
+
+ def cuda(self) -> 'SparseTensor':
+ new_feats = self.feats.cuda()
+ new_coords = self.coords.cuda()
+ return self.replace(new_feats, new_coords)
+
+ def half(self) -> 'SparseTensor':
+ new_feats = self.feats.half()
+ return self.replace(new_feats)
+
+ def float(self) -> 'SparseTensor':
+ new_feats = self.feats.float()
+ return self.replace(new_feats)
+
+ def detach(self) -> 'SparseTensor':
+ new_coords = self.coords.detach()
+ new_feats = self.feats.detach()
+ return self.replace(new_feats, new_coords)
+
+ def dense(self) -> torch.Tensor:
+ if BACKEND == 'torchsparse':
+ return self.data.dense()
+ elif BACKEND == 'spconv':
+ return self.data.dense()
+
+ def reshape(self, *shape) -> 'SparseTensor':
+ new_feats = self.feats.reshape(self.feats.shape[0], *shape)
+ return self.replace(new_feats)
+
+ def unbind(self, dim: int) -> List['SparseTensor']:
+ return sparse_unbind(self, dim)
+
+ def replace(self, feats: torch.Tensor, coords: Optional[torch.Tensor] = None) -> 'SparseTensor':
+ new_shape = [self.shape[0]]
+ new_shape.extend(feats.shape[1:])
+ if BACKEND == 'torchsparse':
+ new_data = SparseTensorData(
+ feats=feats,
+ coords=self.data.coords if coords is None else coords,
+ stride=self.data.stride,
+ spatial_range=self.data.spatial_range,
+ )
+ new_data._caches = self.data._caches
+ elif BACKEND == 'spconv':
+ new_data = SparseTensorData(
+ self.data.features.reshape(self.data.features.shape[0], -1),
+ self.data.indices,
+ self.data.spatial_shape,
+ self.data.batch_size,
+ self.data.grid,
+ self.data.voxel_num,
+ self.data.indice_dict
+ )
+ new_data._features = feats
+ new_data.benchmark = self.data.benchmark
+ new_data.benchmark_record = self.data.benchmark_record
+ new_data.thrust_allocator = self.data.thrust_allocator
+ new_data._timer = self.data._timer
+ new_data.force_algo = self.data.force_algo
+ new_data.int8_scale = self.data.int8_scale
+ if coords is not None:
+ new_data.indices = coords
+ new_tensor = SparseTensor(new_data, shape=torch.Size(new_shape), layout=self.layout, scale=self._scale, spatial_cache=self._spatial_cache)
+ return new_tensor
+
+ @staticmethod
+ def full(aabb, dim, value, dtype=torch.float32, device=None) -> 'SparseTensor':
+ N, C = dim
+ x = torch.arange(aabb[0], aabb[3] + 1)
+ y = torch.arange(aabb[1], aabb[4] + 1)
+ z = torch.arange(aabb[2], aabb[5] + 1)
+ coords = torch.stack(torch.meshgrid(x, y, z, indexing='ij'), dim=-1).reshape(-1, 3)
+ coords = torch.cat([
+ torch.arange(N).view(-1, 1).repeat(1, coords.shape[0]).view(-1, 1),
+ coords.repeat(N, 1),
+ ], dim=1).to(dtype=torch.int32, device=device)
+ feats = torch.full((coords.shape[0], C), value, dtype=dtype, device=device)
+ return SparseTensor(feats=feats, coords=coords)
+
+ def __merge_sparse_cache(self, other: 'SparseTensor') -> dict:
+ new_cache = {}
+ for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())):
+ if k in self._spatial_cache:
+ new_cache[k] = self._spatial_cache[k]
+ if k in other._spatial_cache:
+ if k not in new_cache:
+ new_cache[k] = other._spatial_cache[k]
+ else:
+ new_cache[k].update(other._spatial_cache[k])
+ return new_cache
+
+ def __neg__(self) -> 'SparseTensor':
+ return self.replace(-self.feats)
+
+ def __elemwise__(self, other: Union[torch.Tensor, 'SparseTensor'], op: callable) -> 'SparseTensor':
+ if isinstance(other, torch.Tensor):
+ try:
+ other = torch.broadcast_to(other, self.shape)
+ other = sparse_batch_broadcast(self, other)
+ except:
+ pass
+ if isinstance(other, SparseTensor):
+ other = other.feats
+ new_feats = op(self.feats, other)
+ new_tensor = self.replace(new_feats)
+ if isinstance(other, SparseTensor):
+ new_tensor._spatial_cache = self.__merge_sparse_cache(other)
+ return new_tensor
+
+ def __add__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.add)
+
+ def __radd__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.add)
+
+ def __sub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.sub)
+
+ def __rsub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, lambda x, y: torch.sub(y, x))
+
+ def __mul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.mul)
+
+ def __rmul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.mul)
+
+ def __truediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, torch.div)
+
+ def __rtruediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
+ return self.__elemwise__(other, lambda x, y: torch.div(y, x))
+
+ def __getitem__(self, idx):
+ if isinstance(idx, int):
+ idx = [idx]
+ elif isinstance(idx, slice):
+ idx = range(*idx.indices(self.shape[0]))
+ elif isinstance(idx, torch.Tensor):
+ if idx.dtype == torch.bool:
+ assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
+ idx = idx.nonzero().squeeze(1)
+ elif idx.dtype in [torch.int32, torch.int64]:
+ assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
+ else:
+ raise ValueError(f"Unknown index type: {idx.dtype}")
+ else:
+ raise ValueError(f"Unknown index type: {type(idx)}")
+
+ coords = []
+ feats = []
+ for new_idx, old_idx in enumerate(idx):
+ coords.append(self.coords[self.layout[old_idx]].clone())
+ coords[-1][:, 0] = new_idx
+ feats.append(self.feats[self.layout[old_idx]])
+ coords = torch.cat(coords, dim=0).contiguous()
+ feats = torch.cat(feats, dim=0).contiguous()
+ return SparseTensor(feats=feats, coords=coords)
+
+ def register_spatial_cache(self, key, value) -> None:
+ """
+ Register a spatial cache.
+ The spatial cache can be any thing you want to cache.
+ The registery and retrieval of the cache is based on current scale.
+ """
+ scale_key = str(self._scale)
+ if scale_key not in self._spatial_cache:
+ self._spatial_cache[scale_key] = {}
+ self._spatial_cache[scale_key][key] = value
+
+ def get_spatial_cache(self, key=None):
+ """
+ Get a spatial cache.
+ """
+ scale_key = str(self._scale)
+ cur_scale_cache = self._spatial_cache.get(scale_key, {})
+ if key is None:
+ return cur_scale_cache
+ return cur_scale_cache.get(key, None)
+
+
+def sparse_batch_broadcast(input: SparseTensor, other: torch.Tensor) -> torch.Tensor:
+ """
+ Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation.
+
+ Args:
+ input (torch.Tensor): 1D tensor to broadcast.
+ target (SparseTensor): Sparse tensor to broadcast to.
+ op (callable): Operation to perform after broadcasting. Defaults to torch.add.
+ """
+ coords, feats = input.coords, input.feats
+ broadcasted = torch.zeros_like(feats)
+ for k in range(input.shape[0]):
+ broadcasted[input.layout[k]] = other[k]
+ return broadcasted
+
+
+def sparse_batch_op(input: SparseTensor, other: torch.Tensor, op: callable = torch.add) -> SparseTensor:
+ """
+ Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation.
+
+ Args:
+ input (torch.Tensor): 1D tensor to broadcast.
+ target (SparseTensor): Sparse tensor to broadcast to.
+ op (callable): Operation to perform after broadcasting. Defaults to torch.add.
+ """
+ return input.replace(op(input.feats, sparse_batch_broadcast(input, other)))
+
+
+def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor:
+ """
+ Concatenate a list of sparse tensors.
+
+ Args:
+ inputs (List[SparseTensor]): List of sparse tensors to concatenate.
+ """
+ if dim == 0:
+ start = 0
+ coords = []
+ for input in inputs:
+ coords.append(input.coords.clone())
+ coords[-1][:, 0] += start
+ start += input.shape[0]
+ coords = torch.cat(coords, dim=0)
+ feats = torch.cat([input.feats for input in inputs], dim=0)
+ output = SparseTensor(
+ coords=coords,
+ feats=feats,
+ )
+ else:
+ feats = torch.cat([input.feats for input in inputs], dim=dim)
+ output = inputs[0].replace(feats)
+
+ return output
+
+
+def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]:
+ """
+ Unbind a sparse tensor along a dimension.
+
+ Args:
+ input (SparseTensor): Sparse tensor to unbind.
+ dim (int): Dimension to unbind.
+ """
+ if dim == 0:
+ return [input[i] for i in range(input.shape[0])]
+ else:
+ feats = input.feats.unbind(dim)
+ return [input.replace(f) for f in feats]
diff --git a/Stable3DGen/hi3dgen/modules/sparse/conv/__init__.py b/Stable3DGen/hi3dgen/modules/sparse/conv/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..b5dbc6b00222258b6257652d8aab29be4e60086b
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/conv/__init__.py
@@ -0,0 +1,45 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from .. import BACKEND
+
+
+SPCONV_ALGO = 'auto' # 'auto', 'implicit_gemm', 'native'
+
+def __from_env():
+ import os
+
+ global SPCONV_ALGO
+ env_spconv_algo = os.environ.get('SPCONV_ALGO')
+ if env_spconv_algo is not None and env_spconv_algo in ['auto', 'implicit_gemm', 'native']:
+ SPCONV_ALGO = env_spconv_algo
+ print(f"[SPARSE][CONV] spconv algo: {SPCONV_ALGO}")
+
+
+__from_env()
+
+if BACKEND == 'torchsparse':
+ from .conv_torchsparse import *
+elif BACKEND == 'spconv':
+ from .conv_spconv import *
diff --git a/Stable3DGen/hi3dgen/modules/sparse/conv/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/conv/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..157e07fb1274588078a964865d4402b28e99e525
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/conv/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/conv/__pycache__/conv_spconv.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/conv/__pycache__/conv_spconv.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4cb9608b321e80c7df73f364607f8c3b2939cea8
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/conv/__pycache__/conv_spconv.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/conv/conv_spconv.py b/Stable3DGen/hi3dgen/modules/sparse/conv/conv_spconv.py
new file mode 100755
index 0000000000000000000000000000000000000000..2248b9cc186ab95d1140dfba10107ff9aca6b216
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/conv/conv_spconv.py
@@ -0,0 +1,104 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) Microsoft
+# SPDX-License-Identifier: MIT
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+from .. import DEBUG
+from . import SPCONV_ALGO
+
+class SparseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+ super(SparseConv3d, self).__init__()
+ if 'spconv' not in globals():
+ import spconv.pytorch as spconv
+ algo = None
+ if SPCONV_ALGO == 'native':
+ algo = spconv.ConvAlgo.Native
+ elif SPCONV_ALGO == 'implicit_gemm':
+ algo = spconv.ConvAlgo.MaskImplicitGemm
+ if stride == 1 and (padding is None):
+ self.conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias, indice_key=indice_key, algo=algo)
+ else:
+ self.conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, indice_key=indice_key, algo=algo)
+ self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
+ self.padding = padding
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ spatial_changed = any(s != 1 for s in self.stride) or (self.padding is not None)
+ new_data = self.conv(x.data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ new_layout = None if spatial_changed else x.layout
+
+ if spatial_changed and (x.shape[0] != 1):
+ # spconv was non-1 stride will break the contiguous of the output tensor, sort by the coords
+ fwd = new_data.indices[:, 0].argsort()
+ bwd = torch.zeros_like(fwd).scatter_(0, fwd, torch.arange(fwd.shape[0], device=fwd.device))
+ sorted_feats = new_data.features[fwd]
+ sorted_coords = new_data.indices[fwd]
+ unsorted_data = new_data
+ new_data = spconv.SparseConvTensor(sorted_feats, sorted_coords, unsorted_data.spatial_shape, unsorted_data.batch_size) # type: ignore
+
+ out = SparseTensor(
+ new_data, shape=torch.Size(new_shape), layout=new_layout,
+ scale=tuple([s * stride for s, stride in zip(x._scale, self.stride)]),
+ spatial_cache=x._spatial_cache,
+ )
+
+ if spatial_changed and (x.shape[0] != 1):
+ out.register_spatial_cache(f'conv_{self.stride}_unsorted_data', unsorted_data)
+ out.register_spatial_cache(f'conv_{self.stride}_sort_bwd', bwd)
+
+ return out
+
+
+class SparseInverseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+ super(SparseInverseConv3d, self).__init__()
+ if 'spconv' not in globals():
+ import spconv.pytorch as spconv
+ self.conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, bias=bias, indice_key=indice_key)
+ self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ spatial_changed = any(s != 1 for s in self.stride)
+ if spatial_changed:
+ # recover the original spconv order
+ data = x.get_spatial_cache(f'conv_{self.stride}_unsorted_data')
+ bwd = x.get_spatial_cache(f'conv_{self.stride}_sort_bwd')
+ data = data.replace_feature(x.feats[bwd])
+ if DEBUG:
+ assert torch.equal(data.indices, x.coords[bwd]), 'Recover the original order failed'
+ else:
+ data = x.data
+
+ new_data = self.conv(data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ new_layout = None if spatial_changed else x.layout
+ out = SparseTensor(
+ new_data, shape=torch.Size(new_shape), layout=new_layout,
+ scale=tuple([s // stride for s, stride in zip(x._scale, self.stride)]),
+ spatial_cache=x._spatial_cache,
+ )
+ return out
diff --git a/Stable3DGen/hi3dgen/modules/sparse/conv/conv_torchsparse.py b/Stable3DGen/hi3dgen/modules/sparse/conv/conv_torchsparse.py
new file mode 100755
index 0000000000000000000000000000000000000000..8861042eb5b1ea066a22b5ab370ca50033dfdba9
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/conv/conv_torchsparse.py
@@ -0,0 +1,62 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+
+
+class SparseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+ super(SparseConv3d, self).__init__()
+ if 'torchsparse' not in globals():
+ import torchsparse
+ self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias)
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ out = self.conv(x.data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None)
+ out._spatial_cache = x._spatial_cache
+ out._scale = tuple([s * stride for s, stride in zip(x._scale, self.conv.stride)])
+ return out
+
+
+class SparseInverseConv3d(nn.Module):
+ def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+ super(SparseInverseConv3d, self).__init__()
+ if 'torchsparse' not in globals():
+ import torchsparse
+ self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias, transposed=True)
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ out = self.conv(x.data)
+ new_shape = [x.shape[0], self.conv.out_channels]
+ out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None)
+ out._spatial_cache = x._spatial_cache
+ out._scale = tuple([s // stride for s, stride in zip(x._scale, self.conv.stride)])
+ return out
+
+
+
diff --git a/Stable3DGen/hi3dgen/modules/sparse/linear.py b/Stable3DGen/hi3dgen/modules/sparse/linear.py
new file mode 100755
index 0000000000000000000000000000000000000000..d055ee426ad2372b31df05800813f8ea2c898695
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/linear.py
@@ -0,0 +1,39 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch
+import torch.nn as nn
+from . import SparseTensor
+
+__all__ = [
+ 'SparseLinear'
+]
+
+
+class SparseLinear(nn.Linear):
+ def __init__(self, in_features, out_features, bias=True):
+ super(SparseLinear, self).__init__(in_features, out_features, bias)
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
diff --git a/Stable3DGen/hi3dgen/modules/sparse/nonlinearity.py b/Stable3DGen/hi3dgen/modules/sparse/nonlinearity.py
new file mode 100755
index 0000000000000000000000000000000000000000..3329f84452043995ae1455d583c3257310bd17cf
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/nonlinearity.py
@@ -0,0 +1,59 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch
+import torch.nn as nn
+from . import SparseTensor
+
+__all__ = [
+ 'SparseReLU',
+ 'SparseSiLU',
+ 'SparseGELU',
+ 'SparseActivation'
+]
+
+
+class SparseReLU(nn.ReLU):
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
+
+
+class SparseSiLU(nn.SiLU):
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
+
+
+class SparseGELU(nn.GELU):
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(super().forward(input.feats))
+
+
+class SparseActivation(nn.Module):
+ def __init__(self, activation: nn.Module):
+ super().__init__()
+ self.activation = activation
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ return input.replace(self.activation(input.feats))
+
diff --git a/Stable3DGen/hi3dgen/modules/sparse/norm.py b/Stable3DGen/hi3dgen/modules/sparse/norm.py
new file mode 100755
index 0000000000000000000000000000000000000000..1fa033d91c51f34bb9bab31842aa850bb1421665
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/norm.py
@@ -0,0 +1,82 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch
+import torch.nn as nn
+from . import SparseTensor
+from . import DEBUG
+
+__all__ = [
+ 'SparseGroupNorm',
+ 'SparseLayerNorm',
+ 'SparseGroupNorm32',
+ 'SparseLayerNorm32',
+]
+
+
+class SparseGroupNorm(nn.GroupNorm):
+ def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
+ super(SparseGroupNorm, self).__init__(num_groups, num_channels, eps, affine)
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ nfeats = torch.zeros_like(input.feats)
+ for k in range(input.shape[0]):
+ if DEBUG:
+ assert (input.coords[input.layout[k], 0] == k).all(), f"SparseGroupNorm: batch index mismatch"
+ bfeats = input.feats[input.layout[k]]
+ bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
+ bfeats = super().forward(bfeats)
+ bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
+ nfeats[input.layout[k]] = bfeats
+ return input.replace(nfeats)
+
+
+class SparseLayerNorm(nn.LayerNorm):
+ def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+ super(SparseLayerNorm, self).__init__(normalized_shape, eps, elementwise_affine)
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ nfeats = torch.zeros_like(input.feats)
+ for k in range(input.shape[0]):
+ bfeats = input.feats[input.layout[k]]
+ bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
+ bfeats = super().forward(bfeats)
+ bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
+ nfeats[input.layout[k]] = bfeats
+ return input.replace(nfeats)
+
+
+class SparseGroupNorm32(SparseGroupNorm):
+ """
+ A GroupNorm layer that converts to float32 before the forward pass.
+ """
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ return super().forward(x.float()).type(x.dtype)
+
+class SparseLayerNorm32(SparseLayerNorm):
+ """
+ A LayerNorm layer that converts to float32 before the forward pass.
+ """
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ return super().forward(x.float()).type(x.dtype)
diff --git a/Stable3DGen/hi3dgen/modules/sparse/spatial.py b/Stable3DGen/hi3dgen/modules/sparse/spatial.py
new file mode 100755
index 0000000000000000000000000000000000000000..1fefde1b38298fd96b02eb1a64ba9f1f8a8bc908
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/spatial.py
@@ -0,0 +1,134 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+from . import SparseTensor
+
+__all__ = [
+ 'SparseDownsample',
+ 'SparseUpsample',
+ 'SparseSubdivide'
+]
+
+
+class SparseDownsample(nn.Module):
+ """
+ Downsample a sparse tensor by a factor of `factor`.
+ Implemented as average pooling.
+ """
+ def __init__(self, factor: Union[int, Tuple[int, ...], List[int]]):
+ super(SparseDownsample, self).__init__()
+ self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ DIM = input.coords.shape[-1] - 1
+ factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM
+ assert DIM == len(factor), 'Input coordinates must have the same dimension as the downsample factor.'
+
+ coord = list(input.coords.unbind(dim=-1))
+ for i, f in enumerate(factor):
+ coord[i+1] = coord[i+1] // f
+
+ MAX = [coord[i+1].max().item() + 1 for i in range(DIM)]
+ OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1]
+ code = sum([c * o for c, o in zip(coord, OFFSET)])
+ code, idx = code.unique(return_inverse=True)
+
+ new_feats = torch.scatter_reduce(
+ torch.zeros(code.shape[0], input.feats.shape[1], device=input.feats.device, dtype=input.feats.dtype),
+ dim=0,
+ index=idx.unsqueeze(1).expand(-1, input.feats.shape[1]),
+ src=input.feats,
+ reduce='mean'
+ )
+ new_coords = torch.stack(
+ [code // OFFSET[0]] +
+ [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)],
+ dim=-1
+ )
+ out = SparseTensor(new_feats, new_coords, input.shape,)
+ out._scale = tuple([s // f for s, f in zip(input._scale, factor)])
+ out._spatial_cache = input._spatial_cache
+
+ out.register_spatial_cache(f'upsample_{factor}_coords', input.coords)
+ out.register_spatial_cache(f'upsample_{factor}_layout', input.layout)
+ out.register_spatial_cache(f'upsample_{factor}_idx', idx)
+
+ return out
+
+
+class SparseUpsample(nn.Module):
+ """
+ Upsample a sparse tensor by a factor of `factor`.
+ Implemented as nearest neighbor interpolation.
+ """
+ def __init__(self, factor: Union[int, Tuple[int, int, int], List[int]]):
+ super(SparseUpsample, self).__init__()
+ self.factor = tuple(factor) if isinstance(factor, (list, tuple)) else factor
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ DIM = input.coords.shape[-1] - 1
+ factor = self.factor if isinstance(self.factor, tuple) else (self.factor,) * DIM
+ assert DIM == len(factor), 'Input coordinates must have the same dimension as the upsample factor.'
+
+ new_coords = input.get_spatial_cache(f'upsample_{factor}_coords')
+ new_layout = input.get_spatial_cache(f'upsample_{factor}_layout')
+ idx = input.get_spatial_cache(f'upsample_{factor}_idx')
+ if any([x is None for x in [new_coords, new_layout, idx]]):
+ raise ValueError('Upsample cache not found. SparseUpsample must be paired with SparseDownsample.')
+ new_feats = input.feats[idx]
+ out = SparseTensor(new_feats, new_coords, input.shape, new_layout)
+ out._scale = tuple([s * f for s, f in zip(input._scale, factor)])
+ out._spatial_cache = input._spatial_cache
+ return out
+
+class SparseSubdivide(nn.Module):
+ """
+ Upsample a sparse tensor by a factor of `factor`.
+ Implemented as nearest neighbor interpolation.
+ """
+ def __init__(self):
+ super(SparseSubdivide, self).__init__()
+
+ def forward(self, input: SparseTensor) -> SparseTensor:
+ DIM = input.coords.shape[-1] - 1
+ # upsample scale=2^DIM
+ n_cube = torch.ones([2] * DIM, device=input.device, dtype=torch.int)
+ n_coords = torch.nonzero(n_cube)
+ n_coords = torch.cat([torch.zeros_like(n_coords[:, :1]), n_coords], dim=-1)
+ factor = n_coords.shape[0]
+ assert factor == 2 ** DIM
+ # print(n_coords.shape)
+ new_coords = input.coords.clone()
+ new_coords[:, 1:] *= 2
+ new_coords = new_coords.unsqueeze(1) + n_coords.unsqueeze(0).to(new_coords.dtype)
+
+ new_feats = input.feats.unsqueeze(1).expand(input.feats.shape[0], factor, *input.feats.shape[1:])
+ out = SparseTensor(new_feats.flatten(0, 1), new_coords.flatten(0, 1), input.shape)
+ out._scale = input._scale * 2
+ out._spatial_cache = input._spatial_cache
+ return out
+
diff --git a/Stable3DGen/hi3dgen/modules/sparse/transformer/__init__.py b/Stable3DGen/hi3dgen/modules/sparse/transformer/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..758e2d4bd9ccec609f6aed22773925572f56dea0
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/transformer/__init__.py
@@ -0,0 +1,26 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from .blocks import *
+from .modulated import *
\ No newline at end of file
diff --git a/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..20f2965af371bc6a5088ce99286a5b2321a95973
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/blocks.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/blocks.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ff1739f6a2b34721899b19a81509244d8aa0f192
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/blocks.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/modulated.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/modulated.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6dd1bf05b76ee7f78a65412a5e37821e0d0631f2
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/sparse/transformer/__pycache__/modulated.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/sparse/transformer/blocks.py b/Stable3DGen/hi3dgen/modules/sparse/transformer/blocks.py
new file mode 100755
index 0000000000000000000000000000000000000000..e6bf0867a429b46ad79d8719c3d125d12105396a
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/transformer/blocks.py
@@ -0,0 +1,175 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+from ..basic import SparseTensor
+from ..linear import SparseLinear
+from ..nonlinearity import SparseGELU
+from ..attention import SparseMultiHeadAttention, SerializeMode
+from ...norm import LayerNorm32
+
+
+class SparseFeedForwardNet(nn.Module):
+ def __init__(self, channels: int, mlp_ratio: float = 4.0):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ SparseLinear(channels, int(channels * mlp_ratio)),
+ SparseGELU(approximate="tanh"),
+ SparseLinear(int(channels * mlp_ratio), channels),
+ )
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ return self.mlp(x)
+
+
+class SparseTransformerBlock(nn.Module):
+ """
+ Sparse Transformer block (MSA + FFN).
+ """
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
+ window_size: Optional[int] = None,
+ shift_sequence: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ serialize_mode: Optional[SerializeMode] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qkv_bias: bool = True,
+ ln_affine: bool = False,
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.attn = SparseMultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_sequence=shift_sequence,
+ shift_window=shift_window,
+ serialize_mode=serialize_mode,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.mlp = SparseFeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+
+ def _forward(self, x: SparseTensor) -> SparseTensor:
+ h = x.replace(self.norm1(x.feats))
+ h = self.attn(h)
+ x = x + h
+ h = x.replace(self.norm2(x.feats))
+ h = self.mlp(h)
+ x = x + h
+ return x
+
+ def forward(self, x: SparseTensor) -> SparseTensor:
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+ else:
+ return self._forward(x)
+
+
+class SparseTransformerCrossBlock(nn.Module):
+ """
+ Sparse Transformer cross-attention block (MSA + MCA + FFN).
+ """
+ def __init__(
+ self,
+ channels: int,
+ ctx_channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
+ window_size: Optional[int] = None,
+ shift_sequence: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ serialize_mode: Optional[SerializeMode] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qk_rms_norm_cross: bool = False,
+ qkv_bias: bool = True,
+ ln_affine: bool = False,
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.self_attn = SparseMultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ type="self",
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_sequence=shift_sequence,
+ shift_window=shift_window,
+ serialize_mode=serialize_mode,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.cross_attn = SparseMultiHeadAttention(
+ channels,
+ ctx_channels=ctx_channels,
+ num_heads=num_heads,
+ type="cross",
+ attn_mode="full",
+ qkv_bias=qkv_bias,
+ qk_rms_norm=qk_rms_norm_cross,
+ )
+ self.mlp = SparseFeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+
+ def _forward(self, x: SparseTensor, mod: torch.Tensor, context: torch.Tensor):
+ h = x.replace(self.norm1(x.feats))
+ h = self.self_attn(h)
+ x = x + h
+ h = x.replace(self.norm2(x.feats))
+ h = self.cross_attn(h, context)
+ x = x + h
+ h = x.replace(self.norm3(x.feats))
+ h = self.mlp(h)
+ x = x + h
+ return x
+
+ def forward(self, x: SparseTensor, context: torch.Tensor):
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, context, use_reentrant=False)
+ else:
+ return self._forward(x, context)
diff --git a/Stable3DGen/hi3dgen/modules/sparse/transformer/modulated.py b/Stable3DGen/hi3dgen/modules/sparse/transformer/modulated.py
new file mode 100755
index 0000000000000000000000000000000000000000..e57eb16e798bc826d667394b8888a1bd4532e42d
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/sparse/transformer/modulated.py
@@ -0,0 +1,190 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+from ..basic import SparseTensor
+from ..attention import SparseMultiHeadAttention, SerializeMode
+from ...norm import LayerNorm32
+from .blocks import SparseFeedForwardNet
+
+
+class ModulatedSparseTransformerBlock(nn.Module):
+ """
+ Sparse Transformer block (MSA + FFN) with adaptive layer norm conditioning.
+ """
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
+ window_size: Optional[int] = None,
+ shift_sequence: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ serialize_mode: Optional[SerializeMode] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qkv_bias: bool = True,
+ share_mod: bool = False,
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.share_mod = share_mod
+ self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.attn = SparseMultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_sequence=shift_sequence,
+ shift_window=shift_window,
+ serialize_mode=serialize_mode,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.mlp = SparseFeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+ if not share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(channels, 6 * channels, bias=True)
+ )
+
+ def _forward(self, x: SparseTensor, mod: torch.Tensor) -> SparseTensor:
+ if self.share_mod:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+ else:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+ h = x.replace(self.norm1(x.feats))
+ h = h * (1 + scale_msa) + shift_msa
+ h = self.attn(h)
+ h = h * gate_msa
+ x = x + h
+ h = x.replace(self.norm2(x.feats))
+ h = h * (1 + scale_mlp) + shift_mlp
+ h = self.mlp(h)
+ h = h * gate_mlp
+ x = x + h
+ return x
+
+ def forward(self, x: SparseTensor, mod: torch.Tensor) -> SparseTensor:
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, mod, use_reentrant=False)
+ else:
+ return self._forward(x, mod)
+
+
+class ModulatedSparseTransformerCrossBlock(nn.Module):
+ """
+ Sparse Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
+ """
+ def __init__(
+ self,
+ channels: int,
+ ctx_channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
+ window_size: Optional[int] = None,
+ shift_sequence: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ serialize_mode: Optional[SerializeMode] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qk_rms_norm_cross: bool = False,
+ qkv_bias: bool = True,
+ share_mod: bool = False,
+
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.share_mod = share_mod
+ self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+ self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.self_attn = SparseMultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ type="self",
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_sequence=shift_sequence,
+ shift_window=shift_window,
+ serialize_mode=serialize_mode,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.cross_attn = SparseMultiHeadAttention(
+ channels,
+ ctx_channels=ctx_channels,
+ num_heads=num_heads,
+ type="cross",
+ attn_mode="full",
+ qkv_bias=qkv_bias,
+ qk_rms_norm=qk_rms_norm_cross,
+ )
+ self.mlp = SparseFeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+ if not share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(channels, 6 * channels, bias=True)
+ )
+
+ def _forward(self, x: SparseTensor, mod: torch.Tensor, context: torch.Tensor) -> SparseTensor:
+ if self.share_mod:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+ else:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+ h = x.replace(self.norm1(x.feats))
+ h = h * (1 + scale_msa) + shift_msa
+ h = self.self_attn(h)
+ h = h * gate_msa
+ x = x + h
+ h = x.replace(self.norm2(x.feats))
+ h = self.cross_attn(h, context)
+ x = x + h
+ h = x.replace(self.norm3(x.feats))
+ h = h * (1 + scale_mlp) + shift_mlp
+ h = self.mlp(h)
+ h = h * gate_mlp
+ x = x + h
+ return x
+
+ def forward(self, x: SparseTensor, mod: torch.Tensor, context: torch.Tensor) -> SparseTensor:
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, use_reentrant=False)
+ else:
+ return self._forward(x, mod, context)
diff --git a/Stable3DGen/hi3dgen/modules/spatial.py b/Stable3DGen/hi3dgen/modules/spatial.py
new file mode 100755
index 0000000000000000000000000000000000000000..514ae6362cacd83643eba1c57b017dcbf0194d91
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/spatial.py
@@ -0,0 +1,72 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch
+
+
+def pixel_shuffle_3d(x: torch.Tensor, scale_factor: int) -> torch.Tensor:
+ """
+ 3D pixel shuffle.
+ """
+ B, C, H, W, D = x.shape
+ C_ = C // scale_factor**3
+ x = x.reshape(B, C_, scale_factor, scale_factor, scale_factor, H, W, D)
+ x = x.permute(0, 1, 5, 2, 6, 3, 7, 4)
+ x = x.reshape(B, C_, H*scale_factor, W*scale_factor, D*scale_factor)
+ return x
+
+
+def patchify(x: torch.Tensor, patch_size: int):
+ """
+ Patchify a tensor.
+
+ Args:
+ x (torch.Tensor): (N, C, *spatial) tensor
+ patch_size (int): Patch size
+ """
+ DIM = x.dim() - 2
+ for d in range(2, DIM + 2):
+ assert x.shape[d] % patch_size == 0, f"Dimension {d} of input tensor must be divisible by patch size, got {x.shape[d]} and {patch_size}"
+
+ x = x.reshape(*x.shape[:2], *sum([[x.shape[d] // patch_size, patch_size] for d in range(2, DIM + 2)], []))
+ x = x.permute(0, 1, *([2 * i + 3 for i in range(DIM)] + [2 * i + 2 for i in range(DIM)]))
+ x = x.reshape(x.shape[0], x.shape[1] * (patch_size ** DIM), *(x.shape[-DIM:]))
+ return x
+
+
+def unpatchify(x: torch.Tensor, patch_size: int):
+ """
+ Unpatchify a tensor.
+
+ Args:
+ x (torch.Tensor): (N, C, *spatial) tensor
+ patch_size (int): Patch size
+ """
+ DIM = x.dim() - 2
+ assert x.shape[1] % (patch_size ** DIM) == 0, f"Second dimension of input tensor must be divisible by patch size to unpatchify, got {x.shape[1]} and {patch_size ** DIM}"
+
+ x = x.reshape(x.shape[0], x.shape[1] // (patch_size ** DIM), *([patch_size] * DIM), *(x.shape[-DIM:]))
+ x = x.permute(0, 1, *(sum([[2 + DIM + i, 2 + i] for i in range(DIM)], [])))
+ x = x.reshape(x.shape[0], x.shape[1], *[x.shape[2 + 2 * i] * patch_size for i in range(DIM)])
+ return x
diff --git a/Stable3DGen/hi3dgen/modules/transformer/__init__.py b/Stable3DGen/hi3dgen/modules/transformer/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..758e2d4bd9ccec609f6aed22773925572f56dea0
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/transformer/__init__.py
@@ -0,0 +1,26 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from .blocks import *
+from .modulated import *
\ No newline at end of file
diff --git a/Stable3DGen/hi3dgen/modules/transformer/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/transformer/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6a60f60f03729e3c1cb9963927b2d23769db78c6
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/transformer/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/transformer/__pycache__/blocks.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/transformer/__pycache__/blocks.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dc6c4851df84d5ce99b16b896b60beff6788aa50
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/transformer/__pycache__/blocks.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/transformer/__pycache__/modulated.cpython-310.pyc b/Stable3DGen/hi3dgen/modules/transformer/__pycache__/modulated.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9f7ebaf2e68153540f2d04e9aea5e2c1e6aca506
Binary files /dev/null and b/Stable3DGen/hi3dgen/modules/transformer/__pycache__/modulated.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/modules/transformer/blocks.py b/Stable3DGen/hi3dgen/modules/transformer/blocks.py
new file mode 100755
index 0000000000000000000000000000000000000000..0fac7d781858a5cb14e20224c14ebc95319a9690
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/transformer/blocks.py
@@ -0,0 +1,206 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+from ..attention import MultiHeadAttention
+from ..norm import LayerNorm32
+
+
+class AbsolutePositionEmbedder(nn.Module):
+ """
+ Embeds spatial positions into vector representations.
+ """
+ def __init__(self, channels: int, in_channels: int = 3):
+ super().__init__()
+ self.channels = channels
+ self.in_channels = in_channels
+ self.freq_dim = channels // in_channels // 2
+ self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+ self.freqs = 1.0 / (10000 ** self.freqs)
+
+ def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Create sinusoidal position embeddings.
+
+ Args:
+ x: a 1-D Tensor of N indices
+
+ Returns:
+ an (N, D) Tensor of positional embeddings.
+ """
+ self.freqs = self.freqs.to(x.device)
+ out = torch.outer(x, self.freqs)
+ out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1)
+ return out
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ x (torch.Tensor): (N, D) tensor of spatial positions
+ """
+ N, D = x.shape
+ assert D == self.in_channels, "Input dimension must match number of input channels"
+ embed = self._sin_cos_embedding(x.reshape(-1))
+ embed = embed.reshape(N, -1)
+ if embed.shape[1] < self.channels:
+ embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1], device=embed.device)], dim=-1)
+ return embed
+
+
+class FeedForwardNet(nn.Module):
+ def __init__(self, channels: int, mlp_ratio: float = 4.0):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(channels, int(channels * mlp_ratio)),
+ nn.GELU(approximate="tanh"),
+ nn.Linear(int(channels * mlp_ratio), channels),
+ )
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.mlp(x)
+
+
+class TransformerBlock(nn.Module):
+ """
+ Transformer block (MSA + FFN).
+ """
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_window: Optional[int] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qkv_bias: bool = True,
+ ln_affine: bool = False,
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.attn = MultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_window=shift_window,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.mlp = FeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+
+ def _forward(self, x: torch.Tensor) -> torch.Tensor:
+ h = self.norm1(x)
+ h = self.attn(h)
+ x = x + h
+ h = self.norm2(x)
+ h = self.mlp(h)
+ x = x + h
+ return x
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+ else:
+ return self._forward(x)
+
+
+class TransformerCrossBlock(nn.Module):
+ """
+ Transformer cross-attention block (MSA + MCA + FFN).
+ """
+ def __init__(
+ self,
+ channels: int,
+ ctx_channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qk_rms_norm_cross: bool = False,
+ qkv_bias: bool = True,
+ ln_affine: bool = False,
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+ self.self_attn = MultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ type="self",
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_window=shift_window,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.cross_attn = MultiHeadAttention(
+ channels,
+ ctx_channels=ctx_channels,
+ num_heads=num_heads,
+ type="cross",
+ attn_mode="full",
+ qkv_bias=qkv_bias,
+ qk_rms_norm=qk_rms_norm_cross,
+ )
+ self.mlp = FeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+
+ def _forward(self, x: torch.Tensor, context: torch.Tensor):
+ h = self.norm1(x)
+ h = self.self_attn(h)
+ x = x + h
+ h = self.norm2(x)
+ h = self.cross_attn(h, context)
+ x = x + h
+ h = self.norm3(x)
+ h = self.mlp(h)
+ x = x + h
+ return x
+
+ def forward(self, x: torch.Tensor, context: torch.Tensor):
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, context, use_reentrant=False)
+ else:
+ return self._forward(x, context)
+
\ No newline at end of file
diff --git a/Stable3DGen/hi3dgen/modules/transformer/modulated.py b/Stable3DGen/hi3dgen/modules/transformer/modulated.py
new file mode 100755
index 0000000000000000000000000000000000000000..6b7ae66aa399f5bbf52c7bc93b60eaa9e8041c52
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/transformer/modulated.py
@@ -0,0 +1,181 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+from ..attention import MultiHeadAttention
+from ..norm import LayerNorm32
+from .blocks import FeedForwardNet
+
+
+class ModulatedTransformerBlock(nn.Module):
+ """
+ Transformer block (MSA + FFN) with adaptive layer norm conditioning.
+ """
+ def __init__(
+ self,
+ channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qkv_bias: bool = True,
+ share_mod: bool = False,
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.share_mod = share_mod
+ self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.attn = MultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_window=shift_window,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.mlp = FeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+ if not share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(channels, 6 * channels, bias=True)
+ )
+
+ def _forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
+ if self.share_mod:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+ else:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+ h = self.norm1(x)
+ h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+ h = self.attn(h)
+ h = h * gate_msa.unsqueeze(1)
+ x = x + h
+ h = self.norm2(x)
+ h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+ h = self.mlp(h)
+ h = h * gate_mlp.unsqueeze(1)
+ x = x + h
+ return x
+
+ def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, mod, use_reentrant=False)
+ else:
+ return self._forward(x, mod)
+
+
+class ModulatedTransformerCrossBlock(nn.Module):
+ """
+ Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
+ """
+ def __init__(
+ self,
+ channels: int,
+ ctx_channels: int,
+ num_heads: int,
+ mlp_ratio: float = 4.0,
+ attn_mode: Literal["full", "windowed"] = "full",
+ window_size: Optional[int] = None,
+ shift_window: Optional[Tuple[int, int, int]] = None,
+ use_checkpoint: bool = False,
+ use_rope: bool = False,
+ qk_rms_norm: bool = False,
+ qk_rms_norm_cross: bool = False,
+ qkv_bias: bool = True,
+ share_mod: bool = False,
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.share_mod = share_mod
+ self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+ self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+ self.self_attn = MultiHeadAttention(
+ channels,
+ num_heads=num_heads,
+ type="self",
+ attn_mode=attn_mode,
+ window_size=window_size,
+ shift_window=shift_window,
+ qkv_bias=qkv_bias,
+ use_rope=use_rope,
+ qk_rms_norm=qk_rms_norm,
+ )
+ self.cross_attn = MultiHeadAttention(
+ channels,
+ ctx_channels=ctx_channels,
+ num_heads=num_heads,
+ type="cross",
+ attn_mode="full",
+ qkv_bias=qkv_bias,
+ qk_rms_norm=qk_rms_norm_cross,
+ )
+ self.mlp = FeedForwardNet(
+ channels,
+ mlp_ratio=mlp_ratio,
+ )
+ if not share_mod:
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(channels, 6 * channels, bias=True)
+ )
+
+ def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):
+ if self.share_mod:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mod.chunk(6, dim=1)
+ else:
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+ h = self.norm1(x)
+ h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+ h = self.self_attn(h)
+ h = h * gate_msa.unsqueeze(1)
+ x = x + h
+ h = self.norm2(x)
+ h = self.cross_attn(h, context)
+ x = x + h
+ h = self.norm3(x)
+ h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+ h = self.mlp(h)
+ h = h * gate_mlp.unsqueeze(1)
+ x = x + h
+ return x
+
+ def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor):
+ if self.use_checkpoint:
+ return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, use_reentrant=False)
+ else:
+ return self._forward(x, mod, context)
+
\ No newline at end of file
diff --git a/Stable3DGen/hi3dgen/modules/utils.py b/Stable3DGen/hi3dgen/modules/utils.py
new file mode 100755
index 0000000000000000000000000000000000000000..2e03f85ff2d7ccedcd5f893ecd5739738c028c98
--- /dev/null
+++ b/Stable3DGen/hi3dgen/modules/utils.py
@@ -0,0 +1,78 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch.nn as nn
+from ..modules import sparse as sp
+
+FP16_MODULES = (
+ nn.Conv1d,
+ nn.Conv2d,
+ nn.Conv3d,
+ nn.ConvTranspose1d,
+ nn.ConvTranspose2d,
+ nn.ConvTranspose3d,
+ nn.Linear,
+ sp.SparseConv3d,
+ sp.SparseInverseConv3d,
+ sp.SparseLinear,
+)
+
+def convert_module_to_f16(l):
+ """
+ Convert primitive modules to float16.
+ """
+ if isinstance(l, FP16_MODULES):
+ for p in l.parameters():
+ p.data = p.data.half()
+
+
+def convert_module_to_f32(l):
+ """
+ Convert primitive modules to float32, undoing convert_module_to_f16().
+ """
+ if isinstance(l, FP16_MODULES):
+ for p in l.parameters():
+ p.data = p.data.float()
+
+
+def zero_module(module):
+ """
+ Zero out the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().zero_()
+ return module
+
+
+def scale_module(module, scale):
+ """
+ Scale the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().mul_(scale)
+ return module
+
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
diff --git a/Stable3DGen/hi3dgen/pipelines/__init__.py b/Stable3DGen/hi3dgen/pipelines/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..88901f7c21ced9e89d783fa57baf02a328e4c908
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/__init__.py
@@ -0,0 +1,47 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from . import samplers
+from .hi3dgen import Hi3DGenPipeline
+
+def from_pretrained(path: str):
+ """
+ Load a pipeline from a model folder or a Hugging Face model hub.
+
+ Args:
+ path: The path to the model. Can be either local path or a Hugging Face model name.
+ """
+ import os
+ import json
+ is_local = os.path.exists(f"{path}/pipeline.json")
+
+ if is_local:
+ config_file = f"{path}/pipeline.json"
+ else:
+ from huggingface_hub import hf_hub_download
+ config_file = hf_hub_download(path, "pipeline.json")
+
+ with open(config_file, 'r') as f:
+ config = json.load(f)
+ return globals()[config['name']].from_pretrained(path)
diff --git a/Stable3DGen/hi3dgen/pipelines/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..78951bd7ea5b0fd6fc2356e2a7dea184ec68318d
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/__pycache__/base.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/__pycache__/base.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5be141cfd1a1921a00b6098fee17596432d70f57
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/__pycache__/base.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/__pycache__/hi3dgen.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/__pycache__/hi3dgen.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..82f6996fb15ee87312d1834c0f055802eb1b5217
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/__pycache__/hi3dgen.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/base.py b/Stable3DGen/hi3dgen/pipelines/base.py
new file mode 100755
index 0000000000000000000000000000000000000000..822e752998004302fca3c8db1bbc8d854d05a234
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/base.py
@@ -0,0 +1,90 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import torch.nn as nn
+from .. import models
+
+
+class Pipeline:
+ """
+ A base class for pipelines.
+ """
+ def __init__(
+ self,
+ models: dict[str, nn.Module] = None,
+ ):
+ if models is None:
+ return
+ self.models = models
+ for model in self.models.values():
+ model.eval()
+
+ @staticmethod
+ def from_pretrained(path: str) -> "Pipeline":
+ """
+ Load a pretrained model.
+ """
+ import os
+ import json
+ is_local = os.path.exists(f"{path}/pipeline.json")
+
+ if is_local:
+ config_file = f"{path}/pipeline.json"
+ else:
+ from huggingface_hub import hf_hub_download
+ config_file = hf_hub_download(path, "pipeline.json")
+
+ with open(config_file, 'r') as f:
+ args = json.load(f)['args']
+
+ _models = {
+ k: models.from_pretrained(f"{path}/{v}")
+ for k, v in args['models'].items()
+ }
+
+ new_pipeline = Pipeline(_models)
+ new_pipeline._pretrained_args = args
+ return new_pipeline
+
+ @property
+ def device(self) -> torch.device:
+ for model in self.models.values():
+ if hasattr(model, 'device'):
+ return model.device
+ for model in self.models.values():
+ if hasattr(model, 'parameters'):
+ return next(model.parameters()).device
+ raise RuntimeError("No device found.")
+
+ def to(self, device: torch.device) -> None:
+ for model in self.models.values():
+ model.to(device)
+
+ def cuda(self) -> None:
+ self.to(torch.device("cuda"))
+
+ def cpu(self) -> None:
+ self.to(torch.device("cpu"))
diff --git a/Stable3DGen/hi3dgen/pipelines/hi3dgen.py b/Stable3DGen/hi3dgen/pipelines/hi3dgen.py
new file mode 100755
index 0000000000000000000000000000000000000000..5389fff1d2a7a0fdddc6b441e1311e8c52ba3e2e
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/hi3dgen.py
@@ -0,0 +1,490 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+#
+# Original file was released under MIT, with the full license text
+# available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+#
+# This modified file is released under the same license.
+from typing import *
+from contextlib import contextmanager
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from tqdm import tqdm
+from torchvision import transforms
+from PIL import Image
+from .base import Pipeline
+from . import samplers
+from ..modules import sparse as sp
+import os
+
+class Hi3DGenPipeline(Pipeline):
+
+ def __init__(
+ self,
+ models: dict[str, nn.Module] = None,
+ sparse_structure_sampler: samplers.Sampler = None,
+ slat_sampler: samplers.Sampler = None,
+ slat_normalization: dict = None,
+ image_cond_model: str = None,
+ ):
+ if models is None:
+ return
+ super().__init__(models)
+ self.sparse_structure_sampler = sparse_structure_sampler
+ self.slat_sampler = slat_sampler
+ self.sparse_structure_sampler_params = {}
+ self.slat_sampler_params = {}
+ self.slat_normalization = slat_normalization
+ self._init_image_cond_model(image_cond_model)
+
+ @staticmethod
+ def from_pretrained(path: str) -> "Hi3DGenPipeline":
+ """
+ Load a pretrained model.
+
+ Args:
+ path (str): The path to the model. Can be either local path or a Hugging Face repository.
+ """
+ pipeline = super(Hi3DGenPipeline, Hi3DGenPipeline).from_pretrained(path)
+ new_pipeline = Hi3DGenPipeline()
+ new_pipeline.__dict__ = pipeline.__dict__
+ args = pipeline._pretrained_args
+
+ new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
+ new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
+
+ new_pipeline.slat_sampler = getattr(samplers, args['slat_sampler']['name'])(**args['slat_sampler']['args'])
+ new_pipeline.slat_sampler_params = args['slat_sampler']['params']
+
+ new_pipeline.slat_normalization = args['slat_normalization']
+
+ new_pipeline._init_image_cond_model(args['image_cond_model'])
+
+ return new_pipeline
+
+ def _init_image_cond_model(self, name: str):
+ """
+ Initialize the image conditioning model.
+ """
+ dinov2_model = torch.hub.load(
+ repo_or_dir="/share/project/cwm/houyuan.chen/Stable3DGen/dinov2",
+ model=name,
+ source='local', # <--- 这是最关键的修改
+ pretrained=False
+
+ )
+ dinov2_weight = torch.load("/share/project/cwm/houyuan.chen/Stable3DGen/dinov2_vitl14_reg4_pretrain.pth")
+ dinov2_model.load_state_dict(dinov2_weight)
+
+
+ # try:
+ # dinov2_model = torch.hub.load(os.path.join(torch.hub.get_dir(), 'facebookresearch_dinov2_main'), name, source='local',pretrained=True)
+ # except:
+ # dinov2_model = torch.hub.load('facebookresearch/dinov2', name, pretrained=True)
+ dinov2_model.eval()
+ self.models['image_cond_model'] = dinov2_model
+ transform = transforms.Compose([
+ transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+ ])
+ self.image_cond_model_transform = transform
+
+ def preprocess_image(self, input: Image.Image, resolution=518) -> Image.Image:
+ """
+ Preprocess the input image using BiRefNet for background removal.
+ Includes padding to maintain aspect ratio when resizing to 518x518.
+ """
+ # if has alpha channel, use it directly
+ has_alpha = False
+ if input.mode == 'RGBA':
+ alpha = np.array(input)[:, :, -1]
+ if not np.all(alpha == 255):
+ has_alpha = True
+
+ if has_alpha:
+ output = input
+ else:
+ input = input.convert('RGB')
+ max_size = max(input.size)
+ scale = min(1, 1024 / max_size)
+ if scale < 1:
+ input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
+
+ # Load BiRefNet model if not already loaded
+ if getattr(self, 'birefnet_model', None) is None:
+ self._lazy_load_birefnet()
+
+ # Get mask using BiRefNet
+ mask = self._get_birefnet_mask(input)
+
+ # Convert input to RGBA and apply mask
+ input_rgba = input.convert('RGBA')
+ input_array = np.array(input_rgba)
+ input_array[:, :, 3] = mask * 255 # Apply mask to alpha channel
+ output = Image.fromarray(input_array)
+
+ # Process the output image
+ output_np = np.array(output)
+ alpha = output_np[:, :, 3]
+
+ # Find bounding box of non-transparent pixels
+ bbox = np.argwhere(alpha > 0.8 * 255)
+ if len(bbox) == 0: # Handle case where no foreground is detected
+ return input.convert('RGB')
+
+ bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
+ center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+ size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
+ size = int(size * 1.2)
+
+ # Calculate and apply crop bbox
+ bbox = (
+ int(center[0] - size // 2),
+ int(center[1] - size // 2),
+ int(center[0] + size // 2),
+ int(center[1] + size // 2)
+ )
+
+ # Ensure bbox is within image bounds
+ bbox = (
+ max(0, bbox[0]),
+ max(0, bbox[1]),
+ min(output.width, bbox[2]),
+ min(output.height, bbox[3])
+ )
+
+ output = output.crop(bbox)
+
+ # Add padding to maintain aspect ratio
+ width, height = output.size
+ if width > height:
+ new_height = width
+ padding = (width - height) // 2
+ padded_output = Image.new('RGBA', (width, new_height), (0, 0, 0, 0))
+ padded_output.paste(output, (0, padding))
+ else:
+ new_width = height
+ padding = (height - width) // 2
+ padded_output = Image.new('RGBA', (new_width, height), (0, 0, 0, 0))
+ padded_output.paste(output, (padding, 0))
+
+ # Resize padded image to target size
+ padded_output = padded_output.resize((resolution, resolution), Image.Resampling.LANCZOS)
+
+ # Final processing
+ output = np.array(padded_output).astype(np.float32) / 255
+ output = np.dstack((
+ output[:, :, :3] * output[:, :, 3:4], # RGB channels premultiplied by alpha
+ output[:, :, 3] # Original alpha channel
+ ))
+ output = Image.fromarray((output * 255).astype(np.uint8), mode='RGBA')
+
+ return output
+
+ def _lazy_load_birefnet(self):
+ """Lazy loading of the BiRefNet model"""
+ from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation, AutoModelForImageSegmentation
+ self.birefnet_model = AutoModelForImageSegmentation.from_pretrained(
+ 'weights/BiRefNet',
+ trust_remote_code=True
+ ).to(self.device)
+ self.birefnet_model.eval()
+
+ def _get_birefnet_mask(self, image: Image.Image) -> np.ndarray:
+ """Get object mask using BiRefNet"""
+ image_size = (1024, 1024)
+ transform_image = transforms.Compose([
+ transforms.Resize(image_size),
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+ ])
+
+ input_images = transform_image(image).unsqueeze(0).to(self.device)
+
+ with torch.no_grad():
+ preds = self.birefnet_model(input_images)[-1].sigmoid().cpu()
+
+ pred = preds[0].squeeze()
+ pred_pil = transforms.ToPILImage()(pred)
+ mask = pred_pil.resize(image.size)
+ mask_np = np.array(mask)
+
+ return (mask_np > 128).astype(np.uint8)
+
+ @torch.no_grad()
+ def encode_image(self, image: Union[torch.Tensor, list[Image.Image]]) -> torch.Tensor:
+ """
+ Encode the image.
+
+ Args:
+ image (Union[torch.Tensor, list[Image.Image]]): The image to encode
+
+ Returns:
+ torch.Tensor: The encoded features.
+ """
+ if isinstance(image, torch.Tensor):
+ assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+ elif isinstance(image, list):
+ assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+ image = [i.resize((518, 518), Image.LANCZOS) for i in image]
+ image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+ image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+ image = torch.stack(image).to(self.device)
+ else:
+ raise ValueError(f"Unsupported type of image: {type(image)}")
+
+ image = self.image_cond_model_transform(image).to(self.device)
+ features = self.models['image_cond_model'](image, is_training=True)['x_prenorm']
+ patchtokens = F.layer_norm(features, features.shape[-1:])
+ return patchtokens
+
+ def get_cond(self, image: Union[torch.Tensor, list[Image.Image]]) -> dict:
+ """
+ Get the conditioning information for the model.
+
+ Args:
+ image (Union[torch.Tensor, list[Image.Image]]): The image prompts.
+
+ Returns:
+ dict: The conditioning information
+ """
+ cond = self.encode_image(image)
+ neg_cond = torch.zeros_like(cond)
+ return {
+ 'cond': cond,
+ 'neg_cond': neg_cond,
+ }
+
+ def sample_sparse_structure(
+ self,
+ cond: dict,
+ num_samples: int = 1,
+ sampler_params: dict = {},
+ ) -> torch.Tensor:
+ """
+ Sample sparse structures with the given conditioning.
+
+ Args:
+ cond (dict): The conditioning information.
+ num_samples (int): The number of samples to generate.
+ sampler_params (dict): Additional parameters for the sampler.
+ """
+ # Sample occupancy latent
+ flow_model = self.models['sparse_structure_flow_model']
+ reso = flow_model.resolution
+ noise = torch.randn(num_samples, flow_model.in_channels, reso, reso, reso).to(self.device)
+ sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
+ z_s = self.sparse_structure_sampler.sample(
+ flow_model,
+ noise,
+ **cond,
+ **sampler_params,
+ verbose=True
+ )["samples"]
+
+ # Decode occupancy latent
+ decoder = self.models['sparse_structure_decoder']
+ coords = torch.argwhere(decoder(z_s)>0)[:, [0, 2, 3, 4]].int()
+
+ return coords
+
+ def decode_slat(
+ self,
+ slat: sp.SparseTensor,
+ formats: List[str] = ['mesh',],
+ ) -> dict:
+ """
+ Decode the structured latent.
+
+ Args:
+ slat (sp.SparseTensor): The structured latent.
+ formats (List[str]): The formats to decode the structured latent to.
+
+ Returns:
+ dict: The decoded structured latent.
+ """
+ ret = {}
+ if 'mesh' in formats:
+ ret['mesh'] = self.models['slat_decoder_mesh'](slat)
+ return ret
+
+ def sample_slat(
+ self,
+ cond: dict,
+ coords: torch.Tensor,
+ sampler_params: dict = {},
+ ) -> sp.SparseTensor:
+ """
+ Sample structured latent with the given conditioning.
+
+ Args:
+ cond (dict): The conditioning information.
+ coords (torch.Tensor): The coordinates of the sparse structure.
+ sampler_params (dict): Additional parameters for the sampler.
+ """
+ # Sample structured latent
+ flow_model = self.models['slat_flow_model']
+ noise = sp.SparseTensor(
+ feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
+ coords=coords,
+ )
+ sampler_params = {**self.slat_sampler_params, **sampler_params}
+ slat = self.slat_sampler.sample(
+ flow_model,
+ noise,
+ **cond,
+ **sampler_params,
+ verbose=True
+ )["samples"]
+
+ std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device)
+ mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device)
+ slat = slat * std + mean
+
+ return slat
+
+ @torch.no_grad()
+ def run(
+ self,
+ image: Image.Image,
+ num_samples: int = 1,
+ seed: int = 42,
+ sparse_structure_sampler_params: dict = {},
+ slat_sampler_params: dict = {},
+ formats: List[str] = ['mesh',],
+ preprocess_image: bool = True,
+ ) -> dict:
+ """
+ Run the pipeline.
+
+ Args:
+ image (Image.Image): The image prompt.
+ num_samples (int): The number of samples to generate.
+ sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
+ slat_sampler_params (dict): Additional parameters for the structured latent sampler.
+ preprocess_image (bool): Whether to preprocess the image.
+ """
+ if preprocess_image:
+ image = self.preprocess_image(image)
+ cond = self.get_cond([image])
+ torch.manual_seed(seed)
+ coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
+ slat = self.sample_slat(cond, coords, slat_sampler_params)
+ return self.decode_slat(slat, formats)
+
+ @contextmanager
+ def inject_sampler_multi_image(
+ self,
+ sampler_name: str,
+ num_images: int,
+ num_steps: int,
+ mode: Literal['stochastic', 'multidiffusion'] = 'stochastic',
+ ):
+ """
+ Inject a sampler with multiple images as condition.
+
+ Args:
+ sampler_name (str): The name of the sampler to inject.
+ num_images (int): The number of images to condition on.
+ num_steps (int): The number of steps to run the sampler for.
+ """
+ sampler = getattr(self, sampler_name)
+ setattr(sampler, f'_old_inference_model', sampler._inference_model)
+
+ if mode == 'stochastic':
+ if num_images > num_steps:
+ print(f"\033[93mWarning: number of conditioning images is greater than number of steps for {sampler_name}. "
+ "This may lead to performance degradation.\033[0m")
+
+ cond_indices = (np.arange(num_steps) % num_images).tolist()
+ def _new_inference_model(self, model, x_t, t, cond, **kwargs):
+ cond_idx = cond_indices.pop(0)
+ cond_i = cond[cond_idx:cond_idx+1]
+ return self._old_inference_model(model, x_t, t, cond=cond_i, **kwargs)
+
+ elif mode =='multidiffusion':
+ from .samplers import FlowEulerSampler
+ def _new_inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, cfg_interval, **kwargs):
+ if cfg_interval[0] <= t <= cfg_interval[1]:
+ preds = []
+ for i in range(len(cond)):
+ preds.append(FlowEulerSampler._inference_model(self, model, x_t, t, cond[i:i+1], **kwargs))
+ pred = sum(preds) / len(preds)
+ neg_pred = FlowEulerSampler._inference_model(self, model, x_t, t, neg_cond, **kwargs)
+ return (1 + cfg_strength) * pred - cfg_strength * neg_pred
+ else:
+ preds = []
+ for i in range(len(cond)):
+ preds.append(FlowEulerSampler._inference_model(self, model, x_t, t, cond[i:i+1], **kwargs))
+ pred = sum(preds) / len(preds)
+ return pred
+
+ else:
+ raise ValueError(f"Unsupported mode: {mode}")
+
+ sampler._inference_model = _new_inference_model.__get__(sampler, type(sampler))
+
+ yield
+
+ sampler._inference_model = sampler._old_inference_model
+ delattr(sampler, f'_old_inference_model')
+
+ @torch.no_grad()
+ def run_multi_image(
+ self,
+ images: List[Image.Image],
+ num_samples: int = 1,
+ seed: int = 42,
+ sparse_structure_sampler_params: dict = {},
+ slat_sampler_params: dict = {},
+ formats: List[str] = ['mesh', 'radiance_field'],
+ preprocess_image: bool = True,
+ mode: Literal['stochastic', 'multidiffusion'] = 'stochastic',
+ ) -> dict:
+ """
+ Run the pipeline with multiple images as condition
+
+ Args:
+ images (List[Image.Image]): The multi-view images of the assets
+ num_samples (int): The number of samples to generate.
+ sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
+ slat_sampler_params (dict): Additional parameters for the structured latent sampler.
+ preprocess_image (bool): Whether to preprocess the image.
+ """
+ if preprocess_image:
+ images = [self.preprocess_image(image) for image in images]
+ cond = self.get_cond(images)
+ cond['neg_cond'] = cond['neg_cond'][:1]
+ torch.manual_seed(seed)
+ ss_steps = {**self.sparse_structure_sampler_params, **sparse_structure_sampler_params}.get('steps')
+ with self.inject_sampler_multi_image('sparse_structure_sampler', len(images), ss_steps, mode=mode):
+ coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
+ slat_steps = {**self.slat_sampler_params, **slat_sampler_params}.get('steps')
+ with self.inject_sampler_multi_image('slat_sampler', len(images), slat_steps, mode=mode):
+ slat = self.sample_slat(cond, coords, slat_sampler_params)
+ return self.decode_slat(slat, formats)
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/__init__.py b/Stable3DGen/hi3dgen/pipelines/samplers/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..96cc97ab6b09cf643d9adeb7696f7b27796c1b07
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/samplers/__init__.py
@@ -0,0 +1,26 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from .base import Sampler
+from .flow_euler import FlowEulerSampler, FlowEulerCfgSampler, FlowEulerGuidanceIntervalSampler
\ No newline at end of file
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c031108367ef61d3755be36df930d880b58229a5
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/base.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/base.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5ccaec89884215c8acadd31290f040bdc0902d6c
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/base.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/classifier_free_guidance_mixin.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/classifier_free_guidance_mixin.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8a2edf22b8bd7775c1ecbbbeed970abd80a716cd
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/classifier_free_guidance_mixin.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/flow_euler.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/flow_euler.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8e43d76e07aeb11c02177768adeb77bfadeaa2eb
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/flow_euler.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/guidance_interval_mixin.cpython-310.pyc b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/guidance_interval_mixin.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4149223cb80e2293df35c702c9ad4db59e0dfa19
Binary files /dev/null and b/Stable3DGen/hi3dgen/pipelines/samplers/__pycache__/guidance_interval_mixin.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/base.py b/Stable3DGen/hi3dgen/pipelines/samplers/base.py
new file mode 100755
index 0000000000000000000000000000000000000000..d6f68bb8124b13bf6c9a80783a1786f85da5fe1b
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/samplers/base.py
@@ -0,0 +1,44 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+from abc import ABC, abstractmethod
+
+
+class Sampler(ABC):
+ """
+ A base class for samplers.
+ """
+
+ @abstractmethod
+ def sample(
+ self,
+ model,
+ **kwargs
+ ):
+ """
+ Sample from a model.
+ """
+ pass
+
\ No newline at end of file
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/classifier_free_guidance_mixin.py b/Stable3DGen/hi3dgen/pipelines/samplers/classifier_free_guidance_mixin.py
new file mode 100755
index 0000000000000000000000000000000000000000..566f830fa921b59430d9f724acd4bd7cd41f22a5
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/samplers/classifier_free_guidance_mixin.py
@@ -0,0 +1,36 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+
+
+class ClassifierFreeGuidanceSamplerMixin:
+ """
+ A mixin class for samplers that apply classifier-free guidance.
+ """
+
+ def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, **kwargs):
+ pred = super()._inference_model(model, x_t, t, cond, **kwargs)
+ neg_pred = super()._inference_model(model, x_t, t, neg_cond, **kwargs)
+ return (1 + cfg_strength) * pred - cfg_strength * neg_pred
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/flow_euler.py b/Stable3DGen/hi3dgen/pipelines/samplers/flow_euler.py
new file mode 100755
index 0000000000000000000000000000000000000000..9a222b67cedb46f402fc1f8f8142de6423bd2094
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/samplers/flow_euler.py
@@ -0,0 +1,222 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+import torch
+import numpy as np
+from tqdm import tqdm
+from .base import Sampler
+from .classifier_free_guidance_mixin import ClassifierFreeGuidanceSamplerMixin
+from .guidance_interval_mixin import GuidanceIntervalSamplerMixin
+
+
+class FlowEulerSampler(Sampler):
+ """
+ Generate samples from a flow-matching model using Euler sampling.
+
+ Args:
+ sigma_min: The minimum scale of noise in flow.
+ """
+ def __init__(
+ self,
+ sigma_min: float,
+ ):
+ self.sigma_min = sigma_min
+
+ def _eps_to_xstart(self, x_t, t, eps):
+ assert x_t.shape == eps.shape
+ return (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * eps) / (1 - t)
+
+ def _xstart_to_eps(self, x_t, t, x_0):
+ assert x_t.shape == x_0.shape
+ return (x_t - (1 - t) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t)
+
+ def _v_to_xstart_eps(self, x_t, t, v):
+ assert x_t.shape == v.shape
+ eps = (1 - t) * v + x_t
+ x_0 = (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * v
+ return x_0, eps
+
+ def _inference_model(self, model, x_t, t, cond=None, **kwargs):
+ t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32)
+ return model(x_t, t, cond, **kwargs)
+
+ def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs):
+ pred_v = self._inference_model(model, x_t, t, cond, **kwargs)
+ pred_x_0, pred_eps = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v)
+ return pred_x_0, pred_eps, pred_v
+
+ @torch.no_grad()
+ def sample_once(
+ self,
+ model,
+ x_t,
+ t: float,
+ t_prev: float,
+ cond: Optional[Any] = None,
+ **kwargs
+ ):
+ """
+ Sample x_{t-1} from the model using Euler method.
+
+ Args:
+ model: The model to sample from.
+ x_t: The [N x C x ...] tensor of noisy inputs at time t.
+ t: The current timestep.
+ t_prev: The previous timestep.
+ cond: conditional information.
+ **kwargs: Additional arguments for model inference.
+
+ Returns:
+ a dict containing the following
+ - 'pred_x_prev': x_{t-1}.
+ - 'pred_x_0': a prediction of x_0.
+ """
+ pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
+ pred_x_prev = x_t - (t - t_prev) * pred_v
+ return {"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0}
+
+ @torch.no_grad()
+ def sample(
+ self,
+ model,
+ noise,
+ cond: Optional[Any] = None,
+ steps: int = 50,
+ rescale_t: float = 1.0,
+ verbose: bool = True,
+ **kwargs
+ ):
+ """
+ Generate samples from the model using Euler method.
+
+ Args:
+ model: The model to sample from.
+ noise: The initial noise tensor.
+ cond: conditional information.
+ steps: The number of steps to sample.
+ rescale_t: The rescale factor for t.
+ verbose: If True, show a progress bar.
+ **kwargs: Additional arguments for model_inference.
+
+ Returns:
+ a dict containing the following
+ - 'samples': the model samples.
+ - 'pred_x_t': a list of prediction of x_t.
+ - 'pred_x_0': a list of prediction of x_0.
+ """
+ sample = noise
+ t_seq = np.linspace(1, 0, steps + 1)
+ t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
+ t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
+ ret = {"samples": None, "pred_x_t": [], "pred_x_0": []}
+ for t, t_prev in tqdm(t_pairs, desc="Sampling", disable=not verbose):
+ out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
+ sample = out["pred_x_prev"]
+ ret["pred_x_t"].append(out["pred_x_prev"])
+ ret["pred_x_0"].append(out["pred_x_0"])
+ ret["samples"] = sample
+ return ret
+
+
+class FlowEulerCfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler):
+ """
+ Generate samples from a flow-matching model using Euler sampling with classifier-free guidance.
+ """
+ @torch.no_grad()
+ def sample(
+ self,
+ model,
+ noise,
+ cond,
+ neg_cond,
+ steps: int = 50,
+ rescale_t: float = 1.0,
+ cfg_strength: float = 3.0,
+ verbose: bool = True,
+ **kwargs
+ ):
+ """
+ Generate samples from the model using Euler method.
+
+ Args:
+ model: The model to sample from.
+ noise: The initial noise tensor.
+ cond: conditional information.
+ neg_cond: negative conditional information.
+ steps: The number of steps to sample.
+ rescale_t: The rescale factor for t.
+ cfg_strength: The strength of classifier-free guidance.
+ verbose: If True, show a progress bar.
+ **kwargs: Additional arguments for model_inference.
+
+ Returns:
+ a dict containing the following
+ - 'samples': the model samples.
+ - 'pred_x_t': a list of prediction of x_t.
+ - 'pred_x_0': a list of prediction of x_0.
+ """
+ return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, **kwargs)
+
+
+class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, FlowEulerSampler):
+ """
+ Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval.
+ """
+ @torch.no_grad()
+ def sample(
+ self,
+ model,
+ noise,
+ cond,
+ neg_cond,
+ steps: int = 50,
+ rescale_t: float = 1.0,
+ cfg_strength: float = 3.0,
+ cfg_interval: Tuple[float, float] = (0.0, 1.0),
+ verbose: bool = True,
+ **kwargs
+ ):
+ """
+ Generate samples from the model using Euler method.
+
+ Args:
+ model: The model to sample from.
+ noise: The initial noise tensor.
+ cond: conditional information.
+ neg_cond: negative conditional information.
+ steps: The number of steps to sample.
+ rescale_t: The rescale factor for t.
+ cfg_strength: The strength of classifier-free guidance.
+ cfg_interval: The interval for classifier-free guidance.
+ verbose: If True, show a progress bar.
+ **kwargs: Additional arguments for model_inference.
+
+ Returns:
+ a dict containing the following
+ - 'samples': the model samples.
+ - 'pred_x_t': a list of prediction of x_t.
+ - 'pred_x_0': a list of prediction of x_0.
+ """
+ return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, cfg_strength=cfg_strength, cfg_interval=cfg_interval, **kwargs)
diff --git a/Stable3DGen/hi3dgen/pipelines/samplers/guidance_interval_mixin.py b/Stable3DGen/hi3dgen/pipelines/samplers/guidance_interval_mixin.py
new file mode 100755
index 0000000000000000000000000000000000000000..8ab16a7c706c8e8617a69a31770e3e7da4415e3a
--- /dev/null
+++ b/Stable3DGen/hi3dgen/pipelines/samplers/guidance_interval_mixin.py
@@ -0,0 +1,39 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from typing import *
+
+
+class GuidanceIntervalSamplerMixin:
+ """
+ A mixin class for samplers that apply classifier-free guidance with interval.
+ """
+
+ def _inference_model(self, model, x_t, t, cond, neg_cond, cfg_strength, cfg_interval, **kwargs):
+ if cfg_interval[0] <= t <= cfg_interval[1]:
+ pred = super()._inference_model(model, x_t, t, cond, **kwargs)
+ neg_pred = super()._inference_model(model, x_t, t, neg_cond, **kwargs)
+ return (1 + cfg_strength) * pred - cfg_strength * neg_pred
+ else:
+ return super()._inference_model(model, x_t, t, cond, **kwargs)
diff --git a/Stable3DGen/hi3dgen/representations/__init__.py b/Stable3DGen/hi3dgen/representations/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..cf34fec872183a8b77ca50d2263c1677387a3d0f
--- /dev/null
+++ b/Stable3DGen/hi3dgen/representations/__init__.py
@@ -0,0 +1,25 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from .mesh import MeshExtractResult
diff --git a/Stable3DGen/hi3dgen/representations/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/representations/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..44cd4d91516838c8688cd4fd3a3df010f51444fe
Binary files /dev/null and b/Stable3DGen/hi3dgen/representations/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/representations/mesh/__init__.py b/Stable3DGen/hi3dgen/representations/mesh/__init__.py
new file mode 100755
index 0000000000000000000000000000000000000000..a79be70df5b8cb361fae019782333d3baf2d6ef9
--- /dev/null
+++ b/Stable3DGen/hi3dgen/representations/mesh/__init__.py
@@ -0,0 +1,25 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+from .cube2mesh import SparseFeatures2Mesh, MeshExtractResult
diff --git a/Stable3DGen/hi3dgen/representations/mesh/__pycache__/__init__.cpython-310.pyc b/Stable3DGen/hi3dgen/representations/mesh/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8f8d20d3be1f15327513f51d1e656dcd7ccaee8c
Binary files /dev/null and b/Stable3DGen/hi3dgen/representations/mesh/__pycache__/__init__.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/representations/mesh/__pycache__/cube2mesh.cpython-310.pyc b/Stable3DGen/hi3dgen/representations/mesh/__pycache__/cube2mesh.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..241850433476b8384b06e20ea32468b0ead800c1
Binary files /dev/null and b/Stable3DGen/hi3dgen/representations/mesh/__pycache__/cube2mesh.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/representations/mesh/__pycache__/utils_cube.cpython-310.pyc b/Stable3DGen/hi3dgen/representations/mesh/__pycache__/utils_cube.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..7ae955bf2f9365d114c1ea2f77349242d5b10152
Binary files /dev/null and b/Stable3DGen/hi3dgen/representations/mesh/__pycache__/utils_cube.cpython-310.pyc differ
diff --git a/Stable3DGen/hi3dgen/representations/mesh/cube2mesh.py b/Stable3DGen/hi3dgen/representations/mesh/cube2mesh.py
new file mode 100755
index 0000000000000000000000000000000000000000..e8acd41051da2e2c99d22fa8e64701e00fa536c3
--- /dev/null
+++ b/Stable3DGen/hi3dgen/representations/mesh/cube2mesh.py
@@ -0,0 +1,384 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [jclarkk]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+#
+# Original file was released under MIT, with the full license text
+# available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+#
+# This modified file is released under the same license.
+import torch
+from ...modules.sparse import SparseTensor
+from .utils_cube import *
+import numpy as np
+import trimesh
+import numpy as np
+from skimage import measure
+from typing import Tuple, Optional
+
+class MeshExtractResult:
+ def __init__(self,
+ vertices,
+ faces,
+ vertex_attrs=None,
+ res=64
+ ):
+ self.vertices = vertices
+ self.faces = faces.long()
+ self.vertex_attrs = vertex_attrs
+ self.vertex_normal = self.comput_v_normals(vertices, faces)
+ self.face_normal = self.comput_face_normals(vertices, faces)
+ self.res = res
+ self.success = (vertices.shape[0] != 0 and faces.shape[0] != 0)
+
+ # training only
+ self.tsdf_v = None
+ self.tsdf_s = None
+ self.reg_loss = None
+
+ def comput_face_normals(self, verts, faces):
+ i0 = faces[..., 0].long()
+ i1 = faces[..., 1].long()
+ i2 = faces[..., 2].long()
+
+ v0 = verts[i0, :]
+ v1 = verts[i1, :]
+ v2 = verts[i2, :]
+ face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+ face_normals = torch.nn.functional.normalize(face_normals, dim=1)
+ # print(face_normals.min(), face_normals.max(), face_normals.shape)
+ return face_normals[:, None, :].repeat(1, 3, 1)
+
+ def comput_v_normals(self, verts, faces):
+ i0 = faces[..., 0].long()
+ i1 = faces[..., 1].long()
+ i2 = faces[..., 2].long()
+
+ v0 = verts[i0, :]
+ v1 = verts[i1, :]
+ v2 = verts[i2, :]
+ face_normals = torch.cross(v1 - v0, v2 - v0, dim=-1)
+ v_normals = torch.zeros_like(verts)
+ v_normals.scatter_add_(0, i0[..., None].repeat(1, 3), face_normals)
+ v_normals.scatter_add_(0, i1[..., None].repeat(1, 3), face_normals)
+ v_normals.scatter_add_(0, i2[..., None].repeat(1, 3), face_normals)
+
+ v_normals = torch.nn.functional.normalize(v_normals, dim=1)
+ return v_normals
+
+ def to_trimesh(self, transform_pose=False):
+ vertices = self.vertices.detach().cpu().numpy()
+ faces = self.faces.detach().cpu().numpy()
+
+ if transform_pose:
+ transform_matrix = np.array([
+ [1, 0, 0],
+ [0, 0, -1],
+ [0, 1, 0]
+ ])
+ vertices = vertices @ transform_matrix
+ vertex_normals = self.vertex_normal.detach().cpu().numpy() @ transform_matrix
+ else:
+ vertex_normals = self.vertex_normal.detach().cpu().numpy()
+
+ # Create the trimesh mesh
+ mesh = trimesh.Trimesh(
+ vertices=vertices,
+ faces=faces,
+ face_normals=self.face_normal.detach().cpu().numpy(),
+ vertex_normals=vertex_normals
+ )
+
+ return mesh
+
+class EnhancedMarchingCubes:
+ def __init__(self, device="cuda"):
+ self.device = device
+
+ def __call__(self,
+ voxelgrid_vertices: torch.Tensor,
+ scalar_field: torch.Tensor,
+ voxelgrid_colors: Optional[torch.Tensor] = None,
+ training: bool = False
+ ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Enhanced Marching Cubes implementation that handles deformations and colors
+ """
+ if scalar_field.dim() == 1:
+ grid_size = int(round(scalar_field.shape[0] ** (1 / 3)))
+ scalar_field = scalar_field.reshape(grid_size, grid_size, grid_size)
+ elif scalar_field.dim() > 3:
+ scalar_field = scalar_field.squeeze()
+
+ # Convert to numpy and ensure values are in correct range
+ scalar_np = scalar_field.cpu().numpy()
+
+ if scalar_np.ndim != 3:
+ raise ValueError(f"Expected 3D array, got shape {scalar_np.shape}")
+
+ # Run marching cubes with normalized coordinates
+ vertices, faces, normals, _ = measure.marching_cubes(
+ scalar_np,
+ level=0.0,
+ gradient_direction='ascent'
+ )
+
+ vertices = torch.from_numpy(np.ascontiguousarray(vertices)).float().to(self.device)
+ faces = torch.from_numpy(np.ascontiguousarray(faces)).long().to(self.device)
+
+ # Apply deformations
+ if voxelgrid_vertices is not None:
+ # Reshape and normalize voxelgrid_vertices if needed
+ if voxelgrid_vertices.dim() == 2:
+ voxelgrid_vertices = voxelgrid_vertices.reshape(grid_size, grid_size, grid_size, 3)
+ deformed_vertices = self._apply_deformations(vertices, voxelgrid_vertices)
+ else:
+ deformed_vertices = vertices
+
+ # Handle colors if provided
+ colors = None
+ if voxelgrid_colors is not None:
+ if voxelgrid_colors.dim() == 2:
+ voxelgrid_colors = voxelgrid_colors.reshape(grid_size, grid_size, grid_size, -1)
+ colors = self._interpolate_colors(vertices, voxelgrid_colors)
+ # Ensure colors are in [0, 1] range
+ colors = torch.sigmoid(colors)
+
+ # Compute deviation loss for training
+ deviation_loss = torch.tensor(0.0, device=self.device)
+ if training:
+ deviation_loss = self._compute_deviation_loss(vertices, deformed_vertices)
+
+ faces = faces.flip(dims=[1]) # Reverse the order of vertices in each face, for some reason it's reversed...
+
+ return deformed_vertices, faces, deviation_loss, colors
+
+ def _apply_deformations(self, vertices: torch.Tensor,
+ voxelgrid_vertices: torch.Tensor) -> torch.Tensor:
+ """Apply deformations to vertices using trilinear interpolation"""
+
+ grid_positions = vertices.clone()
+
+ # Scale to grid coordinates
+ grid_coords = grid_positions.long()
+ local_coords = grid_positions - grid_coords.float()
+
+ # Reshape voxelgrid_vertices if needed
+ if voxelgrid_vertices.dim() == 2:
+ # Assuming voxelgrid_vertices is [N, 3]
+ grid_size = int(round(voxelgrid_vertices.shape[0] ** (1 / 3)))
+ voxelgrid_vertices = voxelgrid_vertices.reshape(grid_size, grid_size, grid_size, 3)
+
+ # Ensure coordinates are within bounds
+ grid_coords = torch.clamp(grid_coords, 0, voxelgrid_vertices.shape[0] - 1)
+
+ # Perform trilinear interpolation
+ deformed_vertices = self._trilinear_interpolate(
+ grid_coords, local_coords, voxelgrid_vertices
+ )
+
+ return deformed_vertices
+
+ def _interpolate_colors(self, vertices: torch.Tensor,
+ voxelgrid_colors: torch.Tensor) -> torch.Tensor:
+ """Interpolate colors for vertices"""
+
+ # Get grid positions
+ grid_positions = vertices.clone()
+
+ # Scale to grid coordinates
+ grid_coords = grid_positions.long()
+ local_coords = grid_positions - grid_coords.float()
+
+ # Reshape colors if they're in 2D format
+ if voxelgrid_colors.dim() == 2:
+ grid_size = int(round(voxelgrid_colors.shape[0] ** (1 / 3)))
+ color_channels = voxelgrid_colors.shape[1]
+ voxelgrid_colors = voxelgrid_colors.reshape(grid_size, grid_size, grid_size, color_channels)
+
+ # Ensure coordinates are within bounds
+ grid_coords = torch.clamp(grid_coords, 0, voxelgrid_colors.shape[0] - 1)
+
+ # Perform trilinear interpolation
+ return self._trilinear_interpolate(
+ grid_coords, local_coords, voxelgrid_colors, is_color=True
+ )
+
+ def _trilinear_interpolate(self, grid_coords: torch.Tensor,
+ local_coords: torch.Tensor,
+ values: torch.Tensor,
+ is_color: bool = False) -> torch.Tensor:
+ """Perform trilinear interpolation"""
+ x, y, z = local_coords[:, 0], local_coords[:, 1], local_coords[:, 2]
+
+ if is_color and values.dim() == 2:
+ # Handle flat color array
+ grid_size = int(round(values.shape[0] ** (1 / 3)))
+ color_channels = values.shape[1]
+ values = values.reshape(grid_size, grid_size, grid_size, color_channels)
+
+ # Get corner values with proper indexing based on dimensionality
+ if values.dim() == 4: # For 4D tensors (grid x grid x grid x channels)
+ c000 = values[grid_coords[:, 0], grid_coords[:, 1], grid_coords[:, 2], :]
+ c001 = values[grid_coords[:, 0], grid_coords[:, 1],
+ torch.clamp(grid_coords[:, 2] + 1, 0, values.shape[2] - 1), :]
+ c010 = values[grid_coords[:, 0], torch.clamp(grid_coords[:, 1] + 1, 0, values.shape[1] - 1),
+ grid_coords[:, 2], :]
+ c011 = values[grid_coords[:, 0], torch.clamp(grid_coords[:, 1] + 1, 0, values.shape[1] - 1),
+ torch.clamp(grid_coords[:, 2] + 1, 0, values.shape[2] - 1), :]
+ c100 = values[torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1), grid_coords[:, 1],
+ grid_coords[:, 2], :]
+ c101 = values[torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1), grid_coords[:, 1],
+ torch.clamp(grid_coords[:, 2] + 1, 0, values.shape[2] - 1), :]
+ c110 = values[torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1),
+ torch.clamp(grid_coords[:, 1] + 1, 0, values.shape[1] - 1), grid_coords[:, 2], :]
+ c111 = values[torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1),
+ torch.clamp(grid_coords[:, 1] + 1, 0, values.shape[1] - 1),
+ torch.clamp(grid_coords[:, 2] + 1, 0, values.shape[2] - 1), :]
+ else:
+ c000 = values[grid_coords[:, 0], grid_coords[:, 1], grid_coords[:, 2]]
+ c001 = values[
+ grid_coords[:, 0], grid_coords[:, 1], torch.clamp(grid_coords[:, 2] + 1, 0, values.shape[2] - 1)]
+ c010 = values[
+ grid_coords[:, 0], torch.clamp(grid_coords[:, 1] + 1, 0, values.shape[1] - 1), grid_coords[:, 2]]
+ c011 = values[grid_coords[:, 0], torch.clamp(grid_coords[:, 1] + 1, 0, values.shape[1] - 1), torch.clamp(
+ grid_coords[:, 2] + 1, 0, values.shape[2] - 1)]
+ c100 = values[
+ torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1), grid_coords[:, 1], grid_coords[:, 2]]
+ c101 = values[torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1), grid_coords[:, 1], torch.clamp(
+ grid_coords[:, 2] + 1, 0, values.shape[2] - 1)]
+ c110 = values[
+ torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1), torch.clamp(grid_coords[:, 1] + 1, 0,
+ values.shape[
+ 1] - 1), grid_coords[:, 2]]
+ c111 = values[
+ torch.clamp(grid_coords[:, 0] + 1, 0, values.shape[0] - 1), torch.clamp(grid_coords[:, 1] + 1, 0,
+ values.shape[
+ 1] - 1), torch.clamp(
+ grid_coords[:, 2] + 1, 0, values.shape[2] - 1)]
+
+ # Add channel dimension for 3D tensors if needed
+ if values.dim() == 3:
+ c000, c001, c010, c011 = [c[..., None] if c.dim() == 1 else c for c in [c000, c001, c010, c011]]
+ c100, c101, c110, c111 = [c[..., None] if c.dim() == 1 else c for c in [c100, c101, c110, c111]]
+
+ # Interpolate along x
+ c00 = c000 * (1 - x)[:, None] + c100 * x[:, None]
+ c01 = c001 * (1 - x)[:, None] + c101 * x[:, None]
+ c10 = c010 * (1 - x)[:, None] + c110 * x[:, None]
+ c11 = c011 * (1 - x)[:, None] + c111 * x[:, None]
+
+ # Interpolate along y
+ c0 = c00 * (1 - y)[:, None] + c10 * y[:, None]
+ c1 = c01 * (1 - y)[:, None] + c11 * y[:, None]
+
+ # Interpolate along z
+ return c0 * (1 - z)[:, None] + c1 * z[:, None]
+
+ def _compute_deviation_loss(self, original_vertices: torch.Tensor,
+ deformed_vertices: torch.Tensor) -> torch.Tensor:
+ """Compute deviation loss for training"""
+ return torch.mean((deformed_vertices - original_vertices) ** 2)
+
+class SparseFeatures2Mesh:
+ def __init__(self, device="cuda", res=128, use_color=True):
+ super().__init__()
+ self.device = device
+ self.res = res
+ self.mesh_extractor = EnhancedMarchingCubes(device=device)
+ self.sdf_bias = -1.0 / res
+ verts, cube = construct_dense_grid(self.res, self.device)
+ self.reg_c = cube.to(self.device)
+ self.reg_v = verts.to(self.device)
+ self.use_color = use_color
+ self._calc_layout()
+
+ def _calc_layout(self):
+ LAYOUTS = {
+ 'sdf': {'shape': (8, 1), 'size': 8},
+ 'deform': {'shape': (8, 3), 'size': 8 * 3},
+ 'weights': {'shape': (21,), 'size': 21}
+ }
+ if self.use_color:
+ '''
+ 6 channel color including normal map
+ '''
+ LAYOUTS['color'] = {'shape': (8, 6,), 'size': 8 * 6}
+ self.layouts = LAYOUTS
+ start = 0
+ for k, v in self.layouts.items():
+ v['range'] = (start, start + v['size'])
+ start += v['size']
+ self.feats_channels = start
+
+ def get_layout(self, feats: torch.Tensor, name: str):
+ if name not in self.layouts:
+ return None
+ return feats[:, self.layouts[name]['range'][0]:self.layouts[name]['range'][1]].reshape(-1, *self.layouts[name][
+ 'shape'])
+
+ def __call__(self, cubefeats: SparseTensor, training=False):
+ coords = cubefeats.coords[:, 1:]
+ feats = cubefeats.feats
+
+ sdf, deform, color, weights = [self.get_layout(feats, name)
+ for name in ['sdf', 'deform', 'color', 'weights']]
+ sdf += self.sdf_bias
+ v_attrs = [sdf, deform, color] if self.use_color else [sdf, deform]
+ v_pos, v_attrs, reg_loss = sparse_cube2verts(coords, torch.cat(v_attrs, dim=-1),
+ training=training)
+
+ v_attrs_d = get_dense_attrs(v_pos, v_attrs, res=self.res + 1, sdf_init=True)
+
+ if self.use_color:
+ sdf_d, deform_d, colors_d = (v_attrs_d[..., 0], v_attrs_d[..., 1:4],
+ v_attrs_d[..., 4:])
+ else:
+ sdf_d, deform_d = v_attrs_d[..., 0], v_attrs_d[..., 1:4]
+ colors_d = None
+
+ x_nx3 = get_defomed_verts(self.reg_v, deform_d, self.res)
+
+ vertices, faces, L_dev, colors = self.mesh_extractor(
+ voxelgrid_vertices=x_nx3,
+ scalar_field=sdf_d,
+ voxelgrid_colors=colors_d,
+ training=training
+ )
+
+ mesh = MeshExtractResult(vertices=vertices, faces=faces,
+ vertex_attrs=colors, res=self.res)
+
+ if training:
+ if mesh.success:
+ reg_loss += L_dev.mean() * 0.5
+ reg_loss += (weights[:, :20]).abs().mean() * 0.2
+ mesh.reg_loss = reg_loss
+ mesh.tsdf_v = get_defomed_verts(v_pos, v_attrs[:, 1:4], self.res)
+ mesh.tsdf_s = v_attrs[:, 0]
+
+ return mesh
+
diff --git a/Stable3DGen/hi3dgen/representations/mesh/utils_cube.py b/Stable3DGen/hi3dgen/representations/mesh/utils_cube.py
new file mode 100755
index 0000000000000000000000000000000000000000..29fc592ce3c1288bf5ce61f06c2a5e62d35f9276
--- /dev/null
+++ b/Stable3DGen/hi3dgen/representations/mesh/utils_cube.py
@@ -0,0 +1,85 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# SPDX-License-Identifier: MIT
+import torch
+cube_corners = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1], [
+ 1, 0, 1], [0, 1, 1], [1, 1, 1]], dtype=torch.int)
+cube_neighbor = torch.tensor([[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]])
+cube_edges = torch.tensor([0, 1, 1, 5, 4, 5, 0, 4, 2, 3, 3, 7, 6, 7, 2, 6,
+ 2, 0, 3, 1, 7, 5, 6, 4], dtype=torch.long, requires_grad=False)
+
+def construct_dense_grid(res, device='cuda'):
+ '''construct a dense grid based on resolution'''
+ res_v = res + 1
+ vertsid = torch.arange(res_v ** 3, device=device)
+ coordsid = vertsid.reshape(res_v, res_v, res_v)[:res, :res, :res].flatten()
+ cube_corners_bias = (cube_corners[:, 0] * res_v + cube_corners[:, 1]) * res_v + cube_corners[:, 2]
+ cube_fx8 = (coordsid.unsqueeze(1) + cube_corners_bias.unsqueeze(0).to(device))
+ verts = torch.stack([vertsid // (res_v ** 2), (vertsid // res_v) % res_v, vertsid % res_v], dim=1)
+ return verts, cube_fx8
+
+
+def construct_voxel_grid(coords):
+ verts = (cube_corners.unsqueeze(0).to(coords) + coords.unsqueeze(1)).reshape(-1, 3)
+ verts_unique, inverse_indices = torch.unique(verts, dim=0, return_inverse=True)
+ cubes = inverse_indices.reshape(-1, 8)
+ return verts_unique, cubes
+
+
+def cubes_to_verts(num_verts, cubes, value, reduce='mean'):
+ """
+ Args:
+ cubes [Vx8] verts index for each cube
+ value [Vx8xM] value to be scattered
+ Operation:
+ reduced[cubes[i][j]][k] += value[i][k]
+ """
+ M = value.shape[2] # number of channels
+ reduced = torch.zeros(num_verts, M, device=cubes.device)
+ return torch.scatter_reduce(reduced, 0,
+ cubes.unsqueeze(-1).expand(-1, -1, M).flatten(0, 1),
+ value.flatten(0, 1), reduce=reduce, include_self=False)
+
+def sparse_cube2verts(coords, feats, training=True):
+ new_coords, cubes = construct_voxel_grid(coords)
+ new_feats = cubes_to_verts(new_coords.shape[0], cubes, feats)
+ if training:
+ con_loss = torch.mean((feats - new_feats[cubes]) ** 2)
+ else:
+ con_loss = 0.0
+ return new_coords, new_feats, con_loss
+
+
+def get_dense_attrs(coords : torch.Tensor, feats : torch.Tensor, res : int, sdf_init=True):
+ F = feats.shape[-1]
+ dense_attrs = torch.zeros([res] * 3 + [F], device=feats.device)
+ if sdf_init:
+ dense_attrs[..., 0] = 1 # initial outside sdf value
+ dense_attrs[coords[:, 0], coords[:, 1], coords[:, 2], :] = feats
+ return dense_attrs.reshape(-1, F)
+
+
+def get_defomed_verts(v_pos : torch.Tensor, deform : torch.Tensor, res):
+ return v_pos / res - 0.5 + (1 - 1e-8) / (res * 2) * torch.tanh(deform)
+
\ No newline at end of file
diff --git a/app.py b/app.py
new file mode 100755
index 0000000000000000000000000000000000000000..48a34194cda1aad0ffb9990fcc77be387cccc9a1
--- /dev/null
+++ b/app.py
@@ -0,0 +1,381 @@
+# MIT License
+
+# Copyright (c) Microsoft
+
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+# Copyright (c) [2025] [Microsoft]
+# Copyright (c) [2025] [Chongjie Ye]
+# SPDX-License-Identifier: MIT
+# This file has been modified by Chongjie Ye on 2025/04/10
+# Original file was released under MIT, with the full license text # available at https://github.com/atong01/conditional-flow-matching/blob/1.0.7/LICENSE.
+# This modified file is released under the same license.
+
+import gradio as gr
+import os
+os.environ['SPCONV_ALGO'] = 'native'
+from typing import *
+import torch
+import numpy as np
+from Stable3DGen.hi3dgen.pipelines import Hi3DGenPipeline
+import trimesh
+import tempfile
+from PIL import Image
+import glob
+from src.data import DemoData
+from src.models import LiNo_UniPS
+from torch.utils.data import DataLoader
+import pytorch_lightning as pl
+
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+os.makedirs(TMP_DIR, exist_ok=True)
+
+
+
+def cache_weights(weights_dir: str) -> dict:
+ import os
+ from huggingface_hub import snapshot_download
+
+ os.makedirs(weights_dir, exist_ok=True)
+ model_ids = [
+ "Stable-X/trellis-normal-v0-1",
+ ]
+ cached_paths = {}
+ for model_id in model_ids:
+ print(f"Caching weights for: {model_id}")
+ # Check if the model is already cached
+ local_path = os.path.join(weights_dir, model_id.split("/")[-1])
+ if os.path.exists(local_path):
+ print(f"Already cached at: {local_path}")
+ cached_paths[model_id] = local_path
+ continue
+ # Download the model and cache it
+ print(f"Downloading and caching model: {model_id}")
+ # Use snapshot_download to download the model
+ local_path = snapshot_download(repo_id=model_id, local_dir=os.path.join(weights_dir, model_id.split("/")[-1]), force_download=False)
+ cached_paths[model_id] = local_path
+ print(f"Cached at: {local_path}")
+
+ return cached_paths
+
+def preprocess_mesh(mesh_prompt):
+ print("Processing mesh")
+ trimesh_mesh = trimesh.load_mesh(mesh_prompt)
+ trimesh_mesh.export(mesh_prompt+'.glb')
+ return mesh_prompt+'.glb'
+
+def generate_3d(image, seed=-1,
+ ss_guidance_strength=3, ss_sampling_steps=50,
+ slat_guidance_strength=3, slat_sampling_steps=6,normal_bridge=None):
+ if image is None:
+ return None, None, None
+
+ if seed == -1:
+ seed = np.random.randint(0, MAX_SEED)
+
+ # image = hi3dgen_pipeline.preprocess_image(image, resolution=1024)
+ # normal_image = normal_predictor(image, resolution=768, match_input_resolution=True, data_type='object')
+ if normal_bridge is None:
+ return 0
+ mask = np.float32(np.abs(1 - np.sqrt(np.sum(normal_bridge * normal_bridge, axis=2))) < 0.5)[:,:,None]
+ normal_image = mask * (normal_bridge * 0.5 + 0.5)
+ normal_image = np.concatenate((normal_image,mask),axis=2)*255.0
+ normal_image = Image.fromarray(normal_image.astype(np.uint8),mode="RGBA")
+
+
+ outputs = hi3dgen_pipeline.run(
+ normal_image,
+ seed=seed,
+ formats=["mesh",],
+ preprocess_image=False,
+ sparse_structure_sampler_params={
+ "steps": ss_sampling_steps,
+ "cfg_strength": ss_guidance_strength,
+ },
+ slat_sampler_params={
+ "steps": slat_sampling_steps,
+ "cfg_strength": slat_guidance_strength,
+ },
+ )
+ generated_mesh = outputs['mesh'][0]
+
+ # Save outputs
+ import datetime
+ output_id = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+ os.makedirs(os.path.join(TMP_DIR, output_id), exist_ok=True)
+ mesh_path = f"{TMP_DIR}/{output_id}/mesh.glb"
+
+ # Export mesh
+ trimesh_mesh = generated_mesh.to_trimesh(transform_pose=True)
+
+ trimesh_mesh.export(mesh_path)
+
+ return mesh_path, mesh_path
+
+def predict_normal(input_images,input_mask):
+ test_dataset = DemoData(input_imgs_list=input_images,input_mask=input_mask)
+ test_loader = DataLoader(test_dataset, batch_size=1)
+
+ trainer = pl.Trainer(accelerator="auto", devices=1,precision="bf16-mixed")
+ nml_predict = trainer.predict(model=lino, dataloaders=test_loader)
+
+ nml_output = 0.5 * nml_predict[0] + 0.5
+
+ return ((nml_output*255.0).astype(np.uint8), nml_predict[0])
+
+def convert_mesh(mesh_path, export_format):
+ """Download the mesh in the selected format."""
+ if not mesh_path:
+ return None
+
+ # Create a temporary file to store the mesh data
+ temp_file = tempfile.NamedTemporaryFile(suffix=f".{export_format}", delete=False)
+ temp_file_path = temp_file.name
+
+ new_mesh_path = mesh_path.replace(".glb", f".{export_format}")
+ mesh = trimesh.load_mesh(mesh_path)
+ mesh.export(temp_file_path) # Export to the temporary file
+
+ return temp_file_path # Return the path to the temporary file
+
+def load_example_data(path,numberofimages):
+ path = os.path.join("demo", path)
+ mask_path = os.path.join(path,"mask.png")
+ image_pathes = glob.glob(os.path.join(path, f"L*")) + glob.glob(os.path.join(path, f"0*"))
+ image_pathes = image_pathes[:numberofimages]
+ input_images = []
+ for p in image_pathes:
+ input_images.append(Image.open(p))
+
+ if os.path.exists(mask_path):
+ input_mask = Image.open(mask_path)
+ else:
+ input_mask =Image.fromarray(np.ones_like(np.array(input_images[0])))
+ normal_path = os.path.join(path,"normal.png")
+ if os.path.exists(normal_path):
+ normal_gt = Image.open(normal_path)
+ else:
+ normal_gt = Image.fromarray(np.ones_like(np.array(input_images[0])))
+ return input_mask,input_images,normal_gt
+
+# Create the Gradio interface with improved layout
+with gr.Blocks(css="footer {visibility: hidden}") as demo:
+ gr.Markdown(
+ """
+ Light of Normals: Unified Feature Representation for Universal Photometric Stereo
+ """
+ )
+
+ with gr.Row():
+ gr.Markdown("""
+
+
+ 
+
+
+ 
+
+
+ 
+
+
+
+ """)
+ with gr.Row():
+ gr.Markdown(
+ """
+ LiNo-UniPS is a method for Univeral Photometric Stereo. It predicts the normal map from a given set of images. Key features include:
+
+ * **Light-Agnostic:** Does not require specific lighting parameters as input.
+ * **Arbitrary-Resolution:** Supports inputs of any resolution.
+ * **Mask-Free:** Also supports mask-free scene normal reconstruction.
+ """
+ )
+ with gr.Row():
+ gr.Markdown(
+ """
+ ### Getting Started:
+
+ 1. **Upload Your Data**: Use the "Upload Multi-light Images" button on the left to provide your input. For best results, we recommend providing 6 or more images.
+
+ 2. **Upload Your Mask (Optional)**: A mask is not required for scene reconstruction. However, to reconstruct the normal map for a specific **object**, providing a mask is highly recommended. Use the "Mask" button on the left.
+
+ 3. **Reconstruct**: Click the "Run" button to start the reconstruction process. You can use the slider in "Advanced Settings" to control the number of multi-light images used by LiNo-UniPS. Note: If the selected number exceeds the total number of uploaded images, the maximum available number will be used instead.
+
+ 4. **Visualize**: The result will appear in the "Normal Output" viewer on the right. If you use one of our provided examples that includes a ground truth normal map, it will be displayed in the "Ground Truth" viewer for comparison.
+
+ 5. **Generate Mesh (Optional)**: After the normal map is reconstructed, you can click the "Generate Mesh" button. This will use the predicted normal as a "normal bridge" to generate the corresponding 3D mesh via Hi3DGen. We recommend this step primarily for **objects**, as Hi3DGen is currently an object-level model.
+ """
+ )
+ with gr.Row():
+
+ with gr.Column(scale=1):
+ with gr.Tabs():
+ with gr.Tab("Input Images"):
+
+ with gr.Row():
+ input_mask = gr.Image(
+ label="Mask (Optional)",
+ type="pil",
+ height="300px",
+
+ )
+ input_images = gr.Gallery(
+ label="Upload Multi-light Images",
+ type="numpy",
+ columns=8,
+ object_fit="contain",
+ preview=True,
+ )
+
+
+ model_output = gr.Model3D(
+ label="3D Model Preview (Generated by Hi3DGen)",
+
+ )
+
+ with gr.Row():
+ export_format = gr.Dropdown(
+ choices=["obj", "glb", "ply", "stl"],
+ value="glb",
+ label="File Format",
+ scale=2
+ )
+ download_btn = gr.DownloadButton(
+ label="Export Mesh",
+ interactive=False,
+ scale=1
+ )
+
+ with gr.Column(scale=2):
+ with gr.Tabs():
+ with gr.Tab("LiNo-UniPS Output"):
+ with gr.Row(scale=3):
+ normal_output = gr.Image(label="Normal Output",height=700,)
+ normal_gt = gr.Image(label="Ground Truth",height=700)
+ with gr.Accordion("Advanced Settings", open=True):
+ numberofimages = gr.Slider(0, 100, label="Number of Images", value=16, step=1)
+
+ run_btn = gr.Button("Run", size="lg", variant="primary")
+ gen_shape_btn = gr.Button("Generate Mesh", size="lg", variant="primary")
+
+
+
+
+ seed = gr.Number(np.random.randint(0,1e10),visible=False)
+ ss_guidance_strength =gr.Number(3,visible=False)
+ ss_sampling_steps = gr.Number(50,visible=False)
+ slat_guidance_strength =gr.Number(3.0,visible=False)
+ slat_sampling_steps = gr.Number(6,visible=False)
+ normal_bridge = gr.State()
+
+ gen_shape_btn.click(
+ generate_3d,
+ inputs=[
+ input_images, seed,
+ ss_guidance_strength, ss_sampling_steps,
+ slat_guidance_strength, slat_sampling_steps,
+ normal_bridge
+ ],
+ outputs=[model_output, download_btn]
+ ).then(
+ lambda: gr.Button(interactive=True),
+ outputs=[download_btn],
+ )
+
+ run_btn.click(
+ predict_normal,
+ inputs=[
+ input_images,
+ input_mask
+ ],
+ outputs=[normal_output,normal_bridge],
+)
+
+ def update_download_button(mesh_path, export_format):
+ if not mesh_path:
+ return gr.File.update(value=None, interactive=False)
+
+ download_path = convert_mesh(mesh_path, export_format)
+ return download_path
+
+ export_format.change(
+ update_download_button,
+ inputs=[model_output, export_format],
+ outputs=[download_btn]
+ ).then(
+ lambda: gr.Button(interactive=True),
+ outputs=[download_btn],
+ )
+
+ example_display = gr.Image(visible=False,type="pil",label="Input images")
+ obj_path = gr.Textbox(label = "Name",visible=False)
+ num = gr.Textbox(label = "Maximum number of images",visible=False)
+ is_mask = gr.Textbox(label = "Mask",visible=False)
+ is_gt = gr.Textbox(label = "Normal ground truth",visible=False)
+ image_type = gr.Textbox(label = "Image type",visible=False)
+ image_resolution = gr.Textbox(label = "Image resolution",visible=False)
+
+
+ display_data = [
+
+ [Image.open("demo/basket/demo.png"), "basket", 8, False, False, "Real","960*960"],
+ [Image.open("demo/key/demo.png"), "key", 9, False, False, "Real","512*612"],
+ [Image.open("demo/ball/demo.png"), "ball", 96, True, True, "Real","512*612"],
+ [Image.open("demo/canandwood/demo.png"), "canandwood", 18, True, False, "Real","4032*2268"],
+ [Image.open("demo/cat/demo.png"), "cat", 96, True, True, "Real","512*612"],
+ [Image.open("demo/coins_and_keyboard/demo.png"), "coins_and_keyboard", 12, False, False, "Real","4000*4000"],
+ [Image.open("demo/owl/demo.png"), "owl", 13, True, False, "Real","2400*1600"],
+ [Image.open("demo/rabit/demo.png"), "rabit", 9, True, False, "Real","4000*4000"],
+ [Image.open("demo/reading/demo.png"), "reading", 96, True, True, "Real","512*612"],
+ ]
+ gr.Markdown(
+ """
+
+ Click any row to load an example.
+
+ """
+ )
+ gr.Examples(
+ examples=display_data,
+ inputs=[example_display,obj_path,num,is_mask,is_gt,image_type,image_resolution],
+ label="Examples"
+ )
+ example_display.change(
+ fn=load_example_data,
+ inputs=[obj_path,numberofimages],
+ outputs=[
+ input_mask,
+ input_images,
+ normal_gt
+ ]
+ )
+
+if __name__ == "__main__":
+ # Download and cache the weights
+ #cache_weights(WEIGHTS_DIR)
+
+ hi3dgen_pipeline = Hi3DGenPipeline.from_pretrained("weights/trellis-normal-v0-1")
+ hi3dgen_pipeline.cuda()
+ lino = LiNo_UniPS()
+ lino.from_pretrained("weights/lino/lino.pth")
+
+ demo.launch(share=False, server_name="0.0.0.0")
+
diff --git a/src/data/__init__.py b/src/data/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..35269d39fa4b3a35f0873f38ed469fb9d057f9ea
--- /dev/null
+++ b/src/data/__init__.py
@@ -0,0 +1,2 @@
+from .data_module import DemoData
+from .data_module import TestData
\ No newline at end of file
diff --git a/src/data/__pycache__/__init__.cpython-310.pyc b/src/data/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0eb6f39cdf7cd1bc4a7b3227ad40591a82a1bac7
Binary files /dev/null and b/src/data/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/data/__pycache__/data_module.cpython-310.pyc b/src/data/__pycache__/data_module.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f7ae4bc2331caae006add230fa058313f72c2f96
Binary files /dev/null and b/src/data/__pycache__/data_module.cpython-310.pyc differ
diff --git a/src/data/__pycache__/testdata_nomask_datamodule.cpython-310.pyc b/src/data/__pycache__/testdata_nomask_datamodule.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..88facec5d5e8dd0e18cf73bda59d61da67dab16c
Binary files /dev/null and b/src/data/__pycache__/testdata_nomask_datamodule.cpython-310.pyc differ
diff --git a/src/data/data_module.py b/src/data/data_module.py
new file mode 100644
index 0000000000000000000000000000000000000000..85b3b6bb6574bb17ffb3b3f1db8fbcd3071e9823
--- /dev/null
+++ b/src/data/data_module.py
@@ -0,0 +1,265 @@
+
+import numpy as np
+from torch.utils.data import Dataset
+import cv2
+import glob
+import os
+
+def get_roi(mask, margin=8):
+ """
+ """
+ h0, w0 = mask.shape[:2]
+
+ if mask is not None:
+ rows, cols = np.nonzero(mask)
+ rowmin, rowmax = np.min(rows), np.max(rows)
+ colmin, colmax = np.min(cols), np.max(cols)
+ row, col = rowmax - rowmin, colmax - colmin
+
+ flag = not (rowmin - margin <= 0 or rowmax + margin > h0 or
+ colmin - margin <= 0 or colmax + margin > w0)
+
+ if row > col and flag:
+ r_s, r_e = rowmin - margin, rowmax + margin
+ c_s, c_e = max(colmin - int(0.5 * (row - col)) - margin, 0), \
+ min(colmax + int(0.5 * (row - col)) + margin, w0)
+ elif col >= row and flag:
+ r_s, r_e = max(rowmin - int(0.5 * (col - row)) - margin, 0), \
+ min(rowmax + int(0.5 * (col - row)) + margin, h0)
+ c_s, c_e = colmin - margin, colmax + margin
+ else:
+ r_s, r_e, c_s, c_e = 0, h0, 0, w0
+ else:
+ r_s, r_e, c_s, c_e = 0, h0, 0, w0
+
+ return np.array([h0, w0, r_s, r_e, c_s, c_e])
+
+def crop_and_resize_img(img, roi, max_image_resolution=6000):
+
+
+ h0, w0, r_s, r_e, c_s, c_e = roi
+
+ img = img[r_s:r_e, c_s:c_e, :]
+
+
+ h = max(512, min(max_image_resolution, (max(img.shape[:2]) // 512) * 512))
+ w = h
+
+ img = cv2.resize(img, (w, h), interpolation=cv2.INTER_CUBIC)
+
+
+ bit_depth = 255.0 if img.dtype == np.uint8 else 65535.0 if img.dtype == np.uint16 else 1.0
+ img = np.float32(img) / bit_depth
+
+ return img
+
+def crop_and_resize_mask(mask, roi, max_image_resolution=6000):
+
+
+ h0, w0, r_s, r_e, c_s, c_e = roi
+
+
+ mask = mask[r_s:r_e, c_s:c_e]
+
+ h = max(512, min(max_image_resolution, (max(mask.shape[:2]) // 512) * 512))
+ w = h
+
+
+ mask = np.float32(cv2.resize(mask, (w, h), interpolation=cv2.INTER_CUBIC) > 0.5)
+
+ return mask
+
+class DemoData(Dataset):
+ def __init__(self,input_imgs_list,input_mask):
+ self.input_imgs_list = input_imgs_list
+ self.input_mask = input_mask
+ def __len__(self):
+ return 1
+ def load(self,input_images_list,mask):
+ if mask is None:
+ mask = np.ones_like(input_images_list[0][0])
+ else:
+ mask = np.array(mask)
+ mask = mask[:,:,0]
+ if mask.max() <= 1.0:
+ self.mask_original = mask[:,:,None]
+ else:
+ self.mask_original = mask[:,:,None] / 255.0
+ self.roi = get_roi(mask)
+ for i in range(len(input_images_list)):
+ img = input_images_list[i]
+ input_images_list[i]= crop_and_resize_img(img[0], self.roi)
+ I = np.array(input_images_list)
+ numberofimages,h,w,_ = I.shape
+ mask = crop_and_resize_mask(mask, self.roi)
+ I = np.reshape(I, (-1, h * w, 3))
+ temp = np.mean(I[:, mask.flatten()==1,:], axis=2)
+ mx = np.max(temp, axis=1)
+ temp = mx
+ I /= (temp.reshape(-1,1,1) + 1.0e-6)
+ I = np.transpose(I, (1, 2, 0))
+ I = I.reshape(h, w, 3, numberofimages)
+ mask = (mask.reshape(h, w, 1)).astype(np.float32)
+ h = mask.shape[0]
+ w = mask.shape[1]
+ self.h = h
+ self.w = w
+ self.I = I
+ self.N = np.ones((h, w, 3), np.float32)
+ self.mask = mask
+ return 1
+ def __getitem__(self, idx):
+ self.load(self.input_imgs_list,self.input_mask)
+ return {
+ "imgs":self.I.transpose(2,0,1,3),
+ "mask":self.mask.transpose(2,0,1),
+ "mask_original":self.mask_original.transpose(2,0,1),
+ "roi":self.roi
+ }
+
+class TestData(Dataset):
+ def __init__(
+ self,
+ data_root: list = None,
+ numofimages: int = 16
+
+ ):
+ self.data_root = data_root
+ self.numberOfImages = numofimages
+ self.objlist = []
+ for i in range(len(self.data_root)):
+ with os.scandir(self.data_root[i]) as entries:
+ self.objlist += [entry.path for entry in entries if entry.is_dir()]
+ print(f"[Dataset] => {len(self.objlist)} items selected.")
+ objlist = self.objlist
+ total = len(objlist)
+ indices = list(range(total))
+ self.objlist = [objlist[i] for i in indices]
+ print(f"Test, => {len(self.objlist)} items selected.")
+ def load(self, objlist, dirid):
+ obj_path = objlist[dirid]
+ if "DiLiGenT" in obj_path:
+ nml_path = os.path.join(obj_path, "Normal_gt.png")
+ if "10" not in obj_path: # diligent
+ directlist = sorted(glob.glob(os.path.join(obj_path, f"0*")))
+ else: # diligent100
+ directlist = sorted([
+ path for path in glob.glob(os.path.join(obj_path, "*.png"))
+ if not os.path.basename(path).lower() == "mask.png"
+ ])
+ elif "LUCES" in obj_path:
+ nml_path = os.path.join(obj_path, "normals.png")
+ directlist = sorted([
+ f for i in range(1, 52) for f in glob.glob(os.path.join(obj_path, f"{i:02d}*"))
+ ])
+ elif "Real" in obj_path:
+ nml_path = os.path.join(obj_path, "Normal_gt.png")
+ directlist = sorted(glob.glob(os.path.join(obj_path, f"L*")))
+ else:
+ print(f"error:unknown dataset{obj_path}")
+ return 0
+
+
+ num_images_to_sample = self.numberOfImages
+ if num_images_to_sample is not None and num_images_to_sample < len(directlist):
+ indexset = np.random.permutation(len(directlist))[:num_images_to_sample]
+ else:
+ indexset = range(len(directlist))
+
+ I = None
+ mask = None
+ N = None
+ n_true = None
+
+ for i, indexofimage in enumerate(indexset):
+ img_path = directlist[indexofimage]
+ read_img = cv2.imread(img_path, flags=cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH)
+ if read_img is None:
+ print(f"warning: can not read {img_path}")
+ return 0
+
+ img = cv2.cvtColor(read_img, cv2.COLOR_BGR2RGB)
+ if i == 0:
+
+ mask_path = os.path.join(obj_path, "mask.png")
+ if os.path.exists(mask_path):
+ mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) / 255.0
+ else:
+ mask = np.ones_like(read_img)[:,:,0]
+
+ if os.path.exists(nml_path):
+ bit_depth = 65535.0 if "LUCES" in obj_path else 255.0
+ N = cv2.cvtColor(cv2.imread(nml_path, flags=cv2.IMREAD_ANYCOLOR | cv2.IMREAD_ANYDEPTH), cv2.COLOR_BGR2RGB) / bit_depth
+ N = 2 * N - 1
+ N = N / np.linalg.norm(N, axis=2, keepdims=True)
+ N = N * mask[:, :, np.newaxis]
+ n_true = N
+
+
+ self.roi = get_roi(mask)
+ mask = crop_and_resize_mask(mask, self.roi)
+
+
+ img= crop_and_resize_img(img, self.roi)
+ h, w = img.shape[:2]
+ if i == 0:
+ I = np.zeros((len(indexset), h, w, 3), np.float32)
+ I[i, :, :, :] = img
+
+
+ imgs_ = I.copy()
+ I = np.reshape(I, (-1, h * w, 3))
+
+ """Data Normalization"""
+ temp = np.mean(I[:, mask.flatten()==1,:], axis=2)
+ mean = np.mean(temp, axis=1)
+ mx = np.max(temp, axis=1)
+ scale = np.random.rand(I.shape[0],)
+ temp = (1-scale) * mean + scale * mx
+ imgs_ /= (temp.reshape(-1,1,1,1) + 1.0e-6)
+ I = imgs_
+ I = np.transpose(I, (1, 2, 3, 0))
+ mask = (mask.reshape(h, w, 1)).astype(np.float32)
+ h = mask.shape[0]
+ w = mask.shape[1]
+ self.h = h
+ self.w = w
+ self.I = I #
+ if ("DiLiGenT" in obj_path and "10" in obj_path) or "Real" in obj_path: # diligent100
+ self.N = np.ones((h,w,3,1))
+ else:
+ self.N = n_true[:,:,:,np.newaxis]
+ self.mask = mask
+ self.directlist = directlist
+
+ return 1
+
+
+ def __getitem__(self, index_):
+ objid = index_
+ while 1:
+ success = self.load(self.objlist, objid)
+ if success:
+ break
+ else:
+ objid = np.random.randint(0, len(self.objlist))
+ img = self.I.transpose(2,0,1,3) # 3 h w Nmax
+ nml = self.N.transpose(2,0,1,3) # 3 h w 1
+ objname = os.path.basename(os.path.basename(self.objlist[objid]))
+ numberOfImages = self.numberOfImages
+ try:
+ output = {
+ 'imgs': img,
+ 'nml': nml,
+ "mask":self.mask.transpose(2,0,1),
+ 'directlist': self.directlist,
+ 'objname': objname,
+ 'numberOfImages': numberOfImages,
+ "roi":self.roi
+ }
+ return output
+ except:
+ raise KeyError
+
+ def __len__(self):
+ return len(self.objlist)
\ No newline at end of file
diff --git a/src/models/Net_module.py b/src/models/Net_module.py
new file mode 100755
index 0000000000000000000000000000000000000000..be62a20d8b5d1ab95b5b32ce3aae9be517edb594
--- /dev/null
+++ b/src/models/Net_module.py
@@ -0,0 +1,253 @@
+import torchvision
+import torch
+from torchmetrics import MeanMetric
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .module.utils import *
+from .utils import decompose_tensors
+from .utils import gauss_filter
+import cv2
+import pytorch_lightning as pl
+from src.models.utils.compute_mae import compute_mae_np
+from datetime import datetime
+
+class LiNo_UniPS(pl.LightningModule):
+ def __init__(self,
+ pixel_samples: int = 2048,
+ task_name :str = None):
+ super().__init__()
+ self.pixel_samples = pixel_samples
+ self.task_name = task_name
+ self.input_dim = 4
+ self.image_encoder = ScaleInvariantSpatialLightImageEncoder(self.input_dim, use_efficient_attention=False)
+ self.input_dim = 0
+ self.glc_upsample = GLC_Upsample(256+self.input_dim, num_enc_sab=1, dim_hidden=256, dim_feedforward=1024, use_efficient_attention=True)
+ self.glc_aggregation = GLC_Aggregation(256+self.input_dim, num_agg_transformer=2, dim_aggout=384, dim_feedforward=1024, use_efficient_attention=False)
+ self.img_embedding = nn.Sequential(
+ nn.Linear(3,32),
+ nn.LeakyReLU(),
+ nn.Linear(32, 256)
+ )
+ self.regressor = Regressor(384, num_enc_sab=1, use_efficient_attention=True, dim_feedforward=1024)
+ self.test_mae = MeanMetric()
+ self.test_loss = MeanMetric()
+ def on_test_start(self):
+
+ timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ self.run_save_dir = f'output/{timestamp}/{self.task_name}/results/'
+ os.makedirs(self.run_save_dir, exist_ok=True)
+
+ def from_pretrained(self,pth_path):
+ pretrain_weight = torch.load(pth_path,weights_only=False)
+ self.load_state_dict(pretrain_weight, strict=False)
+ def _prepare_test_inputs(self, batch):
+ img = batch["imgs"].to(torch.bfloat16)
+ self.numberofImages = img.shape[-1]
+ print("number of test images", self.numberofImages)
+ nml = batch["nml"].to(torch.bfloat16)
+ directlist = batch["directlist"]
+ roi = batch.get("roi",None)
+ roi = roi[0].cpu().numpy()
+ return img, nml,directlist,roi
+ def _postprocess_prediction(self, nml_predict_raw, nml_gt_raw, roi):
+
+ h_orig, w_orig, r_s, r_e, c_s, c_e = roi
+ nml_predict = nml_predict_raw.squeeze().permute(1, 2, 0).cpu().numpy()
+ nml_predict = cv2.resize(nml_predict, dsize=(c_e - c_s, r_e - r_s), interpolation=cv2.INTER_AREA)
+ mask = np.float32(np.abs(1 - np.sqrt(np.sum(nml_predict * nml_predict, axis=2))) < 0.5)
+ nml_predict = np.divide(nml_predict, np.linalg.norm(nml_predict, axis=2, keepdims=True) + 1e-12)
+ nml_predict = nml_predict * mask[:, :, np.newaxis]
+ nout = np.zeros((h_orig, w_orig, 3), np.float32)
+ nout[r_s:r_e, c_s:c_e, :] = nml_predict
+
+ nml_gt = nml_gt_raw.squeeze().permute(1, 2, 0).float().cpu().numpy()
+ mask_gt = np.float32(np.abs(1 - np.sqrt(np.sum(nml_gt * nml_gt, axis=2))) < 0.5)
+
+ return nout, nml_gt, mask_gt
+
+ def _calculate_and_log_metrics(self, nout, nml_gt, mask_gt):
+ mse = torch.nn.MSELoss()(torch.tensor(nout).to(self.device), torch.tensor(nml_gt).to(self.device))
+
+ mae, emap = compute_mae_np(nout, nml_gt, mask_gt)
+
+ self.test_loss(mse)
+ self.test_mae(mae)
+ self.log("test/mse", self.test_loss, on_step=False, on_epoch=True, prog_bar=True)
+ self.log("test/mae", self.test_mae, on_step=False, on_epoch=True, prog_bar=True)
+
+ return mse, mae, emap
+
+ def _save_test_results(self, nout, nml_gt, emap, img, loss, mae, directlist, save_dir):
+
+ obj_name_parts = os.path.dirname(directlist[0][0]).split('/')
+ obj_name = obj_name_parts[-1]
+
+
+ save_path = os.path.join(save_dir,f'{self.numberofImages}',f'{obj_name}')
+ os.makedirs(save_path, exist_ok=True)
+ print(f"save to: {save_path}")
+ if ("DiLiGenT_100" not in self.task_name) and ("Real" not in self.task_name):
+ nout_to_save = (nout + 1) / 2
+ nml_gt_to_save = (nml_gt + 1) / 2
+
+ emap_to_save = emap.astype(np.float32).squeeze()
+ thresh = 90
+ emap_to_save[emap_to_save >= thresh] = thresh
+ emap_to_save = emap_to_save / thresh
+
+ plt.imsave(save_path + '/nml_predict.png', np.clip(nout_to_save, 0, 1))
+ plt.imsave(save_path + '/nml_gt.png', np.clip(nml_gt_to_save, 0, 1))
+ plt.imsave(save_path + '/error_map.png', emap_to_save, cmap='jet')
+ torchvision.utils.save_image(img.squeeze(0).permute(3,0,1,2), save_path + '/tiled.png')
+
+ fig, axes = plt.subplots(1, 3, figsize=(12, 4))
+ axes[0].imshow(np.clip(nout_to_save, 0, 1))
+ axes[0].set_title('Prediction'); axes[0].axis('off')
+ axes[1].imshow(np.clip(nml_gt_to_save, 0, 1))
+ axes[1].set_title('Ground Truth'); axes[1].axis('off')
+ axes[2].imshow(emap, cmap='jet')
+ axes[2].set_title('Error Map'); axes[2].axis('off')
+ plt.figtext(0.5, 0.02, f'Loss: {loss:.4f} | MAE: {mae:.4f}', ha='center', fontsize=12)
+ plt.tight_layout()
+ plt.savefig(save_path + '/combined.png', dpi=300)
+ plt.close(fig)
+
+
+ with open(save_path + '/result.txt', 'w') as f:
+ f.write(f"loss: {loss.item()}\n")
+ f.write(f"mae: {mae}\n")
+
+ print(f"Done for {obj_name}")
+ else:
+ if "DiLiGenT_100" in self.task_name:
+ from scipy.io import savemat
+ mat_save_path = os.path.join(os.path.dirname(save_path),"submit")
+ os.makedirs(mat_save_path,exist_ok=True)
+ normal_map = nout
+ savemat(mat_save_path + "/" + obj_name + '.mat', {'Normal_est': normal_map})
+ torchvision.utils.save_image(img.squeeze(0).permute(3,0,1,2), save_path + '/tiled.png')
+ nout = (nout + 1) / 2
+ plt.imsave(save_path + '/nml_predict.png', nout)
+
+ def test_step(self, batch, batch_idx):
+
+ img, nml_gt, directlist, roi = self._prepare_test_inputs(batch)
+
+ nml_predict = self.model_step(batch)
+
+
+ nout, nml_gt, mask_gt = self._postprocess_prediction(nml_predict, nml_gt, roi)
+ if ("DiLiGenT_100" not in self.task_name) and ("Real" not in self.task_name):
+ loss, mae, emap = self._calculate_and_log_metrics(nout, nml_gt, mask_gt)
+ print(f"{os.path.basename(os.path.dirname(directlist[0][0]))} | MAE: {mae:.4f}")
+ self._save_test_results(nout, nml_gt, emap, img, loss, mae, directlist, self.run_save_dir)
+ else:
+ emap,loss,mae = None,None,None
+ self._save_test_results(nout, nml_gt, emap, img, loss, mae, directlist, self.run_save_dir)
+
+ def predict_step(self,batch):
+ roi = batch.get("roi",None)
+ nml_predict = self.model_step(batch)
+ roi = roi[0].cpu().numpy()
+ h_ = roi[0]
+ w_ = roi[1]
+ r_s = roi[2]
+ r_e = roi[3]
+ c_s = roi[4]
+ c_e = roi[5]
+ nml_predict = nml_predict.squeeze().permute(1,2,0).cpu().numpy()
+ nml_predict = cv2.resize(nml_predict, dsize=(c_e-c_s, r_e-r_s), interpolation=cv2.INTER_AREA)
+ nml_predict = np.divide(nml_predict, np.linalg.norm(nml_predict, axis=2, keepdims=True) + 1.0e-12)
+ mask = np.float32(np.abs(1 - np.sqrt(np.sum(nml_predict * nml_predict, axis=2))) < 0.5)
+ nml_predict = nml_predict * mask[:, :, np.newaxis]
+ nout = np.zeros((h_, w_, 3), np.float32)
+ nout[r_s:r_e, c_s:c_e,:] = nml_predict
+ mask = batch["mask_original"].squeeze().cpu().numpy()[:,:,None]
+
+ return nout*mask
+ def model_step(self,batch):
+ I = batch.get("imgs",None)
+ M = batch.get("mask",None)
+ # roi = batch.get("roi",None)
+ B, C, H, W, Nmax = I.shape
+
+ patch_size = 512
+ patches_I = decompose_tensors.divide_tensor_spatial(I.permute(0,4,1,2,3).reshape(-1, C, H, W), block_size=patch_size, method='tile_stride')
+ patches_I = patches_I.reshape(B, Nmax, -1, C, patch_size, patch_size).permute(0, 2, 3, 4, 5, 1)
+ sliding_blocks = patches_I.shape[1]
+ patches_M = decompose_tensors.divide_tensor_spatial(M, block_size=patch_size, method='tile_stride')
+ patches_nml = []
+
+ nImgArray = np.array([Nmax])
+ canonical_resolution = 256
+ for k in range(sliding_blocks):
+ """ Image Encoder at Canonical Resolution """
+ print("please wait for a moment, it may take a while")
+ I = patches_I[:, k, :, :, :, :]
+ M = patches_M[:, k, :, :, :]
+ B, C, H, W, Nmax = I.shape
+ decoder_resolution = H
+ I_enc = I.permute(0, 4, 1, 2, 3)
+ M_enc = M
+ img_index = make_index_list(Nmax, nImgArray)
+ I_enc = I_enc.reshape(-1, I_enc.shape[2], I_enc.shape[3], I_enc.shape[4])
+ M_enc = M_enc.unsqueeze(1).expand(-1, Nmax, -1, -1, -1).reshape(-1, 1, H, W)
+ data = I_enc * M_enc
+ data = data[img_index==1,:,:,:]
+ glc,_= self.image_encoder(data, nImgArray, canonical_resolution)
+ I_dec = []
+ M_dec = []
+ img = I.permute(0, 4, 1, 2, 3)
+ """ Sample Decoder at Original Resokution"""
+ img = img.squeeze()
+ I_dec = F.interpolate(img.float(), size=(decoder_resolution, decoder_resolution), mode='bilinear', align_corners=False).to(torch.bfloat16)
+ M_dec = F.interpolate(M.float(), size=(decoder_resolution, decoder_resolution), mode='nearest').to(torch.bfloat16)
+ decoder_imgsize = (decoder_resolution, decoder_resolution)
+ C = img.shape[1]
+ H = decoder_imgsize[0]
+ W = decoder_imgsize[1]
+ nout = torch.zeros(B, H * W, 3).to(I.device)
+ f_scale = decoder_resolution//canonical_resolution
+ smoothing = gauss_filter.gauss_filter(glc.shape[1], 10 * f_scale+1, 1).to(glc.device)
+ chunk_size = 16
+ processed_chunks = []
+ for glc_chunk in torch.split(glc, chunk_size, dim=0):
+ smoothed_chunk = smoothing(glc_chunk)
+ processed_chunks.append(smoothed_chunk)
+ glc = torch.cat(processed_chunks, dim=0)
+ del M
+ _, _, H, W = I_dec.shape
+ p = 0
+ nout = torch.zeros(B, H * W, 3).to(I.device)
+ conf_out = torch.zeros(B, H * W, 1).to(I.device)
+ for b in range(B):
+ target = range(p, p+nImgArray[b])
+ p = p+nImgArray[b]
+ m_ = M_dec[b, :, :, :].reshape(-1, H * W).permute(1,0)
+ ids = np.nonzero(m_>0)[:,0]
+ ids = ids[np.random.permutation(len(ids))]
+ ids_shuffle = ids[np.random.permutation(len(ids))]
+ num_split = len(ids) // self.pixel_samples + 1
+ idset = np.array_split(ids_shuffle, num_split)
+ o_ = I_dec[target, :, :, :].reshape(nImgArray[b], C, H * W).permute(2,0,1)
+ for ids in idset:
+ o_ids = o_[ids, :, :]
+ glc_ids = glc[target, :, :, :].permute(2,3,0,1).flatten(0,1)[ids,:,:]
+ o_ids = self.img_embedding(o_ids)
+ x = o_ids + glc_ids
+ glc_ids = self.glc_upsample(x)
+ x = o_ids + glc_ids
+ x = self.glc_aggregation(x)
+ x_n, _, _, conf = self.regressor(x, len(ids))
+ x_n = F.normalize(x_n, p=2, dim=-1)
+ nout[b, ids, :] = x_n[b,:,:]
+ conf_out[b, ids, :] = conf[b,:,:].to(torch.float32)
+ nout = nout.reshape(B,H,W,3).permute(0,3,1,2)
+ conf_out = conf_out.reshape(B,H,W,1).permute(0,3,1,2)
+ patches_nml.append(nout)
+
+ patches_nml = torch.stack(patches_nml, dim=1)
+ merged_tensor_nml = decompose_tensors.merge_tensor_spatial(patches_nml.permute(1,0,2,3,4), method='tile_stride')
+ return merged_tensor_nml
diff --git a/src/models/__init__.py b/src/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..733815d09d50fad441ec9adf34b2fbe5145ae242
--- /dev/null
+++ b/src/models/__init__.py
@@ -0,0 +1 @@
+from .Net_module import LiNo_UniPS
\ No newline at end of file
diff --git a/src/models/__pycache__/Net_module.cpython-310.pyc b/src/models/__pycache__/Net_module.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f0856656c4e47f1f125a1dbafa6535c57cc8b438
Binary files /dev/null and b/src/models/__pycache__/Net_module.cpython-310.pyc differ
diff --git a/src/models/__pycache__/__init__.cpython-310.pyc b/src/models/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4225c265b119668961e22a8227ae237d2bfc20b1
Binary files /dev/null and b/src/models/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/aggregator/__init__.py b/src/models/aggregator/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0de9a6d84ebcd41b83446138a733d74fa0ff959c
--- /dev/null
+++ b/src/models/aggregator/__init__.py
@@ -0,0 +1 @@
+from .models.aggregator import Aggregator
\ No newline at end of file
diff --git a/src/models/aggregator/__pycache__/__init__.cpython-310.pyc b/src/models/aggregator/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cd8522f94222c49bc5022bb925aa599f695dd7c9
Binary files /dev/null and b/src/models/aggregator/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/camera_head.cpython-310.pyc b/src/models/aggregator/heads/__pycache__/camera_head.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2da068ef02926ea53139987c6f97c8ba217f2d4e
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/camera_head.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/camera_head.cpython-38.pyc b/src/models/aggregator/heads/__pycache__/camera_head.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9ac4b9dc32263e1d75fc302eea8120a15da33041
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/camera_head.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/dpt_head.cpython-310.pyc b/src/models/aggregator/heads/__pycache__/dpt_head.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a51cb67b5b589994c2acebca13bcfa9d5958d6a2
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/dpt_head.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/dpt_head.cpython-38.pyc b/src/models/aggregator/heads/__pycache__/dpt_head.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9a8d6513e16bf3c55cf29cb706df369fe9e21d25
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/dpt_head.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/head_act.cpython-310.pyc b/src/models/aggregator/heads/__pycache__/head_act.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b2e1674021e1499961f62f4c2a059cb473c6e541
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/head_act.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/head_act.cpython-38.pyc b/src/models/aggregator/heads/__pycache__/head_act.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ac9c2a9ebc2e51819637a3533b90536aff305c17
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/head_act.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/track_head.cpython-310.pyc b/src/models/aggregator/heads/__pycache__/track_head.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a60a3b76f2050de7551d296ed14ac5e29213b713
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/track_head.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/track_head.cpython-38.pyc b/src/models/aggregator/heads/__pycache__/track_head.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5b01ce9d8e941cd4d6689f015fedad0b04cbbb4f
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/track_head.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/utils.cpython-310.pyc b/src/models/aggregator/heads/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..45e3b576717c8ced67cbb3f024f2d512f28492d3
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/utils.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/__pycache__/utils.cpython-38.pyc b/src/models/aggregator/heads/__pycache__/utils.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5f20eff68c2d49843f901faf762545cbe8d4bacf
Binary files /dev/null and b/src/models/aggregator/heads/__pycache__/utils.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/camera_head.py b/src/models/aggregator/heads/camera_head.py
new file mode 100644
index 0000000000000000000000000000000000000000..153b0208c098224ef5bfc3cc3509ba9016970831
--- /dev/null
+++ b/src/models/aggregator/heads/camera_head.py
@@ -0,0 +1,162 @@
+# 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.
+
+import math
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..layers import Mlp
+from ..layers.block import Block
+from ..heads.head_act import activate_pose
+
+
+class CameraHead(nn.Module):
+ """
+ CameraHead predicts camera parameters from token representations using iterative refinement.
+
+ It applies a series of transformer blocks (the "trunk") to dedicated camera tokens.
+ """
+
+ def __init__(
+ self,
+ dim_in: int = 2048,
+ trunk_depth: int = 4,
+ pose_encoding_type: str = "absT_quaR_FoV",
+ num_heads: int = 16,
+ mlp_ratio: int = 4,
+ init_values: float = 0.01,
+ trans_act: str = "linear",
+ quat_act: str = "linear",
+ fl_act: str = "relu", # Field of view activations: ensures FOV values are positive.
+ ):
+ super().__init__()
+
+ if pose_encoding_type == "absT_quaR_FoV":
+ self.target_dim = 9
+ else:
+ raise ValueError(f"Unsupported camera encoding type: {pose_encoding_type}")
+
+ self.trans_act = trans_act
+ self.quat_act = quat_act
+ self.fl_act = fl_act
+ self.trunk_depth = trunk_depth
+
+ # Build the trunk using a sequence of transformer blocks.
+ self.trunk = nn.Sequential(
+ *[
+ Block(
+ dim=dim_in,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ init_values=init_values,
+ )
+ for _ in range(trunk_depth)
+ ]
+ )
+
+ # Normalizations for camera token and trunk output.
+ self.token_norm = nn.LayerNorm(dim_in)
+ self.trunk_norm = nn.LayerNorm(dim_in)
+
+ # Learnable empty camera pose token.
+ self.empty_pose_tokens = nn.Parameter(torch.zeros(1, 1, self.target_dim))
+ self.embed_pose = nn.Linear(self.target_dim, dim_in)
+
+ # Module for producing modulation parameters: shift, scale, and a gate.
+ self.poseLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim_in, 3 * dim_in, bias=True))
+
+ # Adaptive layer normalization without affine parameters.
+ self.adaln_norm = nn.LayerNorm(dim_in, elementwise_affine=False, eps=1e-6)
+ self.pose_branch = Mlp(
+ in_features=dim_in,
+ hidden_features=dim_in // 2,
+ out_features=self.target_dim,
+ drop=0,
+ )
+
+ def forward(self, aggregated_tokens_list: list, num_iterations: int = 4) -> list:
+ """
+ Forward pass to predict camera parameters.
+
+ Args:
+ aggregated_tokens_list (list): List of token tensors from the network;
+ the last tensor is used for prediction.
+ num_iterations (int, optional): Number of iterative refinement steps. Defaults to 4.
+
+ Returns:
+ list: A list of predicted camera encodings (post-activation) from each iteration.
+ """
+ # Use tokens from the last block for camera prediction.
+ tokens = aggregated_tokens_list[-1]
+
+ # Extract the camera tokens
+ pose_tokens = tokens[:, :, 0]
+ pose_tokens = self.token_norm(pose_tokens)
+
+ pred_pose_enc_list = self.trunk_fn(pose_tokens, num_iterations)
+ return pred_pose_enc_list
+
+ def trunk_fn(self, pose_tokens: torch.Tensor, num_iterations: int) -> list:
+ """
+ Iteratively refine camera pose predictions.
+
+ Args:
+ pose_tokens (torch.Tensor): Normalized camera tokens with shape [B, 1, C].
+ num_iterations (int): Number of refinement iterations.
+
+ Returns:
+ list: List of activated camera encodings from each iteration.
+ """
+ B, S, C = pose_tokens.shape # S is expected to be 1.
+ pred_pose_enc = None
+ pred_pose_enc_list = []
+
+ for _ in range(num_iterations):
+ # Use a learned empty pose for the first iteration.
+ if pred_pose_enc is None:
+ module_input = self.embed_pose(self.empty_pose_tokens.expand(B, S, -1))
+ else:
+ # Detach the previous prediction to avoid backprop through time.
+ pred_pose_enc = pred_pose_enc.detach()
+ module_input = self.embed_pose(pred_pose_enc)
+
+ # Generate modulation parameters and split them into shift, scale, and gate components.
+ shift_msa, scale_msa, gate_msa = self.poseLN_modulation(module_input).chunk(3, dim=-1)
+
+ # Adaptive layer normalization and modulation.
+ pose_tokens_modulated = gate_msa * modulate(self.adaln_norm(pose_tokens), shift_msa, scale_msa)
+ pose_tokens_modulated = pose_tokens_modulated + pose_tokens
+
+ pose_tokens_modulated = self.trunk(pose_tokens_modulated)
+ # Compute the delta update for the pose encoding.
+ pred_pose_enc_delta = self.pose_branch(self.trunk_norm(pose_tokens_modulated))
+
+ if pred_pose_enc is None:
+ pred_pose_enc = pred_pose_enc_delta
+ else:
+ pred_pose_enc = pred_pose_enc + pred_pose_enc_delta
+
+ # Apply final activation functions for translation, quaternion, and field-of-view.
+ activated_pose = activate_pose(
+ pred_pose_enc,
+ trans_act=self.trans_act,
+ quat_act=self.quat_act,
+ fl_act=self.fl_act,
+ )
+ pred_pose_enc_list.append(activated_pose)
+
+ return pred_pose_enc_list
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+ """
+ Modulate the input tensor using scaling and shifting parameters.
+ """
+ # modified from https://github.com/facebookresearch/DiT/blob/796c29e532f47bba17c5b9c5eb39b9354b8b7c64/models.py#L19
+ return x * (1 + scale) + shift
diff --git a/src/models/aggregator/heads/dpt_head.py b/src/models/aggregator/heads/dpt_head.py
new file mode 100644
index 0000000000000000000000000000000000000000..60ba169d81767e39f3ab1e9f06ce406645a76f9f
--- /dev/null
+++ b/src/models/aggregator/heads/dpt_head.py
@@ -0,0 +1,545 @@
+# 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.
+
+
+# Inspired by https://github.com/DepthAnything/Depth-Anything-V2
+
+
+import os
+from typing import List, Dict, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .head_act import activate_head
+from .utils import create_uv_grid, position_grid_to_embed
+
+
+class DPTHead(nn.Module):
+ """
+ DPT Head for dense prediction tasks.
+
+ This implementation follows the architecture described in "Vision Transformers for Dense Prediction"
+ (https://arxiv.org/abs/2103.13413). The DPT head processes features from a vision transformer
+ backbone and produces dense predictions by fusing multi-scale features.
+
+ Args:
+ dim_in (int): Input dimension (channels).
+ patch_size (int, optional): Patch size. Default is 14.
+ output_dim (int, optional): Number of output channels. Default is 4.
+ activation (str, optional): Activation type. Default is "inv_log".
+ conf_activation (str, optional): Confidence activation type. Default is "expp1".
+ features (int, optional): Feature channels for intermediate representations. Default is 256.
+ out_channels (List[int], optional): Output channels for each intermediate layer.
+ intermediate_layer_idx (List[int], optional): Indices of layers from aggregated tokens used for DPT.
+ pos_embed (bool, optional): Whether to use positional embedding. Default is True.
+ feature_only (bool, optional): If True, return features only without the last several layers and activation head. Default is False.
+ down_ratio (int, optional): Downscaling factor for the output resolution. Default is 1.
+ """
+
+ def __init__(
+ self,
+ dim_in: int,
+ patch_size: int = 14,
+ output_dim: int = 4,
+ activation: str = "linear",
+ conf_activation: str = "expp1",
+ features: int = 256,
+ out_channels: List[int] = [256, 512, 1024, 1024],
+ intermediate_layer_idx: List[int] = [4, 11, 17, 23],
+ pos_embed: bool = True,
+ feature_only: bool = False,
+ down_ratio: int = 1,
+ ) -> None:
+ super(DPTHead, self).__init__()
+ self.patch_size = patch_size
+ self.activation = activation
+ self.conf_activation = conf_activation
+ self.pos_embed = pos_embed
+ self.feature_only = feature_only
+ self.down_ratio = down_ratio
+ self.intermediate_layer_idx = intermediate_layer_idx
+
+ self.norm = nn.LayerNorm(dim_in)
+
+ # Projection layers for each output channel from tokens.
+ self.projects = nn.ModuleList(
+ [
+ nn.Conv2d(
+ in_channels=dim_in,
+ out_channels=oc,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+ for oc in out_channels
+ ]
+ )
+
+ # Resize layers for upsampling feature maps.
+ self.resize_layers = nn.ModuleList(
+ [
+ nn.ConvTranspose2d(
+ in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=4, stride=4, padding=0
+ ),
+ nn.ConvTranspose2d(
+ in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0
+ ),
+ nn.Identity(),
+ nn.Conv2d(
+ in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=3, stride=2, padding=1
+ ),
+ ]
+ )
+
+ self.scratch = _make_scratch(
+ out_channels,
+ features,
+ expand=False,
+ )
+
+ # Attach additional modules to scratch.
+ self.scratch.stem_transpose = None
+ self.scratch.refinenet1 = _make_fusion_block(features)
+ self.scratch.refinenet2 = _make_fusion_block(features)
+ self.scratch.refinenet3 = _make_fusion_block(features)
+ self.scratch.refinenet4 = _make_fusion_block(features, has_residual=False)
+
+ head_features_1 = features
+ head_features_2 = 32
+
+ if feature_only:
+ self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1, kernel_size=3, stride=1, padding=1)
+ else:
+ self.scratch.output_conv1 = nn.Conv2d(
+ head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1
+ )
+ conv2_in_channels = head_features_1 // 2
+
+ # self.scratch.output_conv2 = nn.Sequential(
+ # nn.Conv2d(conv2_in_channels, head_features_2, kernel_size=3, stride=1, padding=1),
+ # nn.ReLU(inplace=True),
+ # nn.Conv2d(head_features_2, output_dim, kernel_size=1, stride=1, padding=0),
+ # )
+ # self.scratch.output_conv3是为了normal estimation而修改
+ # self.scratch.output_conv3 = nn.Sequential(
+ # # 第1层 3x3: 从 in_channels=1280 -> 512
+ # nn.Conv2d(conv2_in_channels, 512, kernel_size=3, stride=1, padding=1),
+ # nn.ReLU(inplace=True),
+
+ # # 第2层 3x3: 512 -> 512
+ # nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+ # nn.ReLU(inplace=True),
+
+ # # 第3层 3x3: 512 -> 256
+ # nn.Conv2d(512, head_features_2, kernel_size=3, stride=1, padding=1),
+ # nn.ReLU(inplace=True),
+
+ # # 最后一层 1x1: 256 -> 4
+ # nn.Conv2d(head_features_2, output_dim, kernel_size=1, stride=1, padding=0),
+ # )
+ def make_output_conv(
+ self,
+ in_channels: int,
+ output_dim: int = 4
+ ):
+
+ conv_mid1 = in_channels //2
+ conv_mid2 = conv_mid1
+ conv_mid3 = in_channels // 4
+
+ layers = nn.Sequential(
+ # 第1层 3x3: 从 in_channels -> conv_mid1
+ nn.Conv2d(in_channels, conv_mid1, kernel_size=3, stride=1, padding=1),
+ nn.ReLU(inplace=True),
+
+ # 第2层 3x3: conv_mid1 -> conv_mid2
+ nn.Conv2d(conv_mid1, conv_mid2, kernel_size=3, stride=1, padding=1),
+ nn.ReLU(inplace=True),
+
+ # 第3层 3x3: conv_mid2 -> conv_mid3
+ nn.Conv2d(conv_mid2, conv_mid3, kernel_size=3, stride=1, padding=1),
+ nn.ReLU(inplace=True),
+
+ # 第4层 1x1: conv_mid3 -> output_dim
+ nn.Conv2d(conv_mid3, output_dim, kernel_size=1, stride=1, padding=0),
+ ).cuda().to(torch.bfloat16)
+
+ return layers
+ def forward(
+ self,
+ aggregated_tokens_list: List[torch.Tensor],
+ images: torch.Tensor,
+ patch_start_idx: int,
+ frames_chunk_size: int = 10,
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Forward pass through the DPT head, supports processing by chunking frames.
+ Args:
+ aggregated_tokens_list (List[Tensor]): List of token tensors from different transformer layers.
+ images (Tensor): Input images with shape [B, S, 3, H, W], in range [0, 1].
+ patch_start_idx (int): Starting index for patch tokens in the token sequence.
+ Used to separate patch tokens from other tokens (e.g., camera or register tokens).
+ frames_chunk_size (int, optional): Number of frames to process in each chunk.
+ If None or larger than S, all frames are processed at once. Default: 8.
+
+ Returns:
+ Tensor or Tuple[Tensor, Tensor]:
+ - If feature_only=True: Feature maps with shape [B, S, C, H, W]
+ - Otherwise: Tuple of (predictions, confidence) both with shape [B, S, 1, H, W]
+ """
+ B, S, _, H, W = images.shape
+
+ # If frames_chunk_size is not specified or greater than S, process all frames at once
+ if frames_chunk_size is None or frames_chunk_size >= S:
+ return self._forward_impl(aggregated_tokens_list, images, patch_start_idx) #frames_chunk_size:一次最多处理几帧
+
+ # Otherwise, process frames in chunks to manage memory usage
+ assert frames_chunk_size > 0
+
+ # Process frames in batches
+ all_preds = []
+ all_conf = []
+
+ for frames_start_idx in range(0, S, frames_chunk_size):
+ frames_end_idx = min(frames_start_idx + frames_chunk_size, S)
+
+ # Process batch of frames
+ if self.feature_only:
+ chunk_output = self._forward_impl(
+ aggregated_tokens_list, images, patch_start_idx, frames_start_idx, frames_end_idx
+ )
+ all_preds.append(chunk_output)
+ else:
+ chunk_preds, chunk_conf = self._forward_impl(
+ aggregated_tokens_list, images, patch_start_idx, frames_start_idx, frames_end_idx
+ )
+ all_preds.append(chunk_preds)
+ all_conf.append(chunk_conf)
+
+ # Concatenate results along the sequence dimension
+ if self.feature_only:
+ return torch.cat(all_preds, dim=1)
+ else:
+ return torch.cat(all_preds, dim=1), torch.cat(all_conf, dim=1)
+
+ def _forward_impl(
+ self,
+ aggregated_tokens_list: List[torch.Tensor],
+ images: torch.Tensor,
+ patch_start_idx: int,
+ frames_start_idx: int = None,
+ frames_end_idx: int = None,
+ ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+ """
+ Implementation of the forward pass through the DPT head.
+
+ This method processes a specific chunk of frames from the sequence.
+
+ Args:
+ aggregated_tokens_list (List[Tensor]): List of token tensors from different transformer layers.
+ images (Tensor): Input images with shape [B, S, 3, H, W].
+ patch_start_idx (int): Starting index for patch tokens.
+ frames_start_idx (int, optional): Starting index for frames to process.
+ frames_end_idx (int, optional): Ending index for frames to process.
+
+ Returns:
+ Tensor or Tuple[Tensor, Tensor]: Feature maps or (predictions, confidence).
+ """
+ if frames_start_idx is not None and frames_end_idx is not None:
+ images = images[:, frames_start_idx:frames_end_idx].contiguous()
+
+ B, S, _, H, W = images.shape#就是图片
+
+ patch_h, patch_w = H // self.patch_size, W // self.patch_size
+
+ out = []
+ dpt_idx = 0
+ # aggregated_tokens_list: [24,B,3,H*W//14 + 4,C] 24是AABlock的个数
+ for layer_idx in self.intermediate_layer_idx:# 这段代码需要借鉴过来，相当于用transformer多层特征（相同尺寸）模仿conv的mutil-scale特征，有意思
+ x = aggregated_tokens_list[layer_idx][:, :, patch_start_idx:] #[B,3,H*W//14,C]
+ # 为什么只选[4,11,17,23]这几个layer_idx呢？
+ # Select frames if processing a chunk
+ if frames_start_idx is not None and frames_end_idx is not None:
+ x = x[:, frames_start_idx:frames_end_idx]
+
+ x = x.reshape(B * S, -1, x.shape[-1])
+
+ x = self.norm(x)
+
+ x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
+
+ x = self.projects[dpt_idx](x)
+ if self.pos_embed:
+ x = self._apply_pos_embed(x, W, H).to(torch.bfloat16)
+ x = self.resize_layers[dpt_idx](x)
+
+ out.append(x) #一个feature list
+ dpt_idx += 1
+
+ # Fuse features from multiple layers.
+ out = self.scratch_forward(out)
+ # Interpolate fused output to match target image resolution.
+ out = custom_interpolate(
+ out.float(),
+ (int(patch_h * self.patch_size / self.down_ratio), int(patch_w * self.patch_size / self.down_ratio)),
+ mode="bilinear",
+ align_corners=True,
+ ).to(torch.bfloat16) # feature upsample
+
+ if self.pos_embed:
+ out = self._apply_pos_embed(out, W, H).to(torch.bfloat16)
+
+ if self.feature_only:
+ return out.reshape(B, S, *out.shape[1:])
+
+ #out = self.scratch.output_conv2(out)
+ out = out.reshape(B, S, *out.shape[1:]).flatten(1,2) # B S C H W
+ in_channels = out.shape[1]
+ out = self.make_output_conv(in_channels, output_dim=4)(out) # B S 4 H W
+ preds, conf = activate_head(out, activation='linear', conf_activation=self.conf_activation)
+ # preds默认使用的是inv_log(x) = exp(x) - 1
+ preds = preds.permute(0, 3, 1, 2) # B C,H,W
+ preds = F.normalize(preds, dim=1, p=2, eps=1e-12) # B 1 H W
+ # conf[:, None, ...] # B 1 H W
+ return preds, conf.unsqueeze(1) # B 1 H W
+
+ def _apply_pos_embed(self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1) -> torch.Tensor:
+ """
+ Apply positional embedding to tensor x.
+ """
+ patch_w = x.shape[-1]
+ patch_h = x.shape[-2]
+ pos_embed = create_uv_grid(patch_w, patch_h, aspect_ratio=W / H, dtype=x.dtype, device=x.device)
+ pos_embed = position_grid_to_embed(pos_embed, x.shape[1])
+ pos_embed = pos_embed * ratio
+ pos_embed = pos_embed.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1)
+ return x + pos_embed
+
+ def scratch_forward(self, features: List[torch.Tensor]) -> torch.Tensor:
+ """
+ Forward pass through the fusion blocks.
+
+ Args:
+ features (List[Tensor]): List of feature maps from different layers.
+
+ Returns:
+ Tensor: Fused feature map.
+ """
+ layer_1, layer_2, layer_3, layer_4 = features
+
+ layer_1_rn = self.scratch.layer1_rn(layer_1)
+ layer_2_rn = self.scratch.layer2_rn(layer_2)
+ layer_3_rn = self.scratch.layer3_rn(layer_3)
+ layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+ out = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
+ del layer_4_rn, layer_4
+
+ out = self.scratch.refinenet3(out, layer_3_rn, size=layer_2_rn.shape[2:])
+ del layer_3_rn, layer_3
+
+ out = self.scratch.refinenet2(out, layer_2_rn, size=layer_1_rn.shape[2:])
+ del layer_2_rn, layer_2
+
+ out = self.scratch.refinenet1(out, layer_1_rn)
+ del layer_1_rn, layer_1
+
+ out = self.scratch.output_conv1(out)
+ return out
+
+
+################################################################################
+# Modules
+################################################################################
+
+
+def _make_fusion_block(features: int, size: int = None, has_residual: bool = True, groups: int = 1) -> nn.Module:
+ return FeatureFusionBlock(
+ features,
+ nn.ReLU(inplace=True),
+ deconv=False,
+ bn=False,
+ expand=False,
+ align_corners=True,
+ size=size,
+ has_residual=has_residual,
+ groups=groups,
+ )
+
+
+def _make_scratch(in_shape: List[int], out_shape: int, groups: int = 1, expand: bool = False) -> nn.Module:
+ scratch = nn.Module()
+ out_shape1 = out_shape
+ out_shape2 = out_shape
+ out_shape3 = out_shape
+ if len(in_shape) >= 4:
+ out_shape4 = out_shape
+
+ if expand:
+ out_shape1 = out_shape
+ out_shape2 = out_shape * 2
+ out_shape3 = out_shape * 4
+ if len(in_shape) >= 4:
+ out_shape4 = out_shape * 8
+
+ scratch.layer1_rn = nn.Conv2d(
+ in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ )
+ scratch.layer2_rn = nn.Conv2d(
+ in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ )
+ scratch.layer3_rn = nn.Conv2d(
+ in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ )
+ if len(in_shape) >= 4:
+ scratch.layer4_rn = nn.Conv2d(
+ in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ )
+ return scratch
+
+
+class ResidualConvUnit(nn.Module):
+ """Residual convolution module."""
+
+ def __init__(self, features, activation, bn, groups=1):
+ """Init.
+
+ Args:
+ features (int): number of features
+ """
+ super().__init__()
+
+ self.bn = bn
+ self.groups = groups
+ self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+ self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+
+ self.norm1 = None
+ self.norm2 = None
+
+ self.activation = activation
+ self.skip_add = nn.quantized.FloatFunctional()
+
+ def forward(self, x):
+ """Forward pass.
+
+ Args:
+ x (tensor): input
+
+ Returns:
+ tensor: output
+ """
+
+ out = self.activation(x)
+ out = self.conv1(out)
+ if self.norm1 is not None:
+ out = self.norm1(out)
+
+ out = self.activation(out)
+ out = self.conv2(out)
+ if self.norm2 is not None:
+ out = self.norm2(out)
+
+ return self.skip_add.add(out, x)
+
+
+class FeatureFusionBlock(nn.Module):
+ """Feature fusion block."""
+
+ def __init__(
+ self,
+ features,
+ activation,
+ deconv=False,
+ bn=False,
+ expand=False,
+ align_corners=True,
+ size=None,
+ has_residual=True,
+ groups=1,
+ ):
+ """Init.
+
+ Args:
+ features (int): number of features
+ """
+ super(FeatureFusionBlock, self).__init__()
+
+ self.deconv = deconv
+ self.align_corners = align_corners
+ self.groups = groups
+ self.expand = expand
+ out_features = features
+ if self.expand == True:
+ out_features = features // 2
+
+ self.out_conv = nn.Conv2d(
+ features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=self.groups
+ )
+
+ if has_residual:
+ self.resConfUnit1 = ResidualConvUnit(features, activation, bn, groups=self.groups)
+
+ self.has_residual = has_residual
+ self.resConfUnit2 = ResidualConvUnit(features, activation, bn, groups=self.groups)
+
+ self.skip_add = nn.quantized.FloatFunctional()
+ self.size = size
+
+ def forward(self, *xs, size=None):
+ """Forward pass.
+
+ Returns:
+ tensor: output
+ """
+ output = xs[0]
+
+ if self.has_residual:
+ res = self.resConfUnit1(xs[1])
+ output = self.skip_add.add(output, res)
+
+ output = self.resConfUnit2(output)
+
+ if (size is None) and (self.size is None):
+ modifier = {"scale_factor": 2}
+ elif size is None:
+ modifier = {"size": self.size}
+ else:
+ modifier = {"size": size}
+
+ output = custom_interpolate(output.float(), **modifier, mode="bilinear", align_corners=self.align_corners).to(torch.bfloat16)
+ output = self.out_conv(output)
+
+ return output
+
+
+def custom_interpolate(
+ x: torch.Tensor,
+ size: Tuple[int, int] = None,
+ scale_factor: float = None,
+ mode: str = "bilinear",
+ align_corners: bool = True,
+) -> torch.Tensor:
+ """
+ Custom interpolate to avoid INT_MAX issues in nn.functional.interpolate.
+ """
+ if size is None:
+ size = (int(x.shape[-2] * scale_factor), int(x.shape[-1] * scale_factor))
+
+ INT_MAX = 1610612736
+
+ input_elements = size[0] * size[1] * x.shape[0] * x.shape[1]
+
+ if input_elements > INT_MAX:
+ chunks = torch.chunk(x, chunks=(input_elements // INT_MAX) + 1, dim=0)
+ interpolated_chunks = [
+ nn.functional.interpolate(chunk, size=size, mode=mode, align_corners=align_corners) for chunk in chunks
+ ]
+ x = torch.cat(interpolated_chunks, dim=0)
+ return x.contiguous()
+ else:
+ return nn.functional.interpolate(x, size=size, mode=mode, align_corners=align_corners)
diff --git a/src/models/aggregator/heads/head_act.py b/src/models/aggregator/heads/head_act.py
new file mode 100644
index 0000000000000000000000000000000000000000..9a8651bd7cf2214c60629cc680078ac50f41d675
--- /dev/null
+++ b/src/models/aggregator/heads/head_act.py
@@ -0,0 +1,126 @@
+# 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.
+
+
+import torch
+import torch.nn.functional as F
+
+
+def activate_pose(pred_pose_enc, trans_act="linear", quat_act="linear", fl_act="linear"):
+ """
+ Activate pose parameters with specified activation functions.
+
+ Args:
+ pred_pose_enc: Tensor containing encoded pose parameters [translation, quaternion, focal length]
+ trans_act: Activation type for translation component
+ quat_act: Activation type for quaternion component
+ fl_act: Activation type for focal length component
+
+ Returns:
+ Activated pose parameters tensor
+ """
+ T = pred_pose_enc[..., :3]
+ quat = pred_pose_enc[..., 3:7]
+ fl = pred_pose_enc[..., 7:] # or fov
+
+ T = base_pose_act(T, trans_act)
+ quat = base_pose_act(quat, quat_act)
+ fl = base_pose_act(fl, fl_act) # or fov
+
+ pred_pose_enc = torch.cat([T, quat, fl], dim=-1)
+
+ return pred_pose_enc
+
+
+def base_pose_act(pose_enc, act_type="linear"):
+ """
+ Apply basic activation function to pose parameters.
+
+ Args:
+ pose_enc: Tensor containing encoded pose parameters
+ act_type: Activation type ("linear", "inv_log", "exp", "relu")
+
+ Returns:
+ Activated pose parameters
+ """
+ if act_type == "linear":
+ return pose_enc
+ elif act_type == "inv_log":
+ return inverse_log_transform(pose_enc)
+ elif act_type == "exp":
+ return torch.exp(pose_enc)
+ elif act_type == "relu":
+ return F.relu(pose_enc)
+ else:
+ raise ValueError(f"Unknown act_type: {act_type}")
+
+
+def activate_head(out, activation="norm_exp", conf_activation="expp1"):
+ """
+ Process network output to extract 3D points and confidence values.
+
+ Args:
+ out: Network output tensor (B, C, H, W)
+ activation: Activation type for 3D points
+ conf_activation: Activation type for confidence values
+
+ Returns:
+ Tuple of (3D points tensor, confidence tensor)
+ """
+ # Move channels from last dim to the 4th dimension => (B, H, W, C)
+ fmap = out.permute(0, 2, 3, 1) # B,H,W,C expected
+
+ # Split into xyz (first C-1 channels) and confidence (last channel)
+ xyz = fmap[:, :, :, :-1]
+ conf = fmap[:, :, :, -1]
+
+ if activation == "norm_exp":
+ d = xyz.norm(dim=-1, keepdim=True).clamp(min=1e-8)
+ xyz_normed = xyz / d
+ pts3d = xyz_normed * torch.expm1(d)
+ elif activation == "norm":
+ pts3d = xyz / xyz.norm(dim=-1, keepdim=True)
+ elif activation == "exp":
+ pts3d = torch.exp(xyz)
+ elif activation == "relu":
+ pts3d = F.relu(xyz)
+ elif activation == "inv_log":
+ pts3d = inverse_log_transform(xyz)
+ elif activation == "xy_inv_log":
+ xy, z = xyz.split([2, 1], dim=-1)
+ z = inverse_log_transform(z)
+ pts3d = torch.cat([xy * z, z], dim=-1)
+ elif activation == "sigmoid":
+ pts3d = torch.sigmoid(xyz)
+ elif activation == "linear":
+ pts3d = xyz
+
+ else:
+ raise ValueError(f"Unknown activation: {activation}")
+
+ if conf_activation == "expp1":
+ conf_out = 1 + conf.exp()
+ elif conf_activation == "expp0":
+ conf_out = conf.exp()
+ elif conf_activation == "sigmoid":
+ conf_out = torch.sigmoid(conf)
+ else:
+ raise ValueError(f"Unknown conf_activation: {conf_activation}")
+
+ return pts3d, conf_out
+
+
+def inverse_log_transform(y):
+ """
+ Apply inverse log transform: sign(y) * (exp(|y|) - 1)
+
+ Args:
+ y: Input tensor
+
+ Returns:
+ Transformed tensor
+ """
+ return torch.sign(y) * (torch.expm1(torch.abs(y)))
diff --git a/src/models/aggregator/heads/track_head.py b/src/models/aggregator/heads/track_head.py
new file mode 100644
index 0000000000000000000000000000000000000000..9ec7199bd185060989c236997f93b93f4fc77825
--- /dev/null
+++ b/src/models/aggregator/heads/track_head.py
@@ -0,0 +1,108 @@
+# 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.
+
+import torch.nn as nn
+from .dpt_head import DPTHead
+from .track_modules.base_track_predictor import BaseTrackerPredictor
+
+
+class TrackHead(nn.Module):
+ """
+ Track head that uses DPT head to process tokens and BaseTrackerPredictor for tracking.
+ The tracking is performed iteratively, refining predictions over multiple iterations.
+ """
+
+ def __init__(
+ self,
+ dim_in,
+ patch_size=14,
+ features=128,
+ iters=4,
+ predict_conf=True,
+ stride=2,
+ corr_levels=7,
+ corr_radius=4,
+ hidden_size=384,
+ ):
+ """
+ Initialize the TrackHead module.
+
+ Args:
+ dim_in (int): Input dimension of tokens from the backbone.
+ patch_size (int): Size of image patches used in the vision transformer.
+ features (int): Number of feature channels in the feature extractor output.
+ iters (int): Number of refinement iterations for tracking predictions.
+ predict_conf (bool): Whether to predict confidence scores for tracked points.
+ stride (int): Stride value for the tracker predictor.
+ corr_levels (int): Number of correlation pyramid levels
+ corr_radius (int): Radius for correlation computation, controlling the search area.
+ hidden_size (int): Size of hidden layers in the tracker network.
+ """
+ super().__init__()
+
+ self.patch_size = patch_size
+
+ # Feature extractor based on DPT architecture
+ # Processes tokens into feature maps for tracking
+ self.feature_extractor = DPTHead(
+ dim_in=dim_in,
+ patch_size=patch_size,
+ features=features,
+ feature_only=True, # Only output features, no activation
+ down_ratio=2, # Reduces spatial dimensions by factor of 2
+ pos_embed=False,
+ )
+
+ # Tracker module that predicts point trajectories
+ # Takes feature maps and predicts coordinates and visibility
+ self.tracker = BaseTrackerPredictor(
+ latent_dim=features, # Match the output_dim of feature extractor
+ predict_conf=predict_conf,
+ stride=stride,
+ corr_levels=corr_levels,
+ corr_radius=corr_radius,
+ hidden_size=hidden_size,
+ )
+
+ self.iters = iters
+
+ def forward(self, aggregated_tokens_list, images, patch_start_idx, query_points=None, iters=None):
+ """
+ Forward pass of the TrackHead.
+
+ Args:
+ aggregated_tokens_list (list): List of aggregated tokens from the backbone.
+ images (torch.Tensor): Input images of shape (B, S, C, H, W) where:
+ B = batch size, S = sequence length.
+ patch_start_idx (int): Starting index for patch tokens.
+ query_points (torch.Tensor, optional): Initial query points to track.
+ If None, points are initialized by the tracker.
+ iters (int, optional): Number of refinement iterations. If None, uses self.iters.
+
+ Returns:
+ tuple:
+ - coord_preds (torch.Tensor): Predicted coordinates for tracked points.
+ - vis_scores (torch.Tensor): Visibility scores for tracked points.
+ - conf_scores (torch.Tensor): Confidence scores for tracked points (if predict_conf=True).
+ """
+ B, S, _, H, W = images.shape
+
+ # Extract features from tokens
+ # feature_maps has shape (B, S, C, H//2, W//2) due to down_ratio=2
+ feature_maps = self.feature_extractor(aggregated_tokens_list, images, patch_start_idx)
+
+ # Use default iterations if not specified
+ if iters is None:
+ iters = self.iters
+
+ # Perform tracking using the extracted features
+ coord_preds, vis_scores, conf_scores = self.tracker(
+ query_points=query_points,
+ fmaps=feature_maps,
+ iters=iters,
+ )
+
+ return coord_preds, vis_scores, conf_scores
diff --git a/src/models/aggregator/heads/track_modules/__init__.py b/src/models/aggregator/heads/track_modules/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0952fcc3f57e34b3747962e9ebd6fc57aeea63fa
--- /dev/null
+++ b/src/models/aggregator/heads/track_modules/__init__.py
@@ -0,0 +1,5 @@
+# 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.
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/__init__.cpython-310.pyc b/src/models/aggregator/heads/track_modules/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6f5e515ae9e321249c39dd136afa8b01eacb9801
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/__init__.cpython-38.pyc b/src/models/aggregator/heads/track_modules/__pycache__/__init__.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f15cb3d22c0764ab1cc52da68187e8fbfc946d3a
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/__init__.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/base_track_predictor.cpython-310.pyc b/src/models/aggregator/heads/track_modules/__pycache__/base_track_predictor.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6ae954de6ef79378cd2f37035dc5b9813c5ae7b0
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/base_track_predictor.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/base_track_predictor.cpython-38.pyc b/src/models/aggregator/heads/track_modules/__pycache__/base_track_predictor.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..4a17f4ee16b9bd2b7c872d6970b08c7fffc1f59a
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/base_track_predictor.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/blocks.cpython-310.pyc b/src/models/aggregator/heads/track_modules/__pycache__/blocks.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5eb4e59a075bd5d7043666b0cd50dc3ec67858e3
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/blocks.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/blocks.cpython-38.pyc b/src/models/aggregator/heads/track_modules/__pycache__/blocks.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cd883b0f1454656a06f94402062bc845e70a1136
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/blocks.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/modules.cpython-310.pyc b/src/models/aggregator/heads/track_modules/__pycache__/modules.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..de45769bc99361637d73d6ad6869c36283f36061
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/modules.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/modules.cpython-38.pyc b/src/models/aggregator/heads/track_modules/__pycache__/modules.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..195d417765ec1362574bd113c619f3b39db182d9
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/modules.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/utils.cpython-310.pyc b/src/models/aggregator/heads/track_modules/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..27be613add713a68263532cc5b40e99f9d15bdee
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/utils.cpython-310.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/__pycache__/utils.cpython-38.pyc b/src/models/aggregator/heads/track_modules/__pycache__/utils.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..30ed44ae9e4109eb83f9e6034675027ba2c0ec05
Binary files /dev/null and b/src/models/aggregator/heads/track_modules/__pycache__/utils.cpython-38.pyc differ
diff --git a/src/models/aggregator/heads/track_modules/base_track_predictor.py b/src/models/aggregator/heads/track_modules/base_track_predictor.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ce8ec4b66fff236e015d1bcaf85c8237a52be7a
--- /dev/null
+++ b/src/models/aggregator/heads/track_modules/base_track_predictor.py
@@ -0,0 +1,209 @@
+# 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.
+
+import torch
+import torch.nn as nn
+from einops import rearrange, repeat
+
+
+from .blocks import EfficientUpdateFormer, CorrBlock
+from .utils import sample_features4d, get_2d_embedding, get_2d_sincos_pos_embed
+from .modules import Mlp
+
+
+class BaseTrackerPredictor(nn.Module):
+ def __init__(
+ self,
+ stride=1,
+ corr_levels=5,
+ corr_radius=4,
+ latent_dim=128,
+ hidden_size=384,
+ use_spaceatt=True,
+ depth=6,
+ max_scale=518,
+ predict_conf=True,
+ ):
+ super(BaseTrackerPredictor, self).__init__()
+ """
+ The base template to create a track predictor
+
+ Modified from https://github.com/facebookresearch/co-tracker/
+ and https://github.com/facebookresearch/vggsfm
+ """
+
+ self.stride = stride
+ self.latent_dim = latent_dim
+ self.corr_levels = corr_levels
+ self.corr_radius = corr_radius
+ self.hidden_size = hidden_size
+ self.max_scale = max_scale
+ self.predict_conf = predict_conf
+
+ self.flows_emb_dim = latent_dim // 2
+
+ self.corr_mlp = Mlp(
+ in_features=self.corr_levels * (self.corr_radius * 2 + 1) ** 2,
+ hidden_features=self.hidden_size,
+ out_features=self.latent_dim,
+ )
+
+ self.transformer_dim = self.latent_dim + self.latent_dim + self.latent_dim + 4
+
+ self.query_ref_token = nn.Parameter(torch.randn(1, 2, self.transformer_dim))
+
+ space_depth = depth if use_spaceatt else 0
+ time_depth = depth
+
+ self.updateformer = EfficientUpdateFormer(
+ space_depth=space_depth,
+ time_depth=time_depth,
+ input_dim=self.transformer_dim,
+ hidden_size=self.hidden_size,
+ output_dim=self.latent_dim + 2,
+ mlp_ratio=4.0,
+ add_space_attn=use_spaceatt,
+ )
+
+ self.fmap_norm = nn.LayerNorm(self.latent_dim)
+ self.ffeat_norm = nn.GroupNorm(1, self.latent_dim)
+
+ # A linear layer to update track feats at each iteration
+ self.ffeat_updater = nn.Sequential(nn.Linear(self.latent_dim, self.latent_dim), nn.GELU())
+
+ self.vis_predictor = nn.Sequential(nn.Linear(self.latent_dim, 1))
+
+ if predict_conf:
+ self.conf_predictor = nn.Sequential(nn.Linear(self.latent_dim, 1))
+
+ def forward(self, query_points, fmaps=None, iters=6, return_feat=False, down_ratio=1, apply_sigmoid=True):
+ """
+ query_points: B x N x 2, the number of batches, tracks, and xy
+ fmaps: B x S x C x HH x WW, the number of batches, frames, and feature dimension.
+ note HH and WW is the size of feature maps instead of original images
+ """
+ B, N, D = query_points.shape
+ B, S, C, HH, WW = fmaps.shape
+
+ assert D == 2, "Input points must be 2D coordinates"
+
+ # apply a layernorm to fmaps here
+ fmaps = self.fmap_norm(fmaps.permute(0, 1, 3, 4, 2))
+ fmaps = fmaps.permute(0, 1, 4, 2, 3)
+
+ # Scale the input query_points because we may downsample the images
+ # by down_ratio or self.stride
+ # e.g., if a 3x1024x1024 image is processed to a 128x256x256 feature map
+ # its query_points should be query_points/4
+ if down_ratio > 1:
+ query_points = query_points / float(down_ratio)
+
+ query_points = query_points / float(self.stride)
+
+ # Init with coords as the query points
+ # It means the search will start from the position of query points at the reference frames
+ coords = query_points.clone().reshape(B, 1, N, 2).repeat(1, S, 1, 1)
+
+ # Sample/extract the features of the query points in the query frame
+ query_track_feat = sample_features4d(fmaps[:, 0], coords[:, 0])
+
+ # init track feats by query feats
+ track_feats = query_track_feat.unsqueeze(1).repeat(1, S, 1, 1) # B, S, N, C
+ # back up the init coords
+ coords_backup = coords.clone()
+
+ fcorr_fn = CorrBlock(fmaps, num_levels=self.corr_levels, radius=self.corr_radius)
+
+ coord_preds = []
+
+ # Iterative Refinement
+ for _ in range(iters):
+ # Detach the gradients from the last iteration
+ # (in my experience, not very important for performance)
+ coords = coords.detach()
+
+ fcorrs = fcorr_fn.corr_sample(track_feats, coords)
+
+ corr_dim = fcorrs.shape[3]
+ fcorrs_ = fcorrs.permute(0, 2, 1, 3).reshape(B * N, S, corr_dim)
+ fcorrs_ = self.corr_mlp(fcorrs_)
+
+ # Movement of current coords relative to query points
+ flows = (coords - coords[:, 0:1]).permute(0, 2, 1, 3).reshape(B * N, S, 2)
+
+ flows_emb = get_2d_embedding(flows, self.flows_emb_dim, cat_coords=False)
+
+ # (In my trials, it is also okay to just add the flows_emb instead of concat)
+ flows_emb = torch.cat([flows_emb, flows / self.max_scale, flows / self.max_scale], dim=-1)
+
+ track_feats_ = track_feats.permute(0, 2, 1, 3).reshape(B * N, S, self.latent_dim)
+
+ # Concatenate them as the input for the transformers
+ transformer_input = torch.cat([flows_emb, fcorrs_, track_feats_], dim=2)
+
+ # 2D positional embed
+ # TODO: this can be much simplified
+ pos_embed = get_2d_sincos_pos_embed(self.transformer_dim, grid_size=(HH, WW)).to(query_points.device)
+ sampled_pos_emb = sample_features4d(pos_embed.expand(B, -1, -1, -1), coords[:, 0])
+
+ sampled_pos_emb = rearrange(sampled_pos_emb, "b n c -> (b n) c").unsqueeze(1)
+
+ x = transformer_input + sampled_pos_emb
+
+ # Add the query ref token to the track feats
+ query_ref_token = torch.cat(
+ [self.query_ref_token[:, 0:1], self.query_ref_token[:, 1:2].expand(-1, S - 1, -1)], dim=1
+ )
+ x = x + query_ref_token.to(x.device).to(x.dtype)
+
+ # B, N, S, C
+ x = rearrange(x, "(b n) s d -> b n s d", b=B)
+
+ # Compute the delta coordinates and delta track features
+ delta, _ = self.updateformer(x)
+
+ # BN, S, C
+ delta = rearrange(delta, " b n s d -> (b n) s d", b=B)
+ delta_coords_ = delta[:, :, :2]
+ delta_feats_ = delta[:, :, 2:]
+
+ track_feats_ = track_feats_.reshape(B * N * S, self.latent_dim)
+ delta_feats_ = delta_feats_.reshape(B * N * S, self.latent_dim)
+
+ # Update the track features
+ track_feats_ = self.ffeat_updater(self.ffeat_norm(delta_feats_)) + track_feats_
+
+ track_feats = track_feats_.reshape(B, N, S, self.latent_dim).permute(0, 2, 1, 3) # BxSxNxC
+
+ # B x S x N x 2
+ coords = coords + delta_coords_.reshape(B, N, S, 2).permute(0, 2, 1, 3)
+
+ # Force coord0 as query
+ # because we assume the query points should not be changed
+ coords[:, 0] = coords_backup[:, 0]
+
+ # The predicted tracks are in the original image scale
+ if down_ratio > 1:
+ coord_preds.append(coords * self.stride * down_ratio)
+ else:
+ coord_preds.append(coords * self.stride)
+
+ # B, S, N
+ vis_e = self.vis_predictor(track_feats.reshape(B * S * N, self.latent_dim)).reshape(B, S, N)
+ if apply_sigmoid:
+ vis_e = torch.sigmoid(vis_e)
+
+ if self.predict_conf:
+ conf_e = self.conf_predictor(track_feats.reshape(B * S * N, self.latent_dim)).reshape(B, S, N)
+ if apply_sigmoid:
+ conf_e = torch.sigmoid(conf_e)
+ else:
+ conf_e = None
+
+ if return_feat:
+ return coord_preds, vis_e, track_feats, query_track_feat, conf_e
+ else:
+ return coord_preds, vis_e, conf_e
diff --git a/src/models/aggregator/heads/track_modules/blocks.py b/src/models/aggregator/heads/track_modules/blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e7763f4fd8f515662421db192594380dbb574e5
--- /dev/null
+++ b/src/models/aggregator/heads/track_modules/blocks.py
@@ -0,0 +1,246 @@
+# 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 from https://github.com/facebookresearch/co-tracker/
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import bilinear_sampler
+from .modules import Mlp, AttnBlock, CrossAttnBlock, ResidualBlock
+
+
+class EfficientUpdateFormer(nn.Module):
+ """
+ Transformer model that updates track estimates.
+ """
+
+ def __init__(
+ self,
+ space_depth=6,
+ time_depth=6,
+ input_dim=320,
+ hidden_size=384,
+ num_heads=8,
+ output_dim=130,
+ mlp_ratio=4.0,
+ add_space_attn=True,
+ num_virtual_tracks=64,
+ ):
+ super().__init__()
+
+ self.out_channels = 2
+ self.num_heads = num_heads
+ self.hidden_size = hidden_size
+ self.add_space_attn = add_space_attn
+
+ # Add input LayerNorm before linear projection
+ self.input_norm = nn.LayerNorm(input_dim)
+ self.input_transform = torch.nn.Linear(input_dim, hidden_size, bias=True)
+
+ # Add output LayerNorm before final projection
+ self.output_norm = nn.LayerNorm(hidden_size)
+ self.flow_head = torch.nn.Linear(hidden_size, output_dim, bias=True)
+ self.num_virtual_tracks = num_virtual_tracks
+
+ if self.add_space_attn:
+ self.virual_tracks = nn.Parameter(torch.randn(1, num_virtual_tracks, 1, hidden_size))
+ else:
+ self.virual_tracks = None
+
+ self.time_blocks = nn.ModuleList(
+ [
+ AttnBlock(
+ hidden_size,
+ num_heads,
+ mlp_ratio=mlp_ratio,
+ attn_class=nn.MultiheadAttention,
+ )
+ for _ in range(time_depth)
+ ]
+ )
+
+ if add_space_attn:
+ self.space_virtual_blocks = nn.ModuleList(
+ [
+ AttnBlock(
+ hidden_size,
+ num_heads,
+ mlp_ratio=mlp_ratio,
+ attn_class=nn.MultiheadAttention,
+ )
+ for _ in range(space_depth)
+ ]
+ )
+ self.space_point2virtual_blocks = nn.ModuleList(
+ [CrossAttnBlock(hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(space_depth)]
+ )
+ self.space_virtual2point_blocks = nn.ModuleList(
+ [CrossAttnBlock(hidden_size, hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(space_depth)]
+ )
+ assert len(self.time_blocks) >= len(self.space_virtual2point_blocks)
+ self.initialize_weights()
+
+ def initialize_weights(self):
+ def _basic_init(module):
+ if isinstance(module, nn.Linear):
+ torch.nn.init.xavier_uniform_(module.weight)
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+ torch.nn.init.trunc_normal_(self.flow_head.weight, std=0.001)
+
+ self.apply(_basic_init)
+
+ def forward(self, input_tensor, mask=None):
+ # Apply input LayerNorm
+ input_tensor = self.input_norm(input_tensor)
+ tokens = self.input_transform(input_tensor)
+
+ init_tokens = tokens
+
+ B, _, T, _ = tokens.shape
+
+ if self.add_space_attn:
+ virtual_tokens = self.virual_tracks.repeat(B, 1, T, 1)
+ tokens = torch.cat([tokens, virtual_tokens], dim=1)
+
+ _, N, _, _ = tokens.shape
+
+ j = 0
+ for i in range(len(self.time_blocks)):
+ time_tokens = tokens.contiguous().view(B * N, T, -1) # B N T C -> (B N) T C
+
+ time_tokens = self.time_blocks[i](time_tokens)
+
+ tokens = time_tokens.view(B, N, T, -1) # (B N) T C -> B N T C
+ if self.add_space_attn and (i % (len(self.time_blocks) // len(self.space_virtual_blocks)) == 0):
+ space_tokens = tokens.permute(0, 2, 1, 3).contiguous().view(B * T, N, -1) # B N T C -> (B T) N C
+ point_tokens = space_tokens[:, : N - self.num_virtual_tracks]
+ virtual_tokens = space_tokens[:, N - self.num_virtual_tracks :]
+
+ virtual_tokens = self.space_virtual2point_blocks[j](virtual_tokens, point_tokens, mask=mask)
+ virtual_tokens = self.space_virtual_blocks[j](virtual_tokens)
+ point_tokens = self.space_point2virtual_blocks[j](point_tokens, virtual_tokens, mask=mask)
+
+ space_tokens = torch.cat([point_tokens, virtual_tokens], dim=1)
+ tokens = space_tokens.view(B, T, N, -1).permute(0, 2, 1, 3) # (B T) N C -> B N T C
+ j += 1
+
+ if self.add_space_attn:
+ tokens = tokens[:, : N - self.num_virtual_tracks]
+
+ tokens = tokens + init_tokens
+
+ # Apply output LayerNorm before final projection
+ tokens = self.output_norm(tokens)
+ flow = self.flow_head(tokens)
+
+ return flow, None
+
+
+class CorrBlock:
+ def __init__(self, fmaps, num_levels=4, radius=4, multiple_track_feats=False, padding_mode="zeros"):
+ """
+ Build a pyramid of feature maps from the input.
+
+ fmaps: Tensor (B, S, C, H, W)
+ num_levels: number of pyramid levels (each downsampled by factor 2)
+ radius: search radius for sampling correlation
+ multiple_track_feats: if True, split the target features per pyramid level
+ padding_mode: passed to grid_sample / bilinear_sampler
+ """
+ B, S, C, H, W = fmaps.shape
+ self.S, self.C, self.H, self.W = S, C, H, W
+ self.num_levels = num_levels
+ self.radius = radius
+ self.padding_mode = padding_mode
+ self.multiple_track_feats = multiple_track_feats
+
+ # Build pyramid: each level is half the spatial resolution of the previous
+ self.fmaps_pyramid = [fmaps] # level 0 is full resolution
+ current_fmaps = fmaps
+ for i in range(num_levels - 1):
+ B, S, C, H, W = current_fmaps.shape
+ # Merge batch & sequence dimensions
+ current_fmaps = current_fmaps.reshape(B * S, C, H, W)
+ # Avg pool down by factor 2
+ current_fmaps = F.avg_pool2d(current_fmaps, kernel_size=2, stride=2)
+ _, _, H_new, W_new = current_fmaps.shape
+ current_fmaps = current_fmaps.reshape(B, S, C, H_new, W_new)
+ self.fmaps_pyramid.append(current_fmaps)
+
+ # Precompute a delta grid (of shape (2r+1, 2r+1, 2)) for sampling.
+ # This grid is added to the (scaled) coordinate centroids.
+ r = self.radius
+ dx = torch.linspace(-r, r, 2 * r + 1, device=fmaps.device, dtype=fmaps.dtype)
+ dy = torch.linspace(-r, r, 2 * r + 1, device=fmaps.device, dtype=fmaps.dtype)
+ # delta: for every (dy,dx) displacement (i.e. Δx, Δy)
+ self.delta = torch.stack(torch.meshgrid(dy, dx, indexing="ij"), dim=-1) # shape: (2r+1, 2r+1, 2)
+
+ def corr_sample(self, targets, coords):
+ """
+ Instead of storing the entire correlation pyramid, we compute each level's correlation
+ volume, sample it immediately, then discard it. This saves GPU memory.
+
+ Args:
+ targets: Tensor (B, S, N, C) — features for the current targets.
+ coords: Tensor (B, S, N, 2) — coordinates at full resolution.
+
+ Returns:
+ Tensor (B, S, N, L) where L = num_levels * (2*radius+1)**2 (concatenated sampled correlations)
+ """
+ B, S, N, C = targets.shape
+
+ # If you have multiple track features, split them per level.
+ if self.multiple_track_feats:
+ targets_split = torch.split(targets, C // self.num_levels, dim=-1)
+
+ out_pyramid = []
+ for i, fmaps in enumerate(self.fmaps_pyramid):
+ # Get current spatial resolution H, W for this pyramid level.
+ B, S, C, H, W = fmaps.shape
+ # Reshape feature maps for correlation computation:
+ # fmap2s: (B, S, C, H*W)
+ fmap2s = fmaps.view(B, S, C, H * W)
+ # Choose appropriate target features.
+ fmap1 = targets_split[i] if self.multiple_track_feats else targets # shape: (B, S, N, C)
+
+ # Compute correlation directly
+ corrs = compute_corr_level(fmap1, fmap2s, C)
+ corrs = corrs.view(B, S, N, H, W)
+
+ # Prepare sampling grid:
+ # Scale down the coordinates for the current level.
+ centroid_lvl = coords.reshape(B * S * N, 1, 1, 2) / (2**i)
+ # Make sure our precomputed delta grid is on the same device/dtype.
+ delta_lvl = self.delta.to(coords.device).to(coords.dtype)
+ # Now the grid for grid_sample is:
+ # coords_lvl = centroid_lvl + delta_lvl (broadcasted over grid)
+ coords_lvl = centroid_lvl + delta_lvl.view(1, 2 * self.radius + 1, 2 * self.radius + 1, 2)
+
+ # Sample from the correlation volume using bilinear interpolation.
+ # We reshape corrs to (B * S * N, 1, H, W) so grid_sample acts over each target.
+ corrs_sampled = bilinear_sampler(
+ corrs.reshape(B * S * N, 1, H, W), coords_lvl, padding_mode=self.padding_mode
+ )
+ # The sampled output is (B * S * N, 1, 2r+1, 2r+1). Flatten the last two dims.
+ corrs_sampled = corrs_sampled.view(B, S, N, -1) # Now shape: (B, S, N, (2r+1)^2)
+ out_pyramid.append(corrs_sampled)
+
+ # Concatenate all levels along the last dimension.
+ out = torch.cat(out_pyramid, dim=-1).contiguous()
+ return out
+
+
+def compute_corr_level(fmap1, fmap2s, C):
+ # fmap1: (B, S, N, C)
+ # fmap2s: (B, S, C, H*W)
+ corrs = torch.matmul(fmap1, fmap2s) # (B, S, N, H*W)
+ corrs = corrs.view(fmap1.shape[0], fmap1.shape[1], fmap1.shape[2], -1) # (B, S, N, H*W)
+ return corrs / math.sqrt(C)
diff --git a/src/models/aggregator/heads/track_modules/modules.py b/src/models/aggregator/heads/track_modules/modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..4b090ddc4a9db01c8dd3564f9053e1ca9cdde93a
--- /dev/null
+++ b/src/models/aggregator/heads/track_modules/modules.py
@@ -0,0 +1,218 @@
+# 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.
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from functools import partial
+from typing import Callable
+import collections
+from torch import Tensor
+from itertools import repeat
+
+
+# From PyTorch internals
+def _ntuple(n):
+ def parse(x):
+ if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
+ return tuple(x)
+ return tuple(repeat(x, n))
+
+ return parse
+
+
+def exists(val):
+ return val is not None
+
+
+def default(val, d):
+ return val if exists(val) else d
+
+
+to_2tuple = _ntuple(2)
+
+
+class ResidualBlock(nn.Module):
+ """
+ ResidualBlock: construct a block of two conv layers with residual connections
+ """
+
+ def __init__(self, in_planes, planes, norm_fn="group", stride=1, kernel_size=3):
+ super(ResidualBlock, self).__init__()
+
+ self.conv1 = nn.Conv2d(
+ in_planes,
+ planes,
+ kernel_size=kernel_size,
+ padding=1,
+ stride=stride,
+ padding_mode="zeros",
+ )
+ self.conv2 = nn.Conv2d(
+ planes,
+ planes,
+ kernel_size=kernel_size,
+ padding=1,
+ padding_mode="zeros",
+ )
+ self.relu = nn.ReLU(inplace=True)
+
+ num_groups = planes // 8
+
+ if norm_fn == "group":
+ self.norm1 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+ self.norm2 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+ if not stride == 1:
+ self.norm3 = nn.GroupNorm(num_groups=num_groups, num_channels=planes)
+
+ elif norm_fn == "batch":
+ self.norm1 = nn.BatchNorm2d(planes)
+ self.norm2 = nn.BatchNorm2d(planes)
+ if not stride == 1:
+ self.norm3 = nn.BatchNorm2d(planes)
+
+ elif norm_fn == "instance":
+ self.norm1 = nn.InstanceNorm2d(planes)
+ self.norm2 = nn.InstanceNorm2d(planes)
+ if not stride == 1:
+ self.norm3 = nn.InstanceNorm2d(planes)
+
+ elif norm_fn == "none":
+ self.norm1 = nn.Sequential()
+ self.norm2 = nn.Sequential()
+ if not stride == 1:
+ self.norm3 = nn.Sequential()
+ else:
+ raise NotImplementedError
+
+ if stride == 1:
+ self.downsample = None
+ else:
+ self.downsample = nn.Sequential(
+ nn.Conv2d(in_planes, planes, kernel_size=1, stride=stride),
+ self.norm3,
+ )
+
+ def forward(self, x):
+ y = x
+ y = self.relu(self.norm1(self.conv1(y)))
+ y = self.relu(self.norm2(self.conv2(y)))
+
+ if self.downsample is not None:
+ x = self.downsample(x)
+
+ return self.relu(x + y)
+
+
+class Mlp(nn.Module):
+ """MLP as used in Vision Transformer, MLP-Mixer and related networks"""
+
+ def __init__(
+ self,
+ in_features,
+ hidden_features=None,
+ out_features=None,
+ act_layer=nn.GELU,
+ norm_layer=None,
+ bias=True,
+ drop=0.0,
+ use_conv=False,
+ ):
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ bias = to_2tuple(bias)
+ drop_probs = to_2tuple(drop)
+ linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
+
+ self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0])
+ self.act = act_layer()
+ self.drop1 = nn.Dropout(drop_probs[0])
+ self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1])
+ self.drop2 = nn.Dropout(drop_probs[1])
+
+ def forward(self, x):
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.drop1(x)
+ x = self.fc2(x)
+ x = self.drop2(x)
+ return x
+
+
+class AttnBlock(nn.Module):
+ def __init__(
+ self,
+ hidden_size,
+ num_heads,
+ attn_class: Callable[..., nn.Module] = nn.MultiheadAttention,
+ mlp_ratio=4.0,
+ **block_kwargs
+ ):
+ """
+ Self attention block
+ """
+ super().__init__()
+
+ self.norm1 = nn.LayerNorm(hidden_size)
+ self.norm2 = nn.LayerNorm(hidden_size)
+
+ self.attn = attn_class(embed_dim=hidden_size, num_heads=num_heads, batch_first=True, **block_kwargs)
+
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+ self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, drop=0)
+
+ def forward(self, x, mask=None):
+ # Prepare the mask for PyTorch's attention (it expects a different format)
+ # attn_mask = mask if mask is not None else None
+ # Normalize before attention
+ x = self.norm1(x)
+
+ # PyTorch's MultiheadAttention returns attn_output, attn_output_weights
+ # attn_output, _ = self.attn(x, x, x, attn_mask=attn_mask)
+
+ attn_output, _ = self.attn(x, x, x)
+
+ # Add & Norm
+ x = x + attn_output
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+class CrossAttnBlock(nn.Module):
+ def __init__(self, hidden_size, context_dim, num_heads=1, mlp_ratio=4.0, **block_kwargs):
+ """
+ Cross attention block
+ """
+ super().__init__()
+
+ self.norm1 = nn.LayerNorm(hidden_size)
+ self.norm_context = nn.LayerNorm(hidden_size)
+ self.norm2 = nn.LayerNorm(hidden_size)
+
+ self.cross_attn = nn.MultiheadAttention(
+ embed_dim=hidden_size, num_heads=num_heads, batch_first=True, **block_kwargs
+ )
+
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+
+ self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, drop=0)
+
+ def forward(self, x, context, mask=None):
+ # Normalize inputs
+ x = self.norm1(x)
+ context = self.norm_context(context)
+
+ # Apply cross attention
+ # Note: nn.MultiheadAttention returns attn_output, attn_output_weights
+ attn_output, _ = self.cross_attn(x, context, context, attn_mask=mask)
+
+ # Add & Norm
+ x = x + attn_output
+ x = x + self.mlp(self.norm2(x))
+ return x
diff --git a/src/models/aggregator/heads/track_modules/utils.py b/src/models/aggregator/heads/track_modules/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..51d01d39cdc10388a04dab5db7cf409b31bde766
--- /dev/null
+++ b/src/models/aggregator/heads/track_modules/utils.py
@@ -0,0 +1,226 @@
+# 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 from https://github.com/facebookresearch/vggsfm
+# and https://github.com/facebookresearch/co-tracker/tree/main
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from typing import Optional, Tuple, Union
+
+
+def get_2d_sincos_pos_embed(embed_dim: int, grid_size: Union[int, Tuple[int, int]], return_grid=False) -> torch.Tensor:
+ """
+ This function initializes a grid and generates a 2D positional embedding using sine and cosine functions.
+ It is a wrapper of get_2d_sincos_pos_embed_from_grid.
+ Args:
+ - embed_dim: The embedding dimension.
+ - grid_size: The grid size.
+ Returns:
+ - pos_embed: The generated 2D positional embedding.
+ """
+ if isinstance(grid_size, tuple):
+ grid_size_h, grid_size_w = grid_size
+ else:
+ grid_size_h = grid_size_w = grid_size
+ grid_h = torch.arange(grid_size_h, dtype=torch.float)
+ grid_w = torch.arange(grid_size_w, dtype=torch.float)
+ grid = torch.meshgrid(grid_w, grid_h, indexing="xy")
+ grid = torch.stack(grid, dim=0)
+ grid = grid.reshape([2, 1, grid_size_h, grid_size_w])
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+ if return_grid:
+ return (
+ pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2),
+ grid,
+ )
+ return pos_embed.reshape(1, grid_size_h, grid_size_w, -1).permute(0, 3, 1, 2)
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim: int, grid: torch.Tensor) -> torch.Tensor:
+ """
+ This function generates a 2D positional embedding from a given grid using sine and cosine functions.
+
+ Args:
+ - embed_dim: The embedding dimension.
+ - grid: The grid to generate the embedding from.
+
+ Returns:
+ - emb: The generated 2D positional embedding.
+ """
+ assert embed_dim % 2 == 0
+
+ # use half of dimensions to encode grid_h
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
+
+ emb = torch.cat([emb_h, emb_w], dim=2) # (H*W, D)
+ return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: torch.Tensor) -> torch.Tensor:
+ """
+ This function generates a 1D positional embedding from a given grid using sine and cosine functions.
+
+ Args:
+ - embed_dim: The embedding dimension.
+ - pos: The position to generate the embedding from.
+
+ Returns:
+ - emb: The generated 1D positional embedding.
+ """
+ assert embed_dim % 2 == 0
+ omega = torch.arange(embed_dim // 2, dtype=torch.double)
+ omega /= embed_dim / 2.0
+ omega = 1.0 / 10000**omega # (D/2,)
+
+ pos = pos.reshape(-1) # (M,)
+ out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product
+
+ emb_sin = torch.sin(out) # (M, D/2)
+ emb_cos = torch.cos(out) # (M, D/2)
+
+ emb = torch.cat([emb_sin, emb_cos], dim=1) # (M, D)
+ return emb[None].float()
+
+
+def get_2d_embedding(xy: torch.Tensor, C: int, cat_coords: bool = True) -> torch.Tensor:
+ """
+ This function generates a 2D positional embedding from given coordinates using sine and cosine functions.
+
+ Args:
+ - xy: The coordinates to generate the embedding from.
+ - C: The size of the embedding.
+ - cat_coords: A flag to indicate whether to concatenate the original coordinates to the embedding.
+
+ Returns:
+ - pe: The generated 2D positional embedding.
+ """
+ B, N, D = xy.shape
+ assert D == 2
+
+ x = xy[:, :, 0:1]
+ y = xy[:, :, 1:2]
+ div_term = (torch.arange(0, C, 2, device=xy.device, dtype=torch.float32) * (1000.0 / C)).reshape(1, 1, int(C / 2))
+
+ pe_x = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+ pe_y = torch.zeros(B, N, C, device=xy.device, dtype=torch.float32)
+
+ pe_x[:, :, 0::2] = torch.sin(x * div_term)
+ pe_x[:, :, 1::2] = torch.cos(x * div_term)
+
+ pe_y[:, :, 0::2] = torch.sin(y * div_term)
+ pe_y[:, :, 1::2] = torch.cos(y * div_term)
+
+ pe = torch.cat([pe_x, pe_y], dim=2) # (B, N, C*3)
+ if cat_coords:
+ pe = torch.cat([xy, pe], dim=2) # (B, N, C*3+3)
+ return pe
+
+
+def bilinear_sampler(input, coords, align_corners=True, padding_mode="border"):
+ r"""Sample a tensor using bilinear interpolation
+
+ `bilinear_sampler(input, coords)` samples a tensor :attr:`input` at
+ coordinates :attr:`coords` using bilinear interpolation. It is the same
+ as `torch.nn.functional.grid_sample()` but with a different coordinate
+ convention.
+
+ The input tensor is assumed to be of shape :math:`(B, C, H, W)`, where
+ :math:`B` is the batch size, :math:`C` is the number of channels,
+ :math:`H` is the height of the image, and :math:`W` is the width of the
+ image. The tensor :attr:`coords` of shape :math:`(B, H_o, W_o, 2)` is
+ interpreted as an array of 2D point coordinates :math:`(x_i,y_i)`.
+
+ Alternatively, the input tensor can be of size :math:`(B, C, T, H, W)`,
+ in which case sample points are triplets :math:`(t_i,x_i,y_i)`. Note
+ that in this case the order of the components is slightly different
+ from `grid_sample()`, which would expect :math:`(x_i,y_i,t_i)`.
+
+ If `align_corners` is `True`, the coordinate :math:`x` is assumed to be
+ in the range :math:`[0,W-1]`, with 0 corresponding to the center of the
+ left-most image pixel :math:`W-1` to the center of the right-most
+ pixel.
+
+ If `align_corners` is `False`, the coordinate :math:`x` is assumed to
+ be in the range :math:`[0,W]`, with 0 corresponding to the left edge of
+ the left-most pixel :math:`W` to the right edge of the right-most
+ pixel.
+
+ Similar conventions apply to the :math:`y` for the range
+ :math:`[0,H-1]` and :math:`[0,H]` and to :math:`t` for the range
+ :math:`[0,T-1]` and :math:`[0,T]`.
+
+ Args:
+ input (Tensor): batch of input images.
+ coords (Tensor): batch of coordinates.
+ align_corners (bool, optional): Coordinate convention. Defaults to `True`.
+ padding_mode (str, optional): Padding mode. Defaults to `"border"`.
+
+ Returns:
+ Tensor: sampled points.
+ """
+ coords = coords.detach().clone()
+ ############################################################
+ # IMPORTANT:
+ coords = coords.to(input.device).to(input.dtype)
+ ############################################################
+
+ sizes = input.shape[2:]
+
+ assert len(sizes) in [2, 3]
+
+ if len(sizes) == 3:
+ # t x y -> x y t to match dimensions T H W in grid_sample
+ coords = coords[..., [1, 2, 0]]
+
+ if align_corners:
+ scale = torch.tensor(
+ [2 / max(size - 1, 1) for size in reversed(sizes)], device=coords.device, dtype=coords.dtype
+ )
+ else:
+ scale = torch.tensor([2 / size for size in reversed(sizes)], device=coords.device, dtype=coords.dtype)
+
+ coords.mul_(scale) # coords = coords * scale
+ coords.sub_(1) # coords = coords - 1
+
+ return F.grid_sample(input, coords, align_corners=align_corners, padding_mode=padding_mode)
+
+
+def sample_features4d(input, coords):
+ r"""Sample spatial features
+
+ `sample_features4d(input, coords)` samples the spatial features
+ :attr:`input` represented by a 4D tensor :math:`(B, C, H, W)`.
+
+ The field is sampled at coordinates :attr:`coords` using bilinear
+ interpolation. :attr:`coords` is assumed to be of shape :math:`(B, R,
+ 2)`, where each sample has the format :math:`(x_i, y_i)`. This uses the
+ same convention as :func:`bilinear_sampler` with `align_corners=True`.
+
+ The output tensor has one feature per point, and has shape :math:`(B,
+ R, C)`.
+
+ Args:
+ input (Tensor): spatial features.
+ coords (Tensor): points.
+
+ Returns:
+ Tensor: sampled features.
+ """
+
+ B, _, _, _ = input.shape
+
+ # B R 2 -> B R 1 2
+ coords = coords.unsqueeze(2)
+
+ # B C R 1
+ feats = bilinear_sampler(input, coords)
+
+ return feats.permute(0, 2, 1, 3).view(B, -1, feats.shape[1] * feats.shape[3]) # B C R 1 -> B R C
diff --git a/src/models/aggregator/heads/utils.py b/src/models/aggregator/heads/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d7af1f68fa0ce0a48d11a708d53aa20aa8f78ba2
--- /dev/null
+++ b/src/models/aggregator/heads/utils.py
@@ -0,0 +1,108 @@
+# 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.
+
+import torch
+import torch.nn as nn
+
+
+def position_grid_to_embed(pos_grid: torch.Tensor, embed_dim: int, omega_0: float = 100) -> torch.Tensor:
+ """
+ Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC)
+
+ Args:
+ pos_grid: Tensor of shape (H, W, 2) containing 2D coordinates
+ embed_dim: Output channel dimension for embeddings
+
+ Returns:
+ Tensor of shape (H, W, embed_dim) with positional embeddings
+ """
+ H, W, grid_dim = pos_grid.shape
+ assert grid_dim == 2
+ pos_flat = pos_grid.reshape(-1, grid_dim) # Flatten to (H*W, 2)
+
+ # Process x and y coordinates separately
+ emb_x = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 0], omega_0=omega_0) # [1, H*W, D/2]
+ emb_y = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 1], omega_0=omega_0) # [1, H*W, D/2]
+
+ # Combine and reshape
+ emb = torch.cat([emb_x, emb_y], dim=-1) # [1, H*W, D]
+
+ return emb.view(H, W, embed_dim) # [H, W, D]
+
+
+def make_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, omega_0: float = 100) -> torch.Tensor:
+ """
+ This function generates a 1D positional embedding from a given grid using sine and cosine functions.
+
+ Args:
+ - embed_dim: The embedding dimension.
+ - pos: The position to generate the embedding from.
+
+ Returns:
+ - emb: The generated 1D positional embedding.
+ """
+ assert embed_dim % 2 == 0
+ omega = torch.arange(embed_dim // 2, dtype=torch.double, device=pos.device)
+ omega /= embed_dim / 2.0
+ omega = 1.0 / omega_0**omega # (D/2,)
+
+ pos = pos.reshape(-1) # (M,)
+ out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product
+
+ emb_sin = torch.sin(out) # (M, D/2)
+ emb_cos = torch.cos(out) # (M, D/2)
+
+ emb = torch.cat([emb_sin, emb_cos], dim=1) # (M, D)
+ return emb.float()
+
+
+# Inspired by https://github.com/microsoft/moge
+
+
+def create_uv_grid(
+ width: int, height: int, aspect_ratio: float = None, dtype: torch.dtype = None, device: torch.device = None
+) -> torch.Tensor:
+ """
+ Create a normalized UV grid of shape (width, height, 2).
+
+ The grid spans horizontally and vertically according to an aspect ratio,
+ ensuring the top-left corner is at (-x_span, -y_span) and the bottom-right
+ corner is at (x_span, y_span), normalized by the diagonal of the plane.
+
+ Args:
+ width (int): Number of points horizontally.
+ height (int): Number of points vertically.
+ aspect_ratio (float, optional): Width-to-height ratio. Defaults to width/height.
+ dtype (torch.dtype, optional): Data type of the resulting tensor.
+ device (torch.device, optional): Device on which the tensor is created.
+
+ Returns:
+ torch.Tensor: A (width, height, 2) tensor of UV coordinates.
+ """
+ # Derive aspect ratio if not explicitly provided
+ if aspect_ratio is None:
+ aspect_ratio = float(width) / float(height)
+
+ # Compute normalized spans for X and Y
+ diag_factor = (aspect_ratio**2 + 1.0) ** 0.5
+ span_x = aspect_ratio / diag_factor
+ span_y = 1.0 / diag_factor
+
+ # Establish the linspace boundaries
+ left_x = -span_x * (width - 1) / width
+ right_x = span_x * (width - 1) / width
+ top_y = -span_y * (height - 1) / height
+ bottom_y = span_y * (height - 1) / height
+
+ # Generate 1D coordinates
+ x_coords = torch.linspace(left_x, right_x, steps=width, dtype=dtype, device=device)
+ y_coords = torch.linspace(top_y, bottom_y, steps=height, dtype=dtype, device=device)
+
+ # Create 2D meshgrid (width x height) and stack into UV
+ uu, vv = torch.meshgrid(x_coords, y_coords, indexing="xy")
+ uv_grid = torch.stack((uu, vv), dim=-1)
+
+ return uv_grid
diff --git a/src/models/aggregator/layers/__init__.py b/src/models/aggregator/layers/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8120f4bc83066cb3f825ce32daa3b437f88486f1
--- /dev/null
+++ b/src/models/aggregator/layers/__init__.py
@@ -0,0 +1,11 @@
+# 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.
+
+from .mlp import Mlp
+from .patch_embed import PatchEmbed
+from .swiglu_ffn import SwiGLUFFN, SwiGLUFFNFused
+from .block import NestedTensorBlock
+from .attention import MemEffAttention
diff --git a/src/models/aggregator/layers/__pycache__/__init__.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f8d0884210ab863b2b716f4f48a18af03fc35eb9
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/__init__.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/__init__.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..85c98cbb745eb28c8485621c51dfb37d020e8e0c
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/__init__.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/attention.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/attention.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..793e6cb5e3af9ccc724c48b743d318651b0f8158
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/attention.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/attention.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/attention.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b69bdf3c9dc75b0aa4b341c9213242b7c998628e
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/attention.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/block.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/block.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..37f2c9bf00f4536b2e54f3c8983ff47a0cce357b
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/block.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/block.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/block.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fab2ad068337d51597d87a17892ab1855af7dd2d
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/block.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/drop_path.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/drop_path.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..004f501866937212610cd3b992a5366e2b284eec
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/drop_path.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/drop_path.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/drop_path.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3062810ea8cc4463e903920b31457b841e7af135
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/drop_path.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/layer_scale.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/layer_scale.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..34c208408d6097712a34bbc9f8a74513ddbc9eb3
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/layer_scale.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/layer_scale.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/layer_scale.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..84ffff84ae2449e92fe8523323ad30b8657cfcef
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/layer_scale.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/mlp.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/mlp.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..400bdde0cee8944a40ecd1a044a69e88f8a03a05
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/mlp.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/mlp.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/mlp.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..02d50adc369a6e8779c74d157c3a33f0311c0c8f
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/mlp.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/patch_embed.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/patch_embed.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..59a956750218b5bee78b588fb9d9e69fedded562
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/patch_embed.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/patch_embed.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/patch_embed.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8b65b1cbe8ed95b696be54a0889593811f04b181
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/patch_embed.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/rope.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/rope.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2aab7da8a8ca4564d11b27faa6396dd89d410f68
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/rope.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/rope.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/rope.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..5674489e5fb58d943efcde0ce10a5f8668a26f1a
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/rope.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/swiglu_ffn.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/swiglu_ffn.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..46c8553350e84869f3a716daa5f598bf9c725e52
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/swiglu_ffn.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/swiglu_ffn.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/swiglu_ffn.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f16e5ebfbd1343d2bbe14ef4a4dc86aeae755197
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/swiglu_ffn.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/vision_transformer.cpython-310.pyc b/src/models/aggregator/layers/__pycache__/vision_transformer.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3dfe26018cf46c08e723336eccb444799e149e7c
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/vision_transformer.cpython-310.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/vision_transformer.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/vision_transformer.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cca8f93a217893b5cb4e96d9d1c2951047c1968a
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/vision_transformer.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/__pycache__/vision_transformer_Net1.cpython-38.pyc b/src/models/aggregator/layers/__pycache__/vision_transformer_Net1.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..af8d1d09280591c050774f8923980856cdb0fe30
Binary files /dev/null and b/src/models/aggregator/layers/__pycache__/vision_transformer_Net1.cpython-38.pyc differ
diff --git a/src/models/aggregator/layers/attention.py b/src/models/aggregator/layers/attention.py
new file mode 100644
index 0000000000000000000000000000000000000000..72453f5159a874f3a9e05c130e343223d8b9c441
--- /dev/null
+++ b/src/models/aggregator/layers/attention.py
@@ -0,0 +1,114 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import logging
+import os
+import warnings
+
+from torch import unbind
+
+#from xformers.ops import memory_efficient_attention
+
+
+from torch import Tensor
+from torch import nn
+import torch.nn.functional as F
+
+XFORMERS_AVAILABLE = False
+
+
+class Attention(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int = 8,
+ qkv_bias: bool = True,
+ proj_bias: bool = True,
+ attn_drop: float = 0.0,
+ proj_drop: float = 0.0,
+ norm_layer: nn.Module = nn.LayerNorm,
+ qk_norm: bool = False,
+ fused_attn: bool = True, # use F.scaled_dot_product_attention or not
+ rope=None,
+ ) -> 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.scale = self.head_dim**-0.5
+ self.fused_attn = fused_attn
+
+ 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.attn_drop = nn.Dropout(attn_drop)
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
+ self.proj_drop = nn.Dropout(proj_drop)
+ self.rope = rope
+
+ def forward(self, x: Tensor, pos=None) -> 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)
+
+ if self.rope is not None:
+ q = self.rope(q, pos)
+ k = self.rope(k, pos)
+
+ if self.fused_attn:
+ x = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ dropout_p=self.attn_drop.p if self.training else 0.0,
+ )
+ else:
+ q = q * self.scale
+ attn = q @ k.transpose(-2, -1)
+ attn = attn.softmax(dim=-1)
+ attn = self.attn_drop(attn)
+ x = attn @ v
+ # attn_vis = ((q * self.scale) @ k.transpose(-2, -1)).softmax(dim=-1)
+ # import matplotlib.pyplot as plt
+ # for i in range (attn_vis.shape[0]):
+ # for j in range (attn_vis.shape[1]):
+ # for k in range(7):
+ # _ = attn_vis[i][j][k][7:].reshape(1,1,32,32)
+ # _ = F.interpolate(_, size=(256, 256), mode='nearest')
+ # _ = _.squeeze().cpu().detach().numpy()
+ # os.makedirs("/share/project/cwm/houyuan.chen/UPS_Lightning/plot/results/Net2_v15_epoch5_vis_for_goblet_attn_map/", exist_ok=True)
+ # plt.imsave(f"/share/project/cwm/houyuan.chen/UPS_Lightning/plot/results/Net2_v15_epoch5_vis_for_goblet_attn_map/attn_vis_img_{i}_head_{j}_token_{k}.png", _)
+
+
+ x = x.transpose(1, 2).reshape(B, N, C)
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
+
+
+class MemEffAttention(Attention):
+ def forward(self, x: Tensor, attn_bias=None, pos=None) -> Tensor:
+ assert pos is None
+ if not XFORMERS_AVAILABLE:
+ if attn_bias is not None:
+ raise AssertionError("xFormers is required for using nested tensors")
+ return super().forward(x)
+
+ B, N, C = x.shape
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+ q, k, v = unbind(qkv, 2)
+
+ x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+ x = x.reshape([B, N, C])
+
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
diff --git a/src/models/aggregator/layers/block.py b/src/models/aggregator/layers/block.py
new file mode 100644
index 0000000000000000000000000000000000000000..5f89e4da7121effca97151d1d8429586e422346e
--- /dev/null
+++ b/src/models/aggregator/layers/block.py
@@ -0,0 +1,259 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+import logging
+import os
+from typing import Callable, List, Any, Tuple, Dict
+import warnings
+
+import torch
+from torch import nn, Tensor
+
+from .attention import Attention
+from .drop_path import DropPath
+from .layer_scale import LayerScale
+from .mlp import Mlp
+
+
+XFORMERS_AVAILABLE = False
+
+
+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,
+ drop: float = 0.0,
+ attn_drop: float = 0.0,
+ init_values=None,
+ drop_path: float = 0.0,
+ act_layer: Callable[..., nn.Module] = nn.GELU,
+ norm_layer: Callable[..., nn.Module] = nn.LayerNorm,
+ attn_class: Callable[..., nn.Module] = Attention,
+ ffn_layer: Callable[..., nn.Module] = Mlp,
+ qk_norm: bool = False,
+ fused_attn: bool = True, # use F.scaled_dot_product_attention or not
+ rope=None,
+ ) -> None:
+ super().__init__()
+
+ self.norm1 = norm_layer(dim)
+
+ self.attn = attn_class(
+ dim,
+ num_heads=num_heads,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ attn_drop=attn_drop,
+ proj_drop=drop,
+ qk_norm=qk_norm,
+ fused_attn=fused_attn,
+ rope=rope,
+ )
+
+ self.ls1 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+ self.drop_path1 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ self.norm2 = norm_layer(dim)
+ mlp_hidden_dim = int(dim * mlp_ratio)
+ self.mlp = ffn_layer(
+ in_features=dim,
+ hidden_features=mlp_hidden_dim,
+ act_layer=act_layer,
+ drop=drop,
+ bias=ffn_bias,
+ )
+ self.ls2 = LayerScale(dim, init_values=init_values) if init_values else nn.Identity()
+ self.drop_path2 = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ self.sample_drop_ratio = drop_path
+
+ def forward(self, x: Tensor, pos=None) -> Tensor:
+ def attn_residual_func(x: Tensor, pos=None) -> Tensor:
+ return self.ls1(self.attn(self.norm1(x), pos=pos))
+
+ def ffn_residual_func(x: Tensor) -> Tensor:
+ return self.ls2(self.mlp(self.norm2(x)))
+
+ if self.training and self.sample_drop_ratio > 0.1:
+ # the overhead is compensated only for a drop path rate larger than 0.1
+ x = drop_add_residual_stochastic_depth(
+ x,
+ pos=pos,
+ residual_func=attn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ )
+ x = drop_add_residual_stochastic_depth(
+ x,
+ residual_func=ffn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ )
+ elif self.training and self.sample_drop_ratio > 0.0:
+ x = x + self.drop_path1(attn_residual_func(x, pos=pos))
+ x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2
+ else:
+ x = x + attn_residual_func(x, pos=pos)
+ x = x + ffn_residual_func(x)
+ return x
+
+
+def drop_add_residual_stochastic_depth(
+ x: Tensor,
+ residual_func: Callable[[Tensor], Tensor],
+ sample_drop_ratio: float = 0.0,
+ pos=None,
+) -> Tensor:
+ # 1) extract subset using permutation
+ b, n, d = x.shape
+ sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+ brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+ x_subset = x[brange]
+
+ # 2) apply residual_func to get residual
+ if pos is not None:
+ # if necessary, apply rope to the subset
+ pos = pos[brange]
+ residual = residual_func(x_subset, pos=pos)
+ else:
+ residual = residual_func(x_subset)
+
+ x_flat = x.flatten(1)
+ residual = residual.flatten(1)
+
+ residual_scale_factor = b / sample_subset_size
+
+ # 3) add the residual
+ x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+ return x_plus_residual.view_as(x)
+
+
+def get_branges_scales(x, sample_drop_ratio=0.0):
+ b, n, d = x.shape
+ sample_subset_size = max(int(b * (1 - sample_drop_ratio)), 1)
+ brange = (torch.randperm(b, device=x.device))[:sample_subset_size]
+ residual_scale_factor = b / sample_subset_size
+ return brange, residual_scale_factor
+
+
+def add_residual(x, brange, residual, residual_scale_factor, scaling_vector=None):
+ if scaling_vector is None:
+ x_flat = x.flatten(1)
+ residual = residual.flatten(1)
+ x_plus_residual = torch.index_add(x_flat, 0, brange, residual.to(dtype=x.dtype), alpha=residual_scale_factor)
+ else:
+ x_plus_residual = scaled_index_add(
+ x, brange, residual.to(dtype=x.dtype), scaling=scaling_vector, alpha=residual_scale_factor
+ )
+ return x_plus_residual
+
+
+attn_bias_cache: Dict[Tuple, Any] = {}
+
+
+def get_attn_bias_and_cat(x_list, branges=None):
+ """
+ this will perform the index select, cat the tensors, and provide the attn_bias from cache
+ """
+ batch_sizes = [b.shape[0] for b in branges] if branges is not None else [x.shape[0] for x in x_list]
+ all_shapes = tuple((b, x.shape[1]) for b, x in zip(batch_sizes, x_list))
+ if all_shapes not in attn_bias_cache.keys():
+ seqlens = []
+ for b, x in zip(batch_sizes, x_list):
+ for _ in range(b):
+ seqlens.append(x.shape[1])
+ attn_bias = fmha.BlockDiagonalMask.from_seqlens(seqlens)
+ attn_bias._batch_sizes = batch_sizes
+ attn_bias_cache[all_shapes] = attn_bias
+
+ if branges is not None:
+ cat_tensors = index_select_cat([x.flatten(1) for x in x_list], branges).view(1, -1, x_list[0].shape[-1])
+ else:
+ tensors_bs1 = tuple(x.reshape([1, -1, *x.shape[2:]]) for x in x_list)
+ cat_tensors = torch.cat(tensors_bs1, dim=1)
+
+ return attn_bias_cache[all_shapes], cat_tensors
+
+
+def drop_add_residual_stochastic_depth_list(
+ x_list: List[Tensor],
+ residual_func: Callable[[Tensor, Any], Tensor],
+ sample_drop_ratio: float = 0.0,
+ scaling_vector=None,
+) -> Tensor:
+ # 1) generate random set of indices for dropping samples in the batch
+ branges_scales = [get_branges_scales(x, sample_drop_ratio=sample_drop_ratio) for x in x_list]
+ branges = [s[0] for s in branges_scales]
+ residual_scale_factors = [s[1] for s in branges_scales]
+
+ # 2) get attention bias and index+concat the tensors
+ attn_bias, x_cat = get_attn_bias_and_cat(x_list, branges)
+
+ # 3) apply residual_func to get residual, and split the result
+ residual_list = attn_bias.split(residual_func(x_cat, attn_bias=attn_bias)) # type: ignore
+
+ outputs = []
+ for x, brange, residual, residual_scale_factor in zip(x_list, branges, residual_list, residual_scale_factors):
+ outputs.append(add_residual(x, brange, residual, residual_scale_factor, scaling_vector).view_as(x))
+ return outputs
+
+
+class NestedTensorBlock(Block):
+ def forward_nested(self, x_list: List[Tensor]) -> List[Tensor]:
+ """
+ x_list contains a list of tensors to nest together and run
+ """
+ assert isinstance(self.attn, MemEffAttention)
+
+ if self.training and self.sample_drop_ratio > 0.0:
+
+ def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.attn(self.norm1(x), attn_bias=attn_bias)
+
+ def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.mlp(self.norm2(x))
+
+ x_list = drop_add_residual_stochastic_depth_list(
+ x_list,
+ residual_func=attn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ scaling_vector=self.ls1.gamma if isinstance(self.ls1, LayerScale) else None,
+ )
+ x_list = drop_add_residual_stochastic_depth_list(
+ x_list,
+ residual_func=ffn_residual_func,
+ sample_drop_ratio=self.sample_drop_ratio,
+ scaling_vector=self.ls2.gamma if isinstance(self.ls1, LayerScale) else None,
+ )
+ return x_list
+ else:
+
+ def attn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.ls1(self.attn(self.norm1(x), attn_bias=attn_bias))
+
+ def ffn_residual_func(x: Tensor, attn_bias=None) -> Tensor:
+ return self.ls2(self.mlp(self.norm2(x)))
+
+ attn_bias, x = get_attn_bias_and_cat(x_list)
+ x = x + attn_residual_func(x, attn_bias=attn_bias)
+ x = x + ffn_residual_func(x)
+ return attn_bias.split(x)
+
+ def forward(self, x_or_x_list):
+ if isinstance(x_or_x_list, Tensor):
+ return super().forward(x_or_x_list)
+ elif isinstance(x_or_x_list, list):
+ if not XFORMERS_AVAILABLE:
+ raise AssertionError("xFormers is required for using nested tensors")
+ return self.forward_nested(x_or_x_list)
+ else:
+ raise AssertionError
diff --git a/src/models/aggregator/layers/drop_path.py b/src/models/aggregator/layers/drop_path.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d640e0b969b8dcba96260243473700b4e5b24b5
--- /dev/null
+++ b/src/models/aggregator/layers/drop_path.py
@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/drop.py
+
+
+from torch import nn
+
+
+def drop_path(x, drop_prob: float = 0.0, training: bool = False):
+ if drop_prob == 0.0 or not training:
+ return x
+ keep_prob = 1 - drop_prob
+ shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+ if keep_prob > 0.0:
+ random_tensor.div_(keep_prob)
+ output = x * random_tensor
+ return output
+
+
+class DropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob=None):
+ super(DropPath, self).__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, x):
+ return drop_path(x, self.drop_prob, self.training)
diff --git a/src/models/aggregator/layers/layer_scale.py b/src/models/aggregator/layers/layer_scale.py
new file mode 100644
index 0000000000000000000000000000000000000000..51df0d7ce61f2b41fa9e6369f52391dd7fe7d386
--- /dev/null
+++ b/src/models/aggregator/layers/layer_scale.py
@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# Modified from: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L103-L110
+
+from typing import Union
+
+import torch
+from torch import Tensor
+from torch import nn
+
+
+class LayerScale(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ init_values: Union[float, Tensor] = 1e-5,
+ inplace: bool = False,
+ ) -> None:
+ super().__init__()
+ self.inplace = inplace
+ self.gamma = nn.Parameter(init_values * torch.ones(dim))
+
+ def forward(self, x: Tensor) -> Tensor:
+ return x.mul_(self.gamma) if self.inplace else x * self.gamma
diff --git a/src/models/aggregator/layers/mlp.py b/src/models/aggregator/layers/mlp.py
new file mode 100644
index 0000000000000000000000000000000000000000..bbf9432aae9258612caeae910a7bde17999e328e
--- /dev/null
+++ b/src/models/aggregator/layers/mlp.py
@@ -0,0 +1,40 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/mlp.py
+
+
+from typing import Callable, Optional
+
+from torch import Tensor, nn
+
+
+class Mlp(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_features: Optional[int] = None,
+ out_features: Optional[int] = None,
+ act_layer: Callable[..., nn.Module] = nn.GELU,
+ drop: float = 0.0,
+ 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)
+ self.drop = nn.Dropout(drop)
+
+ def forward(self, x: Tensor) -> Tensor:
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.drop(x)
+ x = self.fc2(x)
+ x = self.drop(x)
+ return x
diff --git a/src/models/aggregator/layers/patch_embed.py b/src/models/aggregator/layers/patch_embed.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b7c0804784a42cf80c0297d110dcc68cc85b339
--- /dev/null
+++ b/src/models/aggregator/layers/patch_embed.py
@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/layers/patch_embed.py
+
+from typing import Callable, Optional, Tuple, Union
+
+from torch import Tensor
+import torch.nn as nn
+
+
+def make_2tuple(x):
+ 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)
+
+ Args:
+ img_size: Image size.
+ patch_size: Patch token size.
+ in_chans: Number of input image channels.
+ embed_dim: Number of linear projection output channels.
+ norm_layer: Normalization layer.
+ """
+
+ def __init__(
+ self,
+ img_size: Union[int, Tuple[int, int]] = 224,
+ patch_size: Union[int, Tuple[int, int]] = 16,
+ in_chans: int = 3,
+ embed_dim: int = 768,
+ norm_layer: Optional[Callable] = None,
+ flatten_embedding: bool = True,
+ ) -> None:
+ super().__init__()
+
+ image_HW = make_2tuple(img_size)
+ patch_HW = make_2tuple(patch_size)
+ patch_grid_size = (
+ image_HW[0] // patch_HW[0],
+ image_HW[1] // patch_HW[1],
+ )
+
+ self.img_size = image_HW
+ self.patch_size = patch_HW
+ self.patches_resolution = patch_grid_size
+ self.num_patches = patch_grid_size[0] * patch_grid_size[1]
+
+ self.in_chans = in_chans
+ self.embed_dim = embed_dim
+
+ self.flatten_embedding = flatten_embedding
+
+ 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) # B C H W
+ H, W = x.size(2), x.size(3)
+ x = x.flatten(2).transpose(1, 2) # B HW C
+ x = self.norm(x)
+ if not self.flatten_embedding:
+ x = x.reshape(-1, H, W, self.embed_dim) # B H W C
+ return x
+
+ def flops(self) -> float:
+ Ho, Wo = self.patches_resolution
+ flops = Ho * Wo * self.embed_dim * self.in_chans * (self.patch_size[0] * self.patch_size[1])
+ if self.norm is not None:
+ flops += Ho * Wo * self.embed_dim
+ return flops
diff --git a/src/models/aggregator/layers/rope.py b/src/models/aggregator/layers/rope.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d5d33304e55dbd05687bd86752a47a80e5f82df
--- /dev/null
+++ b/src/models/aggregator/layers/rope.py
@@ -0,0 +1,188 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+
+# Implementation of 2D Rotary Position Embeddings (RoPE).
+
+# This module provides a clean implementation of 2D Rotary Position Embeddings,
+# which extends the original RoPE concept to handle 2D spatial positions.
+
+# Inspired by:
+# https://github.com/meta-llama/codellama/blob/main/llama/model.py
+# https://github.com/naver-ai/rope-vit
+
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Dict, Tuple
+
+
+class PositionGetter:
+ """Generates and caches 2D spatial positions for patches in a grid.
+
+ This class efficiently manages the generation of spatial coordinates for patches
+ in a 2D grid, caching results to avoid redundant computations.
+
+ Attributes:
+ position_cache: Dictionary storing precomputed position tensors for different
+ grid dimensions.
+ """
+
+ def __init__(self):
+ """Initializes the position generator with an empty cache."""
+ self.position_cache: Dict[Tuple[int, int], torch.Tensor] = {}
+
+ def __call__(self, batch_size: int, height: int, width: int, device: torch.device) -> torch.Tensor:
+ """Generates spatial positions for a batch of patches.
+
+ Args:
+ batch_size: Number of samples in the batch.
+ height: Height of the grid in patches.
+ width: Width of the grid in patches.
+ device: Target device for the position tensor.
+
+ Returns:
+ Tensor of shape (batch_size, height*width, 2) containing y,x coordinates
+ for each position in the grid, repeated for each batch item.
+ """
+ if (height, width) not in self.position_cache:
+ y_coords = torch.arange(height, device=device)
+ x_coords = torch.arange(width, device=device)
+ positions = torch.cartesian_prod(y_coords, x_coords)
+ self.position_cache[height, width] = positions
+
+ cached_positions = self.position_cache[height, width]
+ return cached_positions.view(1, height * width, 2).expand(batch_size, -1, -1).clone()
+
+
+class RotaryPositionEmbedding2D(nn.Module):
+ """2D Rotary Position Embedding implementation.
+
+ This module applies rotary position embeddings to input tokens based on their
+ 2D spatial positions. It handles the position-dependent rotation of features
+ separately for vertical and horizontal dimensions.
+
+ Args:
+ frequency: Base frequency for the position embeddings. Default: 100.0
+ scaling_factor: Scaling factor for frequency computation. Default: 1.0
+
+ Attributes:
+ base_frequency: Base frequency for computing position embeddings.
+ scaling_factor: Factor to scale the computed frequencies.
+ frequency_cache: Cache for storing precomputed frequency components.
+ """
+
+ def __init__(self, frequency: float = 100.0, scaling_factor: float = 1.0):
+ """Initializes the 2D RoPE module."""
+ super().__init__()
+ self.base_frequency = frequency
+ self.scaling_factor = scaling_factor
+ self.frequency_cache: Dict[Tuple, Tuple[torch.Tensor, torch.Tensor]] = {}
+
+ def _compute_frequency_components(
+ self, dim: int, seq_len: int, device: torch.device, dtype: torch.dtype
+ ) -> Tuple[torch.Tensor, torch.Tensor]:
+ """Computes frequency components for rotary embeddings.
+
+ Args:
+ dim: Feature dimension (must be even).
+ seq_len: Maximum sequence length.
+ device: Target device for computations.
+ dtype: Data type for the computed tensors.
+
+ Returns:
+ Tuple of (cosine, sine) tensors for frequency components.
+ """
+ cache_key = (dim, seq_len, device, dtype)
+ if cache_key not in self.frequency_cache:
+ # Compute frequency bands
+ exponents = torch.arange(0, dim, 2, device=device).float() / dim
+ inv_freq = 1.0 / (self.base_frequency**exponents)
+
+ # Generate position-dependent frequencies
+ positions = torch.arange(seq_len, device=device, dtype=inv_freq.dtype)
+ angles = torch.einsum("i,j->ij", positions, inv_freq)
+
+ # Compute and cache frequency components
+ angles = angles.to(dtype)
+ angles = torch.cat((angles, angles), dim=-1)
+ cos_components = angles.cos().to(dtype)
+ sin_components = angles.sin().to(dtype)
+ self.frequency_cache[cache_key] = (cos_components, sin_components)
+
+ return self.frequency_cache[cache_key]
+
+ @staticmethod
+ def _rotate_features(x: torch.Tensor) -> torch.Tensor:
+ """Performs feature rotation by splitting and recombining feature dimensions.
+
+ Args:
+ x: Input tensor to rotate.
+
+ Returns:
+ Rotated feature tensor.
+ """
+ feature_dim = x.shape[-1]
+ x1, x2 = x[..., : feature_dim // 2], x[..., feature_dim // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+ def _apply_1d_rope(
+ self, tokens: torch.Tensor, positions: torch.Tensor, cos_comp: torch.Tensor, sin_comp: torch.Tensor
+ ) -> torch.Tensor:
+ """Applies 1D rotary position embeddings along one dimension.
+
+ Args:
+ tokens: Input token features.
+ positions: Position indices.
+ cos_comp: Cosine components for rotation.
+ sin_comp: Sine components for rotation.
+
+ Returns:
+ Tokens with applied rotary position embeddings.
+ """
+ # Embed positions with frequency components
+ cos = F.embedding(positions, cos_comp)[:, None, :, :]
+ sin = F.embedding(positions, sin_comp)[:, None, :, :]
+
+ # Apply rotation
+ return (tokens * cos) + (self._rotate_features(tokens) * sin)
+
+ def forward(self, tokens: torch.Tensor, positions: torch.Tensor) -> torch.Tensor:
+ """Applies 2D rotary position embeddings to input tokens.
+
+ Args:
+ tokens: Input tensor of shape (batch_size, n_heads, n_tokens, dim).
+ The feature dimension (dim) must be divisible by 4.
+ positions: Position tensor of shape (batch_size, n_tokens, 2) containing
+ the y and x coordinates for each token.
+
+ Returns:
+ Tensor of same shape as input with applied 2D rotary position embeddings.
+
+ Raises:
+ AssertionError: If input dimensions are invalid or positions are malformed.
+ """
+ # Validate inputs
+ assert tokens.size(-1) % 2 == 0, "Feature dimension must be even"
+ assert positions.ndim == 3 and positions.shape[-1] == 2, "Positions must have shape (batch_size, n_tokens, 2)"
+
+ # Compute feature dimension for each spatial direction
+ feature_dim = tokens.size(-1) // 2
+
+ # Get frequency components
+ max_position = int(positions.max()) + 1
+ cos_comp, sin_comp = self._compute_frequency_components(feature_dim, max_position, tokens.device, tokens.dtype)
+
+ # Split features for vertical and horizontal processing
+ vertical_features, horizontal_features = tokens.chunk(2, dim=-1)
+
+ # Apply RoPE separately for each dimension
+ vertical_features = self._apply_1d_rope(vertical_features, positions[..., 0], cos_comp, sin_comp)
+ horizontal_features = self._apply_1d_rope(horizontal_features, positions[..., 1], cos_comp, sin_comp)
+
+ # Combine processed features
+ return torch.cat((vertical_features, horizontal_features), dim=-1)
diff --git a/src/models/aggregator/layers/swiglu_ffn.py b/src/models/aggregator/layers/swiglu_ffn.py
new file mode 100644
index 0000000000000000000000000000000000000000..54fe8e90b7bedf6fbdbf09c6215844e3cc63f857
--- /dev/null
+++ b/src/models/aggregator/layers/swiglu_ffn.py
@@ -0,0 +1,72 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+import os
+from typing import Callable, Optional
+import warnings
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+
+
+class SwiGLUFFN(nn.Module):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_features: Optional[int] = None,
+ out_features: Optional[int] = None,
+ act_layer: Callable[..., nn.Module] = None,
+ drop: float = 0.0,
+ bias: bool = True,
+ ) -> None:
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias)
+ self.w3 = nn.Linear(hidden_features, out_features, bias=bias)
+
+ def forward(self, x: Tensor) -> Tensor:
+ x12 = self.w12(x)
+ x1, x2 = x12.chunk(2, dim=-1)
+ hidden = F.silu(x1) * x2
+ return self.w3(hidden)
+
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+# try:
+# if XFORMERS_ENABLED:
+# from xformers.ops import SwiGLU
+
+# XFORMERS_AVAILABLE = True
+# warnings.warn("xFormers is available (SwiGLU)")
+# else:
+# warnings.warn("xFormers is disabled (SwiGLU)")
+# raise ImportError
+# except ImportError:
+SwiGLU = SwiGLUFFN
+XFORMERS_AVAILABLE = False
+
+# warnings.warn("xFormers is not available (SwiGLU)")
+
+
+class SwiGLUFFNFused(SwiGLU):
+ def __init__(
+ self,
+ in_features: int,
+ hidden_features: Optional[int] = None,
+ out_features: Optional[int] = None,
+ act_layer: Callable[..., nn.Module] = None,
+ drop: float = 0.0,
+ bias: bool = True,
+ ) -> None:
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+ super().__init__(
+ in_features=in_features,
+ hidden_features=hidden_features,
+ out_features=out_features,
+ bias=bias,
+ )
diff --git a/src/models/aggregator/layers/vae.py b/src/models/aggregator/layers/vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..65827e887a46f40f00b613e905b7c46c7cbb65dd
--- /dev/null
+++ b/src/models/aggregator/layers/vae.py
@@ -0,0 +1,23 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers import AutoencoderKL
+from . import PatchEmbed
+class VAE(nn.Module):
+ # VAE代替DINOv2作为encoder的实现
+ # Decoder的部分也顺便写了
+ def __init__(self):
+ super(VAE, self).__init__()
+ self.vae = AutoencoderKL.from_pretrained("/home/lihong/UPS_Lightning/vae/vae/stable-diffusion-3.5-large", subfolder="vae").requires_grad_(False) # vae不训
+ def encode(self, x):
+ """
+ x: [B*f,3,H,W](multi-lihgt) or x: [B,3,H,W](nml)
+ """
+ z = self.vae.encode(x).latent_dist.sample() # [B*f,16,64,64]
+ return z
+ def decode(self, latent):
+ """
+ x: [B,16,64,64] nml的latent
+ """
+ decode_nml = self.vae.decode(latent).sample
+ return decode_nml #
\ No newline at end of file
diff --git a/src/models/aggregator/layers/vision_transformer.py b/src/models/aggregator/layers/vision_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..d456325768c0b6221c744514e50e3e1106edc46f
--- /dev/null
+++ b/src/models/aggregator/layers/vision_transformer.py
@@ -0,0 +1,407 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/main/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+from functools import partial
+import math
+import logging
+from typing import Sequence, Tuple, Union, Callable
+
+import torch
+import torch.nn as nn
+from torch.utils.checkpoint import checkpoint
+from torch.nn.init import trunc_normal_
+from . import Mlp, PatchEmbed, SwiGLUFFNFused, MemEffAttention, NestedTensorBlock as Block
+
+logger = logging.getLogger("dinov2")
+
+
+def named_apply(fn: Callable, module: nn.Module, name="", depth_first=True, include_root=False) -> nn.Module:
+ if not depth_first and include_root:
+ fn(module=module, name=name)
+ for child_name, child_module in module.named_children():
+ child_name = ".".join((name, child_name)) if name else child_name
+ named_apply(fn=fn, module=child_module, name=child_name, depth_first=depth_first, include_root=True)
+ if depth_first and include_root:
+ fn(module=module, name=name)
+ return module
+
+
+class BlockChunk(nn.ModuleList):
+ def forward(self, x):
+ for b in self:
+ x = b(x)
+ return x
+
+
+class DinoVisionTransformer(nn.Module):
+ def __init__(
+ self,
+ img_size=224,
+ patch_size=16,
+ in_chans=3,
+ embed_dim=768,
+ depth=12,
+ num_heads=12,
+ mlp_ratio=4.0,
+ qkv_bias=True,
+ ffn_bias=True,
+ proj_bias=True,
+ drop_path_rate=0.0,
+ drop_path_uniform=False,
+ init_values=None, # for layerscale: None or 0 => no layerscale
+ embed_layer=PatchEmbed,
+ act_layer=nn.GELU,
+ block_fn=Block,
+ ffn_layer="mlp",
+ block_chunks=1,
+ num_register_tokens=0,
+ interpolate_antialias=False,
+ interpolate_offset=0.1,
+ qk_norm=False,
+ ):
+ """
+ Args:
+ img_size (int, tuple): input image size
+ patch_size (int, tuple): patch size
+ in_chans (int): number of input channels
+ embed_dim (int): embedding dimension
+ depth (int): depth of transformer
+ num_heads (int): number of attention heads
+ mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+ qkv_bias (bool): enable bias for qkv if True
+ proj_bias (bool): enable bias for proj in attn if True
+ ffn_bias (bool): enable bias for ffn if True
+ drop_path_rate (float): stochastic depth rate
+ drop_path_uniform (bool): apply uniform drop rate across blocks
+ weight_init (str): weight init scheme
+ init_values (float): layer-scale init values
+ embed_layer (nn.Module): patch embedding layer
+ act_layer (nn.Module): MLP activation layer
+ block_fn (nn.Module): transformer block class
+ ffn_layer (str): "mlp", "swiglu", "swiglufused" or "identity"
+ block_chunks: (int) split block sequence into block_chunks units for FSDP wrap
+ num_register_tokens: (int) number of extra cls tokens (so-called "registers")
+ interpolate_antialias: (str) flag to apply anti-aliasing when interpolating positional embeddings
+ interpolate_offset: (float) work-around offset to apply when interpolating positional embeddings
+ """
+ super().__init__()
+ norm_layer = partial(nn.LayerNorm, eps=1e-6)
+
+ # tricky but makes it work
+ self.use_checkpoint = False
+ #
+
+ self.num_features = self.embed_dim = embed_dim # num_features for consistency with other models
+ self.num_tokens = 1
+ self.n_blocks = depth
+ self.num_heads = num_heads
+ self.patch_size = patch_size
+ self.num_register_tokens = num_register_tokens
+ self.interpolate_antialias = interpolate_antialias
+ self.interpolate_offset = interpolate_offset
+
+ self.patch_embed = embed_layer(img_size=img_size, patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+ num_patches = self.patch_embed.num_patches
+
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+ self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, embed_dim))
+ assert num_register_tokens >= 0
+ self.register_tokens = (
+ nn.Parameter(torch.zeros(1, num_register_tokens, embed_dim)) if num_register_tokens else None
+ )
+
+ if drop_path_uniform is True:
+ dpr = [drop_path_rate] * depth
+ else:
+ dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] # stochastic depth decay rule
+
+ if ffn_layer == "mlp":
+ logger.info("using MLP layer as FFN")
+ ffn_layer = Mlp
+ elif ffn_layer == "swiglufused" or ffn_layer == "swiglu":
+ logger.info("using SwiGLU layer as FFN")
+ ffn_layer = SwiGLUFFNFused
+ elif ffn_layer == "identity":
+ logger.info("using Identity layer as FFN")
+
+ def f(*args, **kwargs):
+ return nn.Identity()
+
+ ffn_layer = f
+ else:
+ raise NotImplementedError
+
+ blocks_list = [
+ block_fn(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ ffn_bias=ffn_bias,
+ drop_path=dpr[i],
+ norm_layer=norm_layer,
+ act_layer=act_layer,
+ ffn_layer=ffn_layer,
+ init_values=init_values,
+ qk_norm=qk_norm,
+ )
+ for i in range(depth)
+ ]
+ if block_chunks > 0:
+ self.chunked_blocks = True
+ chunked_blocks = []
+ chunksize = depth // block_chunks
+ for i in range(0, depth, chunksize):
+ # this is to keep the block index consistent if we chunk the block list
+ chunked_blocks.append([nn.Identity()] * i + blocks_list[i : i + chunksize])
+ self.blocks = nn.ModuleList([BlockChunk(p) for p in chunked_blocks])
+ else:
+ self.chunked_blocks = False
+ self.blocks = nn.ModuleList(blocks_list)
+
+ self.norm = norm_layer(embed_dim)
+ self.head = nn.Identity()
+
+ self.mask_token = nn.Parameter(torch.zeros(1, embed_dim))
+
+ self.init_weights()
+
+ def init_weights(self):
+ trunc_normal_(self.pos_embed, std=0.02)
+ nn.init.normal_(self.cls_token, std=1e-6)
+ if self.register_tokens is not None:
+ nn.init.normal_(self.register_tokens, std=1e-6)
+ named_apply(init_weights_vit_timm, self)
+
+ def interpolate_pos_encoding(self, x, w, h):
+ previous_dtype = x.dtype
+ npatch = x.shape[1] - 1
+ N = self.pos_embed.shape[1] - 1
+ if npatch == N and w == h:
+ return self.pos_embed
+ pos_embed = self.pos_embed.float()
+ class_pos_embed = pos_embed[:, 0]
+ patch_pos_embed = pos_embed[:, 1:]
+ dim = x.shape[-1]
+ w0 = w // self.patch_size
+ h0 = h // self.patch_size
+ M = int(math.sqrt(N)) # Recover the number of patches in each dimension
+ assert N == M * M
+ kwargs = {}
+ if self.interpolate_offset:
+ # Historical kludge: add a small number to avoid floating point error in the interpolation, see https://github.com/facebookresearch/dino/issues/8
+ # Note: still needed for backward-compatibility, the underlying operators are using both output size and scale factors
+ sx = float(w0 + self.interpolate_offset) / M
+ sy = float(h0 + self.interpolate_offset) / M
+ kwargs["scale_factor"] = (sx, sy)
+ else:
+ # Simply specify an output size instead of a scale factor
+ kwargs["size"] = (w0, h0)
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed.reshape(1, M, M, dim).permute(0, 3, 1, 2),
+ mode="bicubic",
+ antialias=self.interpolate_antialias,
+ **kwargs,
+ )
+ assert (w0, h0) == patch_pos_embed.shape[-2:]
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+ return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1).to(previous_dtype)
+
+ def prepare_tokens_with_masks(self, x, masks=None):
+ B, nc, w, h = x.shape
+ x = self.patch_embed(x)
+ if masks is not None:
+ x = torch.where(masks.unsqueeze(-1), self.mask_token.to(x.dtype).unsqueeze(0), x)
+
+ x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1)
+ x = x + self.interpolate_pos_encoding(x, w, h)
+
+ if self.register_tokens is not None:
+ x = torch.cat(
+ (
+ x[:, :1],
+ self.register_tokens.expand(x.shape[0], -1, -1),
+ x[:, 1:],
+ ),
+ dim=1,
+ )
+
+ return x
+
+ def forward_features_list(self, x_list, masks_list):
+ x = [self.prepare_tokens_with_masks(x, masks) for x, masks in zip(x_list, masks_list)]
+
+ for blk in self.blocks:
+ if self.use_checkpoint:
+ x = checkpoint(blk, x, use_reentrant=self.use_reentrant)
+ else:
+ x = blk(x)
+
+ all_x = x
+ output = []
+ for x, masks in zip(all_x, masks_list):
+ x_norm = self.norm(x)
+ output.append(
+ {
+ "x_norm_clstoken": x_norm[:, 0],
+ "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+ "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+ "x_prenorm": x,
+ "masks": masks,
+ }
+ )
+ return output
+
+ def forward_features(self, x, masks=None):
+ if isinstance(x, list):
+ return self.forward_features_list(x, masks)
+
+ x = self.prepare_tokens_with_masks(x, masks) #tokenize
+
+ for blk in self.blocks:
+ if self.use_checkpoint:
+ x = checkpoint(blk, x, use_reentrant=self.use_reentrant)
+ else:
+ x = blk(x)
+
+ x_norm = self.norm(x)
+ return {
+ "x_norm_clstoken": x_norm[:, 0],
+ "x_norm_regtokens": x_norm[:, 1 : self.num_register_tokens + 1],
+ "x_norm_patchtokens": x_norm[:, self.num_register_tokens + 1 :],
+ "x_prenorm": x,
+ "masks": masks,
+ }
+
+ def _get_intermediate_layers_not_chunked(self, x, n=1):
+ x = self.prepare_tokens_with_masks(x)
+ # If n is an int, take the n last blocks. If it's a list, take them
+ output, total_block_len = [], len(self.blocks)
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+ for i, blk in enumerate(self.blocks):
+ x = blk(x)
+ if i in blocks_to_take:
+ output.append(x)
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+ return output
+
+ def _get_intermediate_layers_chunked(self, x, n=1):
+ x = self.prepare_tokens_with_masks(x)
+ output, i, total_block_len = [], 0, len(self.blocks[-1])
+ # If n is an int, take the n last blocks. If it's a list, take them
+ blocks_to_take = range(total_block_len - n, total_block_len) if isinstance(n, int) else n
+ for block_chunk in self.blocks:
+ for blk in block_chunk[i:]: # Passing the nn.Identity()
+ x = blk(x)
+ if i in blocks_to_take:
+ output.append(x)
+ i += 1
+ assert len(output) == len(blocks_to_take), f"only {len(output)} / {len(blocks_to_take)} blocks found"
+ return output
+
+ def get_intermediate_layers(
+ self,
+ x: torch.Tensor,
+ n: Union[int, Sequence] = 1, # Layers or n last layers to take
+ reshape: bool = False,
+ return_class_token: bool = False,
+ norm=True,
+ ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]]]:
+ if self.chunked_blocks:
+ outputs = self._get_intermediate_layers_chunked(x, n)
+ else:
+ outputs = self._get_intermediate_layers_not_chunked(x, n)
+ if norm:
+ outputs = [self.norm(out) for out in outputs]
+ class_tokens = [out[:, 0] for out in outputs]
+ outputs = [out[:, 1 + self.num_register_tokens :] for out in outputs]
+ if reshape:
+ B, _, w, h = x.shape
+ outputs = [
+ out.reshape(B, w // self.patch_size, h // self.patch_size, -1).permute(0, 3, 1, 2).contiguous()
+ for out in outputs
+ ]
+ if return_class_token:
+ return tuple(zip(outputs, class_tokens))
+ return tuple(outputs)
+
+ def forward(self, *args, is_training=True, **kwargs):
+ ret = self.forward_features(*args, **kwargs)
+ if is_training:
+ return ret
+ else:
+ return self.head(ret["x_norm_clstoken"])
+
+
+def init_weights_vit_timm(module: nn.Module, name: str = ""):
+ """ViT weight initialization, original timm impl (for reproducibility)"""
+ if isinstance(module, nn.Linear):
+ trunc_normal_(module.weight, std=0.02)
+ if module.bias is not None:
+ nn.init.zeros_(module.bias)
+
+
+def vit_small(patch_size=16, num_register_tokens=0, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=384,
+ depth=12,
+ num_heads=6,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
+
+
+def vit_base(patch_size=16, num_register_tokens=0, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=768,
+ depth=12,
+ num_heads=12,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
+
+
+def vit_large(patch_size=16, num_register_tokens=0, **kwargs):
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=1024,
+ depth=24,
+ num_heads=16,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
+
+
+def vit_giant2(patch_size=16, num_register_tokens=0, **kwargs):
+ """
+ Close to ViT-giant, with embed-dim 1536 and 24 heads => embed-dim per head 64
+ """
+ model = DinoVisionTransformer(
+ patch_size=patch_size,
+ embed_dim=1536,
+ depth=40,
+ num_heads=24,
+ mlp_ratio=4,
+ block_fn=partial(Block, attn_class=MemEffAttention),
+ num_register_tokens=num_register_tokens,
+ **kwargs,
+ )
+ return model
diff --git a/src/models/aggregator/models/__pycache__/aggregator.cpython-310.pyc b/src/models/aggregator/models/__pycache__/aggregator.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..88506bcea1550cacb58e7c37492988cde6a5b17c
Binary files /dev/null and b/src/models/aggregator/models/__pycache__/aggregator.cpython-310.pyc differ
diff --git a/src/models/aggregator/models/aggregator.py b/src/models/aggregator/models/aggregator.py
new file mode 100644
index 0000000000000000000000000000000000000000..814bb5a575887d800f77fd96885e13791854cf6c
--- /dev/null
+++ b/src/models/aggregator/models/aggregator.py
@@ -0,0 +1,417 @@
+
+
+
+
+
+
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Union, List, Dict, Any
+
+import cv2
+import torchvision
+from ..layers import PatchEmbed
+from ..layers.block import Block
+from ..layers.rope import RotaryPositionEmbedding2D, PositionGetter
+from ..layers.vision_transformer import vit_small, vit_base, vit_large, vit_giant2
+
+
+
+cv2.setNumThreads(1)
+logger = logging.getLogger(__name__)
+
+_RESNET_MEAN = [0.485, 0.456, 0.406]
+_RESNET_STD = [0.229, 0.224, 0.225]
+
+
+
+
+
+
+
+
+
+
+class Aggregator(nn.Module):
+ """
+ The Aggregator applies alternating-attention over input frames,
+ as described in VGGT: Visual Geometry Grounded Transformer.
+
+
+ Args:
+ img_size (int): Image size in pixels.
+ patch_size (int): Size of each patch for PatchEmbed.
+ embed_dim (int): Dimension of the token embeddings.
+ depth (int): Number of blocks.
+ num_heads (int): Number of attention heads.
+ mlp_ratio (float): Ratio of MLP hidden dim to embedding dim.
+ num_register_tokens (int): Number of register tokens.
+ block_fn (nn.Module): The block type used for attention (Block by default).
+ qkv_bias (bool): Whether to include bias in QKV projections.
+ proj_bias (bool): Whether to include bias in the output projection.
+ ffn_bias (bool): Whether to include bias in MLP layers.
+ patch_embed (str): Type of patch embed. e.g., "conv" or "dinov2_vitl14_reg".
+ aa_order (list[str]): The order of alternating attention, e.g. ["frame", "global"].
+ aa_block_size (int): How many blocks to group under each attention type before switching. If not necessary, set to 1.
+ qk_norm (bool): Whether to apply QK normalization.
+ rope_freq (int): Base frequency for rotary embedding. -1 to disable.
+ init_values (float): Init scale for layer scale.
+ """
+
+ def __init__(
+ self,
+ img_size=518,
+ patch_size=14,
+ embed_dim=1024,
+ depth=4,
+ num_heads=6,
+ mlp_ratio=4.0,
+ num_register_tokens=4,
+ block_fn=Block,
+ qkv_bias=True,
+ proj_bias=True,
+ ffn_bias=True,
+ patch_embed="dinov2_vitl14_reg",
+ aa_order=["frame", "light_axis_1","global","light_axis_2"],
+ aa_block_size=1,
+ qk_norm=True,
+ rope_freq=100,
+ init_values=0.01,
+ ):
+ super().__init__()
+
+ self.__build_patch_embed__(patch_embed, img_size, patch_size, num_register_tokens, embed_dim=embed_dim)
+
+
+ self.rope = RotaryPositionEmbedding2D(frequency=rope_freq) if rope_freq > 0 else None
+ self.position_getter = PositionGetter() if self.rope is not None else None
+ self.patch_start_idx = 7
+ self.frame_blocks = nn.ModuleList(
+ [
+ block_fn(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ ffn_bias=ffn_bias,
+ init_values=init_values,
+ qk_norm=qk_norm,
+ rope=self.rope,
+ )
+ for _ in range(depth)
+ ]
+ )
+
+ self.global_blocks = nn.ModuleList(
+ [
+ block_fn(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ ffn_bias=ffn_bias,
+ init_values=init_values,
+ qk_norm=qk_norm,
+ rope=self.rope,
+ )
+ for _ in range(depth)
+ ]
+ )
+ self.light_axis_1_blocks = nn.ModuleList(
+ [
+ block_fn(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ ffn_bias=ffn_bias,
+ init_values=init_values,
+ qk_norm=qk_norm,
+ rope=self.rope,
+ )
+ for _ in range(depth)
+ ]
+ )
+ self.light_axis_2_blocks = nn.ModuleList(
+ [
+ block_fn(
+ dim=embed_dim,
+ num_heads=num_heads,
+ mlp_ratio=mlp_ratio,
+ qkv_bias=qkv_bias,
+ proj_bias=proj_bias,
+ ffn_bias=ffn_bias,
+ init_values=init_values,
+ qk_norm=qk_norm,
+ rope=self.rope,
+ )
+ for _ in range(depth)
+ ]
+ )
+ self.depth = depth
+ self.aa_order = aa_order
+ self.patch_size = patch_size
+ self.aa_block_size = aa_block_size
+
+
+ if self.depth % self.aa_block_size != 0:
+ raise ValueError(f"depth ({depth}) must be divisible by aa_block_size ({aa_block_size})")
+
+ self.aa_block_num = self.depth // self.aa_block_size
+
+
+
+ def __build_patch_embed__(
+ self,
+ patch_embed,
+ img_size,
+ patch_size,
+ num_register_tokens,
+ interpolate_antialias=True,
+ interpolate_offset=0.0,
+ block_chunks=0,
+ init_values=1.0,
+ embed_dim=1024,
+ ):
+ """
+ Build the patch embed layer. If 'conv', we use a
+ simple PatchEmbed conv layer. Otherwise, we use a vision transformer.
+ """
+
+ if "conv" in patch_embed:
+ self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_chans=3, embed_dim=embed_dim)
+ else:
+ vit_models = {
+ "dinov2_vitl14_reg": vit_large,
+ "dinov2_vitb14_reg": vit_base,
+ "dinov2_vits14_reg": vit_small,
+ "dinov2_vitg2_reg": vit_giant2,
+ }
+
+ self.patch_embed = vit_models[patch_embed](
+ img_size=img_size,
+ patch_size=patch_size,
+ num_register_tokens=num_register_tokens,
+ interpolate_antialias=interpolate_antialias,
+ interpolate_offset=interpolate_offset,
+ block_chunks=block_chunks,
+ init_values=init_values
+
+ )
+
+
+
+
+ if hasattr(self.patch_embed, "mask_token"):
+ self.patch_embed.mask_token.requires_grad_(False)
+
+ def forward(
+ self,
+ images: torch.Tensor,
+ ) -> Tuple[List[torch.Tensor], int]:
+
+ B, S, C_in, H, W = images.shape
+
+
+
+
+
+
+
+
+
+
+ images = images.view(B * S, C_in, H, W)
+ patch_tokens = self.patch_embed(images)
+
+ if isinstance(patch_tokens, dict):
+ register_tokens = patch_tokens["x_norm_regtokens"]
+ patch_tokens = patch_tokens["x_norm_patchtokens"]
+
+ light_tokens = register_tokens[:,:3,:]
+
+
+ patch_tokens = torch.cat((light_tokens,register_tokens,patch_tokens), dim=1)
+ _, P, C = patch_tokens.shape
+ tokens = patch_tokens
+
+ pos = None
+ if self.rope is not None:
+ pos = self.position_getter(B * S, H // self.patch_size, W // self.patch_size, device=images.device)
+
+ if self.patch_start_idx > 0:
+
+
+ pos = pos + 1
+ pos_special = torch.zeros(B * S, self.patch_start_idx, 2).to(images.device).to(pos.dtype)
+ pos = torch.cat([pos_special, pos], dim=1)
+
+
+ _, P, C = tokens.shape
+
+ frame_idx = 0
+ global_idx = 0
+ light_axis_idx_1 = 0
+ light_axis_idx_2 = 0
+ output_list = []
+
+ for _ in range(self.aa_block_num):
+ for attn_type in self.aa_order:
+ if attn_type == "frame":
+ tokens, frame_idx, frame_intermediates = self._process_frame_attention(
+ tokens, B, S, P, C, frame_idx, pos=pos
+ )
+ elif attn_type == "global":
+ tokens, global_idx, global_intermediates = self._process_global_attention(
+ tokens, B, S, P, C, global_idx, pos=pos
+ )
+ elif attn_type == "light_axis_1":
+ tokens, light_axis_idx_1, light_axis_intermediates_1 = self._process_light_axis_1_attention(
+ tokens, B, S, P, C, light_axis_idx_1, pos=pos
+ )
+ elif attn_type == "light_axis_2":
+ tokens, light_axis_idx_2, light_axis_intermediates_2 = self._process_light_axis_2_attention(
+ tokens, B, S, P, C, light_axis_idx_2, pos=pos
+ )
+ else:
+ raise ValueError(f"Unknown attention type: {attn_type}")
+
+ for i in range(len(frame_intermediates)):
+
+ concat_inter = torch.cat([frame_intermediates[i],light_axis_intermediates_1[i], global_intermediates[i],light_axis_intermediates_2[i]], dim=-1)
+ output_list.append(concat_inter)
+
+ del concat_inter
+ del frame_intermediates
+ del global_intermediates
+ del light_axis_intermediates_1
+ del light_axis_intermediates_2
+ if self.depth == 6:
+ output_list.pop(1)
+ output_list.pop(2)
+ elif self.depth == 8:
+ output_list.pop(1)
+ output_list.pop(2)
+ output_list.pop(3)
+ output_list.pop(4)
+ return output_list, self.patch_start_idx
+
+ def _process_frame_attention(self, tokens, B, S, P, C, frame_idx, pos=None):
+ """
+ Process frame attention blocks. We keep tokens in shape (B*S, P, C).
+ """
+
+ if tokens.shape != (B * S, P, C):
+ tokens = tokens.view(B, S, P, C).view(B * S, P, C)
+
+ if pos is not None and pos.shape != (B * S, P, 2):
+ pos = pos.view(B, S, P, 2).view(B * S, P, 2)
+
+ intermediates = []
+
+
+ for _ in range(self.aa_block_size):
+ tokens = self.frame_blocks[frame_idx](tokens, pos=pos)
+ frame_idx += 1
+ intermediates.append(tokens.view(B, S, P, C))
+
+ return tokens, frame_idx, intermediates
+
+ def _process_global_attention(self, tokens, B, S, P, C, global_idx, pos=None):
+ """
+ Process global attention blocks. We keep tokens in shape (B, S*P, C).
+ """
+ if tokens.shape != (B, S * P, C):
+ tokens = tokens.view(B, S, P, C).view(B, S * P, C)
+
+ if pos is not None and pos.shape != (B, S * P, 2):
+ pos = pos.view(B, S, P, 2).view(B, S * P, 2)
+
+ intermediates = []
+
+
+ for _ in range(self.aa_block_size):
+ tokens = self.global_blocks[global_idx](tokens, pos=pos)
+ global_idx += 1
+ intermediates.append(tokens.view(B, S, P, C))
+
+ return tokens, global_idx, intermediates
+
+ def _process_light_axis_1_attention(self, tokens, B, S, P, C, light_axis_idx, pos=None):
+ """
+ Process light axis attention blocks. We keep tokens in shape (B*P, S, C).
+ """
+
+
+
+
+
+ tokens = tokens.view(B, S, P, C).permute(0,2,1,3).reshape(B * P, S, C)
+
+ if pos is not None and pos.shape != (B * P, S, 2):
+ pos = pos.view(B, S, P, 2).permute(0,2,1,3).reshape(B * P, S, 2)
+
+ intermediates = []
+
+
+ for _ in range(self.aa_block_size):
+ tokens = self.light_axis_1_blocks[light_axis_idx](tokens, pos=pos)
+ light_axis_idx += 1
+ tokens = tokens.view(B, P, S, C).permute(0,2,1,3)
+ intermediates.append(tokens)
+ tokens = tokens.reshape(B * S, P, C)
+
+ return tokens, light_axis_idx, intermediates
+
+ def _process_light_axis_2_attention(self, tokens, B, S, P, C, light_axis_idx, pos=None):
+ """
+ Process light axis attention blocks. We keep tokens in shape (B*P, S, C).
+ """
+
+
+ tokens = tokens.view(B, S, P, C).permute(0,2,1,3).reshape(B * P, S, C)
+
+ if pos is not None and pos.shape != (B * P, S, 2):
+ pos = pos.view(B, S, P, 2).permute(0,2,1,3).reshape(B * P, S, 2)
+
+ intermediates = []
+
+
+ for _ in range(self.aa_block_size):
+ tokens = self.light_axis_2_blocks[light_axis_idx](tokens, pos=pos)
+ light_axis_idx += 1
+ tokens = tokens.view(B, P, S, C).permute(0,2,1,3)
+ intermediates.append(tokens)
+ tokens = tokens.reshape(B * S, P, C)
+
+ return tokens, light_axis_idx, intermediates
+
+def slice_expand_and_flatten(token_tensor, B, S):
+ """
+ Processes specialized tokens with shape (1, 2, X, C) for multi-frame processing:
+ 1) Uses the first position (index=0) for the first frame only
+ 2) Uses the second position (index=1) for all remaining frames (S-1 frames)
+ 3) Expands both to match batch size B
+ 4) Concatenates to form (B, S, X, C) where each sequence has 1 first-position token
+ followed by (S-1) second-position tokens
+ 5) Flattens to (B*S, X, C) for processing
+
+ Returns:
+ torch.Tensor: Processed tokens with shape (B*S, X, C)
+ """
+
+
+ query = token_tensor[:, 0:1, ...].expand(B, 1, *token_tensor.shape[2:])
+
+ others = token_tensor[:, 1:, ...].expand(B, S - 1, *token_tensor.shape[2:])
+
+ combined = torch.cat([query, others], dim=1)
+
+
+ combined = combined.view(B * S, *combined.shape[2:])
+ return combined
\ No newline at end of file
diff --git a/src/models/aggregator/utils/__pycache__/geometry.cpython-310.pyc b/src/models/aggregator/utils/__pycache__/geometry.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cbffdf122f5f3e1dc515e7545ce09745834bae2b
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/geometry.cpython-310.pyc differ
diff --git a/src/models/aggregator/utils/__pycache__/geometry.cpython-38.pyc b/src/models/aggregator/utils/__pycache__/geometry.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0ae73088845cb8fd4553da764cfcfa813a0bf2e0
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/geometry.cpython-38.pyc differ
diff --git a/src/models/aggregator/utils/__pycache__/load_fn.cpython-310.pyc b/src/models/aggregator/utils/__pycache__/load_fn.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8df22e16de38b75560521474008e92130125df9d
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/load_fn.cpython-310.pyc differ
diff --git a/src/models/aggregator/utils/__pycache__/load_fn.cpython-38.pyc b/src/models/aggregator/utils/__pycache__/load_fn.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b06e6a3e3b6cc91562a513e683c3c82923fddc12
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/load_fn.cpython-38.pyc differ
diff --git a/src/models/aggregator/utils/__pycache__/pose_enc.cpython-310.pyc b/src/models/aggregator/utils/__pycache__/pose_enc.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2f5e81f40409e53261b6f59baa7868ad7938e435
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/pose_enc.cpython-310.pyc differ
diff --git a/src/models/aggregator/utils/__pycache__/pose_enc.cpython-38.pyc b/src/models/aggregator/utils/__pycache__/pose_enc.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..51d9b6f5d8dc0a3811932e7d33befed52cf77855
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/pose_enc.cpython-38.pyc differ
diff --git a/src/models/aggregator/utils/__pycache__/rotation.cpython-310.pyc b/src/models/aggregator/utils/__pycache__/rotation.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1f4df9b4782fece538c5b187c60b0029abe59ac2
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/rotation.cpython-310.pyc differ
diff --git a/src/models/aggregator/utils/__pycache__/rotation.cpython-38.pyc b/src/models/aggregator/utils/__pycache__/rotation.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..341e610a54581c55f68f814d32ef1bbdb32bd898
Binary files /dev/null and b/src/models/aggregator/utils/__pycache__/rotation.cpython-38.pyc differ
diff --git a/src/models/aggregator/utils/geometry.py b/src/models/aggregator/utils/geometry.py
new file mode 100644
index 0000000000000000000000000000000000000000..a85e6482d18889cb6c85e44992b164926b7e627e
--- /dev/null
+++ b/src/models/aggregator/utils/geometry.py
@@ -0,0 +1,166 @@
+# 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.
+
+import os
+import torch
+import numpy as np
+
+
+def unproject_depth_map_to_point_map(
+ depth_map: np.ndarray, extrinsics_cam: np.ndarray, intrinsics_cam: np.ndarray
+) -> np.ndarray:
+ """
+ Unproject a batch of depth maps to 3D world coordinates.
+
+ Args:
+ depth_map (np.ndarray): Batch of depth maps of shape (S, H, W, 1) or (S, H, W)
+ extrinsics_cam (np.ndarray): Batch of camera extrinsic matrices of shape (S, 3, 4)
+ intrinsics_cam (np.ndarray): Batch of camera intrinsic matrices of shape (S, 3, 3)
+
+ Returns:
+ np.ndarray: Batch of 3D world coordinates of shape (S, H, W, 3)
+ """
+ if isinstance(depth_map, torch.Tensor):
+ depth_map = depth_map.cpu().numpy()
+ if isinstance(extrinsics_cam, torch.Tensor):
+ extrinsics_cam = extrinsics_cam.cpu().numpy()
+ if isinstance(intrinsics_cam, torch.Tensor):
+ intrinsics_cam = intrinsics_cam.cpu().numpy()
+
+ world_points_list = []
+ for frame_idx in range(depth_map.shape[0]):
+ cur_world_points, _, _ = depth_to_world_coords_points(
+ depth_map[frame_idx].squeeze(-1), extrinsics_cam[frame_idx], intrinsics_cam[frame_idx]
+ )
+ world_points_list.append(cur_world_points)
+ world_points_array = np.stack(world_points_list, axis=0)
+
+ return world_points_array
+
+
+def depth_to_world_coords_points(
+ depth_map: np.ndarray,
+ extrinsic: np.ndarray,
+ intrinsic: np.ndarray,
+ eps=1e-8,
+) :
+ """
+ Convert a depth map to world coordinates.
+
+ Args:
+ depth_map (np.ndarray): Depth map of shape (H, W).
+ intrinsic (np.ndarray): Camera intrinsic matrix of shape (3, 3).
+ extrinsic (np.ndarray): Camera extrinsic matrix of shape (3, 4). OpenCV camera coordinate convention, cam from world.
+
+ Returns:
+ tuple[np.ndarray, np.ndarray]: World coordinates (H, W, 3) and valid depth mask (H, W).
+ """
+ if depth_map is None:
+ return None, None, None
+
+ # Valid depth mask
+ point_mask = depth_map > eps
+
+ # Convert depth map to camera coordinates
+ cam_coords_points = depth_to_cam_coords_points(depth_map, intrinsic)
+
+ # Multiply with the inverse of extrinsic matrix to transform to world coordinates
+ # extrinsic_inv is 4x4 (note closed_form_inverse_OpenCV is batched, the output is (N, 4, 4))
+ cam_to_world_extrinsic = closed_form_inverse_se3(extrinsic[None])[0]
+
+ R_cam_to_world = cam_to_world_extrinsic[:3, :3]
+ t_cam_to_world = cam_to_world_extrinsic[:3, 3]
+
+ # Apply the rotation and translation to the camera coordinates
+ world_coords_points = np.dot(cam_coords_points, R_cam_to_world.T) + t_cam_to_world # HxWx3, 3x3 -> HxWx3
+ # world_coords_points = np.einsum("ij,hwj->hwi", R_cam_to_world, cam_coords_points) + t_cam_to_world
+
+ return world_coords_points, cam_coords_points, point_mask
+
+
+def depth_to_cam_coords_points(depth_map: np.ndarray, intrinsic: np.ndarray) :
+ """
+ Convert a depth map to camera coordinates.
+
+ Args:
+ depth_map (np.ndarray): Depth map of shape (H, W).
+ intrinsic (np.ndarray): Camera intrinsic matrix of shape (3, 3).
+
+ Returns:
+ tuple[np.ndarray, np.ndarray]: Camera coordinates (H, W, 3)
+ """
+ H, W = depth_map.shape
+ assert intrinsic.shape == (3, 3), "Intrinsic matrix must be 3x3"
+ assert intrinsic[0, 1] == 0 and intrinsic[1, 0] == 0, "Intrinsic matrix must have zero skew"
+
+ # Intrinsic parameters
+ fu, fv = intrinsic[0, 0], intrinsic[1, 1]
+ cu, cv = intrinsic[0, 2], intrinsic[1, 2]
+
+ # Generate grid of pixel coordinates
+ u, v = np.meshgrid(np.arange(W), np.arange(H))
+
+ # Unproject to camera coordinates
+ x_cam = (u - cu) * depth_map / fu
+ y_cam = (v - cv) * depth_map / fv
+ z_cam = depth_map
+
+ # Stack to form camera coordinates
+ cam_coords = np.stack((x_cam, y_cam, z_cam), axis=-1).astype(np.float32)
+
+ return cam_coords
+
+
+def closed_form_inverse_se3(se3, R=None, T=None):
+ """
+ Compute the inverse of each 4x4 (or 3x4) SE3 matrix in a batch.
+
+ If `R` and `T` are provided, they must correspond to the rotation and translation
+ components of `se3`. Otherwise, they will be extracted from `se3`.
+
+ Args:
+ se3: Nx4x4 or Nx3x4 array or tensor of SE3 matrices.
+ R (optional): Nx3x3 array or tensor of rotation matrices.
+ T (optional): Nx3x1 array or tensor of translation vectors.
+
+ Returns:
+ Inverted SE3 matrices with the same type and device as `se3`.
+
+ Shapes:
+ se3: (N, 4, 4)
+ R: (N, 3, 3)
+ T: (N, 3, 1)
+ """
+ # Check if se3 is a numpy array or a torch tensor
+ is_numpy = isinstance(se3, np.ndarray)
+
+ # Validate shapes
+ if se3.shape[-2:] != (4, 4) and se3.shape[-2:] != (3, 4):
+ raise ValueError(f"se3 must be of shape (N,4,4), got {se3.shape}.")
+
+ # Extract R and T if not provided
+ if R is None:
+ R = se3[:, :3, :3] # (N,3,3)
+ if T is None:
+ T = se3[:, :3, 3:] # (N,3,1)
+
+ # Transpose R
+ if is_numpy:
+ # Compute the transpose of the rotation for NumPy
+ R_transposed = np.transpose(R, (0, 2, 1))
+ # -R^T t for NumPy
+ top_right = -np.matmul(R_transposed, T)
+ inverted_matrix = np.tile(np.eye(4), (len(R), 1, 1))
+ else:
+ R_transposed = R.transpose(1, 2) # (N,3,3)
+ top_right = -torch.bmm(R_transposed, T) # (N,3,1)
+ inverted_matrix = torch.eye(4, 4)[None].repeat(len(R), 1, 1)
+ inverted_matrix = inverted_matrix.to(R.dtype).to(R.device)
+
+ inverted_matrix[:, :3, :3] = R_transposed
+ inverted_matrix[:, :3, 3:] = top_right
+
+ return inverted_matrix
diff --git a/src/models/aggregator/utils/load_fn.py b/src/models/aggregator/utils/load_fn.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c5676a328655efa6dffb03b51792bc0914a1612
--- /dev/null
+++ b/src/models/aggregator/utils/load_fn.py
@@ -0,0 +1,146 @@
+# 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.
+
+import torch
+from PIL import Image
+from torchvision import transforms as TF
+
+
+def load_and_preprocess_images(image_path_list, mode="crop"):
+ """
+ A quick start function to load and preprocess images for model input.
+ This assumes the images should have the same shape for easier batching, but our model can also work well with different shapes.
+
+ Args:
+ image_path_list (list): List of paths to image files
+ mode (str, optional): Preprocessing mode, either "crop" or "pad".
+ - "crop" (default): Sets width to 518px and center crops height if needed.
+ - "pad": Preserves all pixels by making the largest dimension 518px
+ and padding the smaller dimension to reach a square shape.
+
+ Returns:
+ torch.Tensor: Batched tensor of preprocessed images with shape (N, 3, H, W)
+
+ Raises:
+ ValueError: If the input list is empty or if mode is invalid
+
+ Notes:
+ - Images with different dimensions will be padded with white (value=1.0)
+ - A warning is printed when images have different shapes
+ - When mode="crop": The function ensures width=518px while maintaining aspect ratio
+ and height is center-cropped if larger than 518px
+ - When mode="pad": The function ensures the largest dimension is 518px while maintaining aspect ratio
+ and the smaller dimension is padded to reach a square shape (518x518)
+ - Dimensions are adjusted to be divisible by 14 for compatibility with model requirements
+ """
+ # Check for empty list
+ if len(image_path_list) == 0:
+ raise ValueError("At least 1 image is required")
+
+ # Validate mode
+ if mode not in ["crop", "pad"]:
+ raise ValueError("Mode must be either 'crop' or 'pad'")
+
+ images = []
+ shapes = set()
+ to_tensor = TF.ToTensor()
+ target_size = 518
+
+ # First process all images and collect their shapes
+ for image_path in image_path_list:
+
+ # Open image
+ img = Image.open(image_path)
+
+ # If there's an alpha channel, blend onto white background:
+ if img.mode == "RGBA":
+ # Create white background
+ background = Image.new("RGBA", img.size, (255, 255, 255, 255))
+ # Alpha composite onto the white background
+ img = Image.alpha_composite(background, img)
+
+ # Now convert to "RGB" (this step assigns white for transparent areas)
+ img = img.convert("RGB")
+
+ width, height = img.size
+
+ if mode == "pad":
+ # Make the largest dimension 518px while maintaining aspect ratio
+ if width >= height:
+ new_width = target_size
+ new_height = round(height * (new_width / width) / 14) * 14 # Make divisible by 14
+ else:
+ new_height = target_size
+ new_width = round(width * (new_height / height) / 14) * 14 # Make divisible by 14
+ else: # mode == "crop"
+ # Original behavior: set width to 518px
+ new_width = target_size
+ # Calculate height maintaining aspect ratio, divisible by 14
+ new_height = round(height * (new_width / width) / 14) * 14
+
+ # Resize with new dimensions (width, height)
+ img = img.resize((new_width, new_height), Image.Resampling.BICUBIC)
+ img = to_tensor(img) # Convert to tensor (0, 1)
+
+ # Center crop height if it's larger than 518 (only in crop mode)
+ if mode == "crop" and new_height > target_size:
+ start_y = (new_height - target_size) // 2
+ img = img[:, start_y : start_y + target_size, :]
+
+ # For pad mode, pad to make a square of target_size x target_size
+ if mode == "pad":
+ h_padding = target_size - img.shape[1]
+ w_padding = target_size - img.shape[2]
+
+ if h_padding > 0 or w_padding > 0:
+ pad_top = h_padding // 2
+ pad_bottom = h_padding - pad_top
+ pad_left = w_padding // 2
+ pad_right = w_padding - pad_left
+
+ # Pad with white (value=1.0)
+ img = torch.nn.functional.pad(
+ img, (pad_left, pad_right, pad_top, pad_bottom), mode="constant", value=1.0
+ )
+
+ shapes.add((img.shape[1], img.shape[2]))
+ images.append(img)
+
+ # Check if we have different shapes
+ # In theory our model can also work well with different shapes
+ if len(shapes) > 1:
+ print(f"Warning: Found images with different shapes: {shapes}")
+ # Find maximum dimensions
+ max_height = max(shape[0] for shape in shapes)
+ max_width = max(shape[1] for shape in shapes)
+
+ # Pad images if necessary
+ padded_images = []
+ for img in images:
+ h_padding = max_height - img.shape[1]
+ w_padding = max_width - img.shape[2]
+
+ if h_padding > 0 or w_padding > 0:
+ pad_top = h_padding // 2
+ pad_bottom = h_padding - pad_top
+ pad_left = w_padding // 2
+ pad_right = w_padding - pad_left
+
+ img = torch.nn.functional.pad(
+ img, (pad_left, pad_right, pad_top, pad_bottom), mode="constant", value=1.0
+ )
+ padded_images.append(img)
+ images = padded_images
+
+ images = torch.stack(images) # concatenate images N,3,H,W
+
+ # Ensure correct shape when single image
+ if len(image_path_list) == 1:
+ # Verify shape is (1, C, H, W)
+ if images.dim() == 3:
+ images = images.unsqueeze(0)
+
+ return images
diff --git a/src/models/aggregator/utils/pose_enc.py b/src/models/aggregator/utils/pose_enc.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f98b0878cb13451b8cdb80074349cbf2644c5fa
--- /dev/null
+++ b/src/models/aggregator/utils/pose_enc.py
@@ -0,0 +1,130 @@
+# 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.
+
+import torch
+from .rotation import quat_to_mat, mat_to_quat
+
+
+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)
+ fx = (W / 2.0) / torch.tan(fov_w / 2.0)
+ intrinsics = torch.zeros(pose_encoding.shape[:2] + (3, 3), device=pose_encoding.device)
+ 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
diff --git a/src/models/aggregator/utils/rotation.py b/src/models/aggregator/utils/rotation.py
new file mode 100644
index 0000000000000000000000000000000000000000..657583e6915437c824c192d51939990b589a14fa
--- /dev/null
+++ b/src/models/aggregator/utils/rotation.py
@@ -0,0 +1,138 @@
+# 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 from PyTorch3D, https://github.com/facebookresearch/pytorch3d
+
+import torch
+import numpy as np
+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)
diff --git a/src/models/aggregator/utils/visual_track.py b/src/models/aggregator/utils/visual_track.py
new file mode 100644
index 0000000000000000000000000000000000000000..796c114ccba00b5f7850e04b9444a6cd5c44b154
--- /dev/null
+++ b/src/models/aggregator/utils/visual_track.py
@@ -0,0 +1,239 @@
+# 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.
+
+import cv2
+import torch
+import numpy as np
+import os
+
+
+def color_from_xy(x, y, W, H, cmap_name="hsv"):
+ """
+ Map (x, y) -> color in (R, G, B).
+ 1) Normalize x,y to [0,1].
+ 2) Combine them into a single scalar c in [0,1].
+ 3) Use matplotlib's colormap to convert c -> (R,G,B).
+
+ You can customize step 2, e.g., c = (x + y)/2, or some function of (x, y).
+ """
+ import matplotlib.cm
+ import matplotlib.colors
+
+ x_norm = x / max(W - 1, 1)
+ y_norm = y / max(H - 1, 1)
+ # Simple combination:
+ c = (x_norm + y_norm) / 2.0
+
+ cmap = matplotlib.cm.get_cmap(cmap_name)
+ # cmap(c) -> (r,g,b,a) in [0,1]
+ rgba = cmap(c)
+ r, g, b = rgba[0], rgba[1], rgba[2]
+ return (r, g, b) # in [0,1], RGB order
+
+
+def get_track_colors_by_position(tracks_b, vis_mask_b=None, image_width=None, image_height=None, cmap_name="hsv"):
+ """
+ Given all tracks in one sample (b), compute a (N,3) array of RGB color values
+ in [0,255]. The color is determined by the (x,y) position in the first
+ visible frame for each track.
+
+ Args:
+ tracks_b: Tensor of shape (S, N, 2). (x,y) for each track in each frame.
+ vis_mask_b: (S, N) boolean mask; if None, assume all are visible.
+ image_width, image_height: used for normalizing (x, y).
+ cmap_name: for matplotlib (e.g., 'hsv', 'rainbow', 'jet').
+
+ Returns:
+ track_colors: np.ndarray of shape (N, 3), each row is (R,G,B) in [0,255].
+ """
+ S, N, _ = tracks_b.shape
+ track_colors = np.zeros((N, 3), dtype=np.uint8)
+
+ if vis_mask_b is None:
+ # treat all as visible
+ vis_mask_b = torch.ones(S, N, dtype=torch.bool, device=tracks_b.device)
+
+ for i in range(N):
+ # Find first visible frame for track i
+ visible_frames = torch.where(vis_mask_b[:, i])[0]
+ if len(visible_frames) == 0:
+ # track is never visible; just assign black or something
+ track_colors[i] = (0, 0, 0)
+ continue
+
+ first_s = int(visible_frames[0].item())
+ # use that frame's (x,y)
+ x, y = tracks_b[first_s, i].tolist()
+
+ # map (x,y) -> (R,G,B) in [0,1]
+ r, g, b = color_from_xy(x, y, W=image_width, H=image_height, cmap_name=cmap_name)
+ # scale to [0,255]
+ r, g, b = int(r * 255), int(g * 255), int(b * 255)
+ track_colors[i] = (r, g, b)
+
+ return track_colors
+
+
+def visualize_tracks_on_images(
+ images,
+ tracks,
+ track_vis_mask=None,
+ out_dir="track_visuals_concat_by_xy",
+ image_format="CHW", # "CHW" or "HWC"
+ normalize_mode="[0,1]",
+ cmap_name="hsv", # e.g. "hsv", "rainbow", "jet"
+ frames_per_row=4, # New parameter for grid layout
+ save_grid=True, # Flag to control whether to save the grid image
+):
+ """
+ Visualizes frames in a grid layout with specified frames per row.
+ Each track's color is determined by its (x,y) position
+ in the first visible frame (or frame 0 if always visible).
+ Finally convert the BGR result to RGB before saving.
+ Also saves each individual frame as a separate PNG file.
+
+ Args:
+ images: torch.Tensor (S, 3, H, W) if CHW or (S, H, W, 3) if HWC.
+ tracks: torch.Tensor (S, N, 2), last dim = (x, y).
+ track_vis_mask: torch.Tensor (S, N) or None.
+ out_dir: folder to save visualizations.
+ image_format: "CHW" or "HWC".
+ normalize_mode: "[0,1]", "[-1,1]", or None for direct raw -> 0..255
+ cmap_name: a matplotlib colormap name for color_from_xy.
+ frames_per_row: number of frames to display in each row of the grid.
+ save_grid: whether to save all frames in one grid image.
+
+ Returns:
+ None (saves images in out_dir).
+ """
+
+ if len(tracks.shape) == 4:
+ tracks = tracks.squeeze(0)
+ images = images.squeeze(0)
+ if track_vis_mask is not None:
+ track_vis_mask = track_vis_mask.squeeze(0)
+
+ import matplotlib
+
+ matplotlib.use("Agg") # for non-interactive (optional)
+
+ os.makedirs(out_dir, exist_ok=True)
+
+ S = images.shape[0]
+ _, N, _ = tracks.shape # (S, N, 2)
+
+ # Move to CPU
+ images = images.cpu().clone()
+ tracks = tracks.cpu().clone()
+ if track_vis_mask is not None:
+ track_vis_mask = track_vis_mask.cpu().clone()
+
+ # Infer H, W from images shape
+ if image_format == "CHW":
+ # e.g. images[s].shape = (3, H, W)
+ H, W = images.shape[2], images.shape[3]
+ else:
+ # e.g. images[s].shape = (H, W, 3)
+ H, W = images.shape[1], images.shape[2]
+
+ # Pre-compute the color for each track i based on first visible position
+ track_colors_rgb = get_track_colors_by_position(
+ tracks, # shape (S, N, 2)
+ vis_mask_b=track_vis_mask if track_vis_mask is not None else None,
+ image_width=W,
+ image_height=H,
+ cmap_name=cmap_name,
+ )
+
+ # We'll accumulate each frame's drawn image in a list
+ frame_images = []
+
+ for s in range(S):
+ # shape => either (3, H, W) or (H, W, 3)
+ img = images[s]
+
+ # Convert to (H, W, 3)
+ if image_format == "CHW":
+ img = img.permute(1, 2, 0) # (H, W, 3)
+ # else "HWC", do nothing
+
+ img = img.numpy().astype(np.float32)
+
+ # Scale to [0,255] if needed
+ if normalize_mode == "[0,1]":
+ img = np.clip(img, 0, 1) * 255.0
+ elif normalize_mode == "[-1,1]":
+ img = (img + 1.0) * 0.5 * 255.0
+ img = np.clip(img, 0, 255.0)
+ # else no normalization
+
+ # Convert to uint8
+ img = img.astype(np.uint8)
+
+ # For drawing in OpenCV, convert to BGR
+ img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+
+ # Draw each visible track
+ cur_tracks = tracks[s] # shape (N, 2)
+ if track_vis_mask is not None:
+ valid_indices = torch.where(track_vis_mask[s])[0]
+ else:
+ valid_indices = range(N)
+
+ cur_tracks_np = cur_tracks.numpy()
+ for i in valid_indices:
+ x, y = cur_tracks_np[i]
+ pt = (int(round(x)), int(round(y)))
+
+ # track_colors_rgb[i] is (R,G,B). For OpenCV circle, we need BGR
+ R, G, B = track_colors_rgb[i]
+ color_bgr = (int(B), int(G), int(R))
+ cv2.circle(img_bgr, pt, radius=3, color=color_bgr, thickness=-1)
+
+ # Convert back to RGB for consistent final saving:
+ img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+
+ # Save individual frame
+ frame_path = os.path.join(out_dir, f"frame_{s:04d}.png")
+ # Convert to BGR for OpenCV imwrite
+ frame_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)
+ cv2.imwrite(frame_path, frame_bgr)
+
+ frame_images.append(img_rgb)
+
+ # Only create and save the grid image if save_grid is True
+ if save_grid:
+ # Calculate grid dimensions
+ num_rows = (S + frames_per_row - 1) // frames_per_row # Ceiling division
+
+ # Create a grid of images
+ grid_img = None
+ for row in range(num_rows):
+ start_idx = row * frames_per_row
+ end_idx = min(start_idx + frames_per_row, S)
+
+ # Concatenate this row horizontally
+ row_img = np.concatenate(frame_images[start_idx:end_idx], axis=1)
+
+ # If this row has fewer than frames_per_row images, pad with black
+ if end_idx - start_idx < frames_per_row:
+ padding_width = (frames_per_row - (end_idx - start_idx)) * W
+ padding = np.zeros((H, padding_width, 3), dtype=np.uint8)
+ row_img = np.concatenate([row_img, padding], axis=1)
+
+ # Add this row to the grid
+ if grid_img is None:
+ grid_img = row_img
+ else:
+ grid_img = np.concatenate([grid_img, row_img], axis=0)
+
+ out_path = os.path.join(out_dir, "tracks_grid.png")
+ # Convert back to BGR for OpenCV imwrite
+ grid_img_bgr = cv2.cvtColor(grid_img, cv2.COLOR_RGB2BGR)
+ cv2.imwrite(out_path, grid_img_bgr)
+ print(f"[INFO] Saved color-by-XY track visualization grid -> {out_path}")
+
+ print(f"[INFO] Saved {S} individual frames to {out_dir}/frame_*.png")
diff --git a/src/models/module/__pycache__/utils.cpython-310.pyc b/src/models/module/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a44f4140bba3ec9ed8980affec7644014e858b59
Binary files /dev/null and b/src/models/module/__pycache__/utils.cpython-310.pyc differ
diff --git a/src/models/module/utils.py b/src/models/module/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..c001a93c9d8b560775ec2bf7d0a578eb432fbdfe
--- /dev/null
+++ b/src/models/module/utils.py
@@ -0,0 +1,610 @@
+from typing import Tuple
+import torch
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..utils.model_utils import *
+from ..utils import transformer
+from ..utils.ind2sub import *
+from ..utils.decompose_tensors import *
+from ..utils.utils import *
+from einops import rearrange
+from ..aggregator import Aggregator
+import pywt
+
+def create_wavelet_filter(wave, in_size, out_size, type=torch.float):
+ w = pywt.Wavelet(wave)
+ dec_hi = torch.tensor(w.dec_hi[::-1], dtype=type)
+ dec_lo = torch.tensor(w.dec_lo[::-1], dtype=type)
+ dec_filters = torch.stack([dec_lo.unsqueeze(0) * dec_lo.unsqueeze(1),
+ dec_lo.unsqueeze(0) * dec_hi.unsqueeze(1),
+ dec_hi.unsqueeze(0) * dec_lo.unsqueeze(1),
+ dec_hi.unsqueeze(0) * dec_hi.unsqueeze(1)], dim=0)
+
+ dec_filters = dec_filters[:, None].repeat(in_size, 1, 1, 1)
+
+ rec_hi = torch.tensor(w.rec_hi[::-1], dtype=type).flip(dims=[0])
+ rec_lo = torch.tensor(w.rec_lo[::-1], dtype=type).flip(dims=[0])
+ rec_filters = torch.stack([rec_lo.unsqueeze(0) * rec_lo.unsqueeze(1),
+ rec_lo.unsqueeze(0) * rec_hi.unsqueeze(1),
+ rec_hi.unsqueeze(0) * rec_lo.unsqueeze(1),
+ rec_hi.unsqueeze(0) * rec_hi.unsqueeze(1)], dim=0)
+
+ rec_filters = rec_filters[:, None].repeat(out_size, 1, 1, 1)
+
+ return dec_filters, rec_filters
+
+def wavelet_transform(x, filters):
+ b, c, h, w = x.shape
+ pad = (filters.shape[2] // 2 - 1, filters.shape[3] // 2 - 1)
+ x = F.conv2d(x, filters, stride=2, groups=c, padding=pad)
+ x = x.reshape(b, c, 4, h // 2, w // 2)
+ return x
+
+
+def inverse_wavelet_transform(x, filters):
+ b, c, _, h_half, w_half = x.shape
+ pad = (filters.shape[2] // 2 - 1, filters.shape[3] // 2 - 1)
+ x = x.reshape(b, c * 4, h_half, w_half)
+ x = F.conv_transpose2d(x, filters, stride=2, groups=c, padding=pad)
+ return x
+
+class ResidualConvUnit(nn.Module):
+ """Residual convolution module."""
+
+ def __init__(self, features, activation, bn, groups=1):
+ """Init.
+
+ Args:
+ features (int): number of features
+ """
+ super().__init__()
+
+ self.bn = bn
+ self.groups = groups
+ self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+ self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True, groups=self.groups)
+
+ self.norm1 = None
+ self.norm2 = None
+
+ self.activation = activation
+ self.skip_add = nn.quantized.FloatFunctional()
+
+ def forward(self, x):
+ """Forward pass.
+
+ Args:
+ x (tensor): input
+
+ Returns:
+ tensor: output
+ """
+
+ out = self.activation(x)
+ out = self.conv1(out)
+ if self.norm1 is not None:
+ out = self.norm1(out)
+
+ out = self.activation(out)
+ out = self.conv2(out)
+ if self.norm2 is not None:
+ out = self.norm2(out)
+
+ return self.skip_add.add(out, x)
+
+
+class FeatureFusionBlock(nn.Module):
+ """Feature fusion block."""
+
+ def __init__(
+ self,
+ features,
+ activation,
+ deconv=False,
+ bn=False,
+ expand=False,
+ align_corners=True,
+ size=None,
+ has_residual=True,
+ groups=1,
+ ):
+ """Init.
+
+ Args:
+ features (int): number of features
+ """
+ super(FeatureFusionBlock, self).__init__()
+
+ self.deconv = deconv
+ self.align_corners = align_corners
+ self.groups = groups
+ self.expand = expand
+ out_features = features
+ if self.expand == True:
+ out_features = features // 2
+
+ self.out_conv = nn.Conv2d(
+ features, out_features, kernel_size=1, stride=1, padding=0, bias=True, groups=self.groups
+ )
+
+ if has_residual:
+ self.resConfUnit1 = ResidualConvUnit(features, activation, bn, groups=self.groups)
+
+ self.has_residual = has_residual
+ self.resConfUnit2 = ResidualConvUnit(features, activation, bn, groups=self.groups)
+
+ self.skip_add = nn.quantized.FloatFunctional()
+ self.size = size
+
+ def forward(self, *xs, size=None):
+ """Forward pass.
+
+ Returns:
+ tensor: output
+ """
+ output = xs[0]
+
+ if self.has_residual:
+ res = self.resConfUnit1(xs[1])
+ output = self.skip_add.add(output, res)
+
+ output = self.resConfUnit2(output)
+
+ if (size is None) and (self.size is None):
+ modifier = {"scale_factor": 2}
+ elif size is None:
+ modifier = {"size": self.size}
+ else:
+ modifier = {"size": size}
+
+ output = custom_interpolate(output.float(), **modifier, mode="bilinear", align_corners=self.align_corners).to(torch.bfloat16)
+ output = self.out_conv(output)
+
+ return output
+
+
+def custom_interpolate(
+ x: torch.Tensor,
+ size: Tuple[int, int] = None,
+ scale_factor: float = None,
+ mode: str = "bilinear",
+ align_corners: bool = True,
+) -> torch.Tensor:
+ """
+ Custom interpolate to avoid INT_MAX issues in nn.functional.interpolate.
+ """
+ if size is None:
+ size = (int(x.shape[-2] * scale_factor), int(x.shape[-1] * scale_factor))
+
+ INT_MAX = 1610612736
+
+ input_elements = size[0] * size[1] * x.shape[0] * x.shape[1]
+
+ if input_elements > INT_MAX:
+ chunks = torch.chunk(x, chunks=(input_elements // INT_MAX) + 1, dim=0)
+ interpolated_chunks = [
+ nn.functional.interpolate(chunk, size=size, mode=mode, align_corners=align_corners) for chunk in chunks
+ ]
+ x = torch.cat(interpolated_chunks, dim=0)
+ return x.contiguous()
+ else:
+ return nn.functional.interpolate(x, size=size, mode=mode, align_corners=align_corners)
+
+def _make_scratch(in_shape, out_shape: int, groups: int = 1, expand: bool = False) -> nn.Module:
+ """
+
+ """
+ scratch = nn.Module()
+
+
+ activation_function = nn.LeakyReLU
+
+
+ out_shape1 = out_shape
+ out_shape2 = out_shape
+ out_shape3 = out_shape
+ if len(in_shape) >= 4:
+ out_shape4 = out_shape
+
+ if expand:
+ out_shape1 = out_shape
+ out_shape2 = out_shape * 2
+ out_shape3 = out_shape * 4
+ if len(in_shape) >= 4:
+ out_shape4 = out_shape * 8
+
+
+ scratch.layer1_rn = nn.Sequential(
+ nn.Conv2d(
+ in_shape[0], out_shape1, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ ),
+ activation_function()
+ )
+ scratch.layer2_rn = nn.Sequential(
+ nn.Conv2d(
+ in_shape[1], out_shape2, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ ),
+ activation_function()
+ )
+ scratch.layer3_rn = nn.Sequential(
+ nn.Conv2d(
+ in_shape[2], out_shape3, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ ),
+ activation_function()
+ )
+ if len(in_shape) >= 4:
+
+ scratch.layer4_rn = nn.Sequential(
+ nn.Conv2d(
+ in_shape[3], out_shape4, kernel_size=3, stride=1, padding=1, bias=False, groups=groups
+ ),
+ activation_function()
+ )
+
+ return scratch
+
+
+def _make_fusion_block(features: int, size: int = None, has_residual: bool = True, groups: int = 1) -> nn.Module:
+ return FeatureFusionBlock(
+ features,
+ nn.LeakyReLU(inplace=False),
+ deconv=False,
+ bn=False,
+ expand=False,
+ align_corners=True,
+ size=size,
+ has_residual=has_residual,
+ groups=groups,
+ )
+
+
+class ImageFeatureExtractor(nn.Module):
+ def __init__(self, depth=4,img_size=256, patch_size=8, embed_dim=384):
+ super(ImageFeatureExtractor, self).__init__()
+ self.aggregator = Aggregator(img_size, patch_size, embed_dim,depth=depth,patch_embed="dinov2_vits14_reg")
+
+
+ def forward(self, x, nImgArray):
+
+
+ feat_list, normal_patch_start_idx = self.aggregator(x)
+
+ return torch.stack(feat_list,dim=0).permute(1,2,0,3,4).flatten(0,1),normal_patch_start_idx
+
+
+class ImageFeatureFusion(nn.Module):
+ def __init__(self,
+ in_channels,
+ use_efficient_attention=False,
+ out_channels = [256, 512, 1024, 1024],
+ features = 256,
+ ):
+ super(ImageFeatureFusion, self).__init__()
+ _, self.iwt_filter = create_wavelet_filter('db1', 384, 384, torch.bfloat16)
+ self.pixel_shuffle = nn.PixelShuffle(2)
+ self.norm = nn.LayerNorm(in_channels)
+ self.projects = nn.ModuleList(
+ [
+ nn.Sequential(
+
+ nn.Conv2d(
+ in_channels=in_channels // 4,
+ out_channels=oc,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ bias=True
+ ),
+ nn.LeakyReLU()
+
+ )
+ for oc in out_channels
+ ]
+ )
+
+ self.resize_layers = nn.ModuleList(
+ [
+ nn.Sequential(
+ nn.ConvTranspose2d(
+ in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=2, stride=2, padding=0
+ ),
+ nn.LeakyReLU(),
+ nn.ConvTranspose2d(
+ in_channels=out_channels[0], out_channels=out_channels[0], kernel_size=2, stride=2, padding=0
+ ),
+ nn.LeakyReLU()
+ ),
+ nn.Sequential(
+ nn.ConvTranspose2d(
+ in_channels=out_channels[1], out_channels=out_channels[1], kernel_size=2, stride=2, padding=0
+ ),
+ nn.LeakyReLU()
+ ),
+ nn.Sequential(
+ nn.Conv2d(
+ in_channels=out_channels[2], out_channels=out_channels[2], kernel_size=1, stride=1, padding=0
+ ),
+ nn.LeakyReLU()
+ ),
+ nn.Sequential(
+ nn.Conv2d(
+ in_channels=out_channels[3], out_channels=out_channels[3], kernel_size=2, stride=2, padding=0
+ ),
+ nn.LeakyReLU()
+ )
+
+
+
+
+
+
+ ]
+)
+
+
+ self.scratch = _make_scratch(
+ out_channels,
+ features,
+ expand=False,
+ )
+
+
+ self.scratch.stem_transpose = None
+ self.scratch.refinenet1 = _make_fusion_block(features)
+ self.scratch.refinenet2 = _make_fusion_block(features)
+ self.scratch.refinenet3 = _make_fusion_block(features)
+ self.scratch.refinenet4 = _make_fusion_block(features, has_residual=False)
+
+ head_features_1 = features
+
+ self.scratch.output_conv1 = nn.Conv2d(
+ head_features_1, head_features_1 , kernel_size=3, stride=2, padding=1
+ )
+
+
+
+ def _apply_pos_embed(self, x: torch.Tensor, W: int, H: int, ratio: float = 0.1) -> torch.Tensor:
+ """
+ Apply positional embedding to tensor x.
+ """
+ patch_w = x.shape[-1]
+ patch_h = x.shape[-2]
+ pos_embed = create_uv_grid(patch_w, patch_h, aspect_ratio=W / H, dtype=x.dtype, device=x.device)
+ pos_embed = position_grid_to_embed(pos_embed, x.shape[1])
+ pos_embed = pos_embed * ratio
+ pos_embed = pos_embed.permute(2, 0, 1)[None].expand(x.shape[0], -1, -1, -1)
+ return x + pos_embed
+
+ def scratch_forward(self, features) -> torch.Tensor:
+ """
+ Forward pass through the fusion blocks.
+
+ Args:
+ features (List[Tensor]): List of feature maps from different layers.
+
+ Returns:
+ Tensor: Fused feature map.
+ """
+ layer_1, layer_2, layer_3, layer_4 = features
+
+ layer_1_rn = self.scratch.layer1_rn(layer_1)
+ layer_2_rn = self.scratch.layer2_rn(layer_2)
+ layer_3_rn = self.scratch.layer3_rn(layer_3)
+ layer_4_rn = self.scratch.layer4_rn(layer_4)
+
+ out = self.scratch.refinenet4(layer_4_rn, size=layer_3_rn.shape[2:])
+ del layer_4_rn, layer_4
+
+ out = self.scratch.refinenet3(out, layer_3_rn, size=layer_2_rn.shape[2:])
+ del layer_3_rn, layer_3
+
+ out = self.scratch.refinenet2(out, layer_2_rn, size=layer_1_rn.shape[2:])
+ del layer_2_rn, layer_2
+
+ out = self.scratch.refinenet1(out, layer_1_rn)
+ del layer_1_rn, layer_1
+
+ out = self.scratch.output_conv1(out)
+ return out
+ def forward(self,
+ glc: torch.Tensor,
+ nImgArray: torch.Tensor,
+ chunk_size: int = 6
+ ) -> torch.Tensor:
+
+ B = glc.shape[0]
+
+ # 如果不需要分块（总批次大小小于或等于块大小），则直接调用核心实现
+ if chunk_size is None or chunk_size >= B:
+ return self._forward_impl(glc, nImgArray)
+
+ # 否则，进行分块处理
+ all_outputs = []
+ # 以 chunk_size 为步长进行循环
+ for start_idx in range(0, B, chunk_size):
+ # 计算当前块的结束索引
+ end_idx = min(start_idx + chunk_size, B)
+
+ # 从大的输入张量中切出当前要处理的小块
+ glc_chunk = glc[start_idx:end_idx]
+
+ # 注意：如果 nImgArray 也与批次相关，也需要进行切片
+ # nImgArray_chunk = nImgArray[start_idx:end_idx]
+
+ # 调用核心实现函数来处理这个小块
+ chunk_output = self._forward_impl(glc_chunk, nImgArray)
+
+ all_outputs.append(chunk_output)
+
+ # 将所有小块的处理结果，沿着批次维度（dim=0）重新拼接起来
+ final_output = torch.cat(all_outputs, dim=0)
+
+ return final_output
+
+ def _forward_impl(self, glc: torch.Tensor, nImgArray: torch.Tensor) -> torch.Tensor:
+ """
+ 这是核心实现方法，处理一个数据块（chunk）。
+ 这里的代码就是您提供的原始 forward 方法的主体。
+ """
+ self.iwt_filter = self.iwt_filter.to(glc.device)
+ B, layer_num, N, C = glc.shape # 这里的 B 现在是 chunk_size
+ out = []
+
+ for layer in range(layer_num):
+ x = glc[:, layer, :, :] # [B, N, C]
+ x = self.norm(x)
+ x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], int(N**0.5), int(N**0.5))) # [B,C,sqrt(N),sqrt(N)]
+
+ x = self.pixel_shuffle(x) # [B, C, H, W] -> [B, C/4, H*2, W*2]
+
+ x = self.projects[layer](x)
+ x = self._apply_pos_embed(x, 256, 256).to(torch.bfloat16)
+ x = self.resize_layers[layer](x)
+
+ out.append(x)
+
+ out = self.scratch_forward(out)
+ out = self._apply_pos_embed(out, 256, 256).to(torch.bfloat16) # [B, 256, 64, 64]
+
+ return out
+
+
+
+
+
+
+
+class ScaleInvariantSpatialLightImageEncoder(nn.Module):
+ def __init__(self, input_nc, depth=4, use_efficient_attention=False):
+ super(ScaleInvariantSpatialLightImageEncoder, self).__init__()
+ out_channels = (96, 192, 384, 768)
+ self.backbone = ImageFeatureExtractor(depth=depth)
+ self.fusion = ImageFeatureFusion(in_channels=1536, use_efficient_attention=use_efficient_attention)
+ self.feat_dim = 256
+ self.wt_filter, _ = create_wavelet_filter('db1', 3, 3, torch.bfloat16)
+ _, self.iwt_filter = create_wavelet_filter('db1', self.feat_dim, self.feat_dim, torch.bfloat16)
+
+
+ def forward(self, x, nImgArray, canonical_resolution):
+ N, C, H, W = x.shape
+ B = N//nImgArray[0]
+ mosaic_scale = H // canonical_resolution
+ K = mosaic_scale * mosaic_scale
+ self.wt_filter = self.wt_filter.to(x.device)
+ self.iwt_filter = self.iwt_filter.to(x.device)
+
+ """ (1a) resizing x to (Hc, Wc)"""
+ x_resized = F.interpolate(x.float(), size= (canonical_resolution, canonical_resolution), mode='bilinear', align_corners=True).to(torch.bfloat16)
+ x_resized = x_resized.view(len(nImgArray), int(nImgArray[0]), C, x_resized.shape[2], x_resized.shape[3])
+
+
+ """ (1b) decomposing x into K x K of (Hc, Wc) non-overlapped blocks (stride)"""
+ x_wt = wavelet_transform(x, self.wt_filter).permute(0, 2, 1, 3, 4)
+
+ x_wt = x_wt.reshape(B,nImgArray[0],K,3,canonical_resolution,canonical_resolution).flatten(1,2).flatten(0,1)
+ x_wt = x_wt.view(len(nImgArray), K * int(nImgArray[0]), C, x_wt.shape[2], x_wt.shape[3])
+
+ """ (2a) feature extraction """
+ aggregated_tokens_list, patch_start_idx = self.backbone(x_resized,nImgArray)
+ light_tokens_resized = aggregated_tokens_list[:,:,:patch_start_idx - 4,:]
+ light_tokens_resized = rearrange(light_tokens_resized,'(B f) layer num c -> B f layer num c',B = B)
+ x = self.fusion(aggregated_tokens_list[:,:,patch_start_idx:,:], nImgArray)
+ f_resized_grid = F.interpolate(x.reshape(N, self.feat_dim, canonical_resolution, canonical_resolution).float() , size= (H, W), mode='bilinear', align_corners=True).to(torch.bfloat16)
+
+ """ (2b) feature extraction (grid) """
+
+ aggregated_tokens_list, patch_start_idx = self.backbone(x_wt,nImgArray)
+ light_tokens_wt = aggregated_tokens_list[:,:,:patch_start_idx - 4,:]
+ light_tokens_wt = rearrange(light_tokens_wt,'(B f k) layer num c -> B f k layer num c',B = B, f=nImgArray[0])
+ light_tokens = torch.cat((light_tokens_resized.unsqueeze(2), light_tokens_wt), dim=2)
+
+
+ x = self.fusion(aggregated_tokens_list[:,:,patch_start_idx:,:], nImgArray)
+ x = rearrange(x, '(f k) c h w -> f c k h w ',k=4)
+ x = inverse_wavelet_transform(x, self.iwt_filter)
+
+ """ (3) upsample """
+ glc = (f_resized_grid + x)
+
+
+
+ return glc,light_tokens
+
+
+
+class GLC_Upsample(nn.Module):
+ def __init__(self, input_nc, num_enc_sab=1, dim_hidden=256, dim_feedforward=1024, use_efficient_attention=False):
+ super(GLC_Upsample, self).__init__()
+ self.comm = transformer.CommunicationBlock(input_nc, num_enc_sab = num_enc_sab, dim_hidden=dim_hidden, ln=True, dim_feedforward = dim_feedforward,use_efficient_attention=False)
+
+ def forward(self, x):
+ x = self.comm(x)
+ return x
+
+class GLC_Aggregation(nn.Module):
+ def __init__(self, input_nc, num_agg_transformer=2, dim_aggout=384, dim_feedforward=1024, use_efficient_attention=False):
+ super(GLC_Aggregation, self).__init__()
+ self.aggregation = transformer.AggregationBlock(dim_input = input_nc, num_enc_sab = num_agg_transformer, num_outputs = 1, dim_hidden=dim_aggout, dim_feedforward = dim_feedforward, num_heads=8, ln=True, attention_dropout=0.1, use_efficient_attention=use_efficient_attention)
+
+ def forward(self, x):
+ x = self.aggregation(x)
+ return x
+
+
+
+
+class Regressor(nn.Module):
+ def __init__(self, input_nc, num_enc_sab=1, use_efficient_attention=False, dim_feedforward=256, output='normal'):
+ super(Regressor, self).__init__()
+
+ self.comm = transformer.CommunicationBlock(input_nc, num_enc_sab = num_enc_sab, dim_hidden=input_nc, ln=True, dim_feedforward = dim_feedforward, use_efficient_attention=use_efficient_attention)
+ if output == 'normal':
+ self.prediction_normal = PredictionHead(input_nc, 3, confidence=True)
+ self.target = output
+
+ def forward(self, x, num_sample_set):
+ """Standard forward
+ INPUT: img [Num_Pix, F]
+ OUTPUT: [Num_Pix, 3]"""
+ if x.shape[0] % num_sample_set == 0:
+ x_ = x.reshape(-1, num_sample_set, x.shape[1])
+ x_ = self.comm(x_)
+ x = x_.reshape(-1, x.shape[1])
+ else:
+ ids = list(range(x.shape[0]))
+ num_split = len(ids) // num_sample_set
+ x_1 = x[:(num_split)*num_sample_set, :].reshape(-1, num_sample_set, x.shape[1])
+ x_1 = self.comm(x_1).reshape(-1, x.shape[1])
+ x_2 = x[(num_split)*num_sample_set:,:].reshape(1, -1, x.shape[1])
+ x_2 = self.comm(x_2).reshape(-1, x.shape[1])
+ x = torch.cat([x_1, x_2], dim=0)
+ if self.target == 'normal':
+ x_n, conf = self.prediction_normal(x.reshape(x.shape[0]//num_sample_set, num_sample_set, -1))
+ x_ = []
+ return x_n, x_, x, conf
+
+
+
+class PredictionHead(nn.Module):
+ def __init__(self, dim_input, dim_output, confidence=False):
+
+
+ super(PredictionHead, self).__init__()
+ modules_regression = []
+ modules_regression.append(nn.Linear(dim_input, dim_input//2))
+ modules_regression.append(nn.ReLU())
+ self.out_layer = nn.Linear(dim_input//2, dim_output)
+ if confidence:
+ self.confi_layer = nn.Linear(dim_input//2, 1)
+
+ self.regression = nn.Sequential(*modules_regression)
+
+ def forward(self, x):
+ h = self.regression(x)
+ ret = self.out_layer(h)
+ if hasattr(self, 'confi_layer'):
+ confidence = self.confi_layer(h)
+ else:
+ confidence = torch.zeros_like([ret.shape[0], 1])
+ return ret, torch.sigmoid(confidence)
diff --git a/src/models/utils/__init__.py b/src/models/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..46723d8b8b6ae0983b3922e00afa8842c26fb484
--- /dev/null
+++ b/src/models/utils/__init__.py
@@ -0,0 +1,7 @@
+from . import model_utils
+from . import transformer
+from . import ind2sub
+from . import decompose_tensors
+from . import utils
+from . import gauss_filter
+from . import compute_mae
\ No newline at end of file
diff --git a/src/models/utils/__pycache__/__init__.cpython-310.pyc b/src/models/utils/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2f5567fdf8d40f7757be377e37ef6cc5e41d57ad
Binary files /dev/null and b/src/models/utils/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/__init__.cpython-38.pyc b/src/models/utils/__pycache__/__init__.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0fe980d77e668eec1d9443aa789b62ad33748712
Binary files /dev/null and b/src/models/utils/__pycache__/__init__.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/compute_mae.cpython-310.pyc b/src/models/utils/__pycache__/compute_mae.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..77e0835e9556108eeca5fd2223d703a9cd1f3a94
Binary files /dev/null and b/src/models/utils/__pycache__/compute_mae.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/convnext.cpython-310.pyc b/src/models/utils/__pycache__/convnext.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..8d6b3911b129ee7e09cd898414a131e229c33087
Binary files /dev/null and b/src/models/utils/__pycache__/convnext.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/convnext.cpython-38.pyc b/src/models/utils/__pycache__/convnext.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..f598022df6429340c3a76bc0cda5cbd6d9e2a41a
Binary files /dev/null and b/src/models/utils/__pycache__/convnext.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/decompose_tensors.cpython-310.pyc b/src/models/utils/__pycache__/decompose_tensors.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..2dbaecda46091ba4be8f6e783cb73225110e460a
Binary files /dev/null and b/src/models/utils/__pycache__/decompose_tensors.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/decompose_tensors.cpython-38.pyc b/src/models/utils/__pycache__/decompose_tensors.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..79fd25c72c35b052e1000e66c4c2e21179cd95cb
Binary files /dev/null and b/src/models/utils/__pycache__/decompose_tensors.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/fpn.cpython-310.pyc b/src/models/utils/__pycache__/fpn.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..73d7c6a2a55899509e78c5f26640b01b34bf218e
Binary files /dev/null and b/src/models/utils/__pycache__/fpn.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/fpn.cpython-38.pyc b/src/models/utils/__pycache__/fpn.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b35bca92209b4a457e8691c6a8d96b48a138a52a
Binary files /dev/null and b/src/models/utils/__pycache__/fpn.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/gauss_filter.cpython-310.pyc b/src/models/utils/__pycache__/gauss_filter.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..95238b2d0867523f15d0a13ed30246718a792f3f
Binary files /dev/null and b/src/models/utils/__pycache__/gauss_filter.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/gauss_filter.cpython-38.pyc b/src/models/utils/__pycache__/gauss_filter.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..3336cd35c935e3434fa803ea40cc44dca8ca06d5
Binary files /dev/null and b/src/models/utils/__pycache__/gauss_filter.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/ind2sub.cpython-310.pyc b/src/models/utils/__pycache__/ind2sub.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b99f9a37042959f135e58335932882960331e8d7
Binary files /dev/null and b/src/models/utils/__pycache__/ind2sub.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/ind2sub.cpython-38.pyc b/src/models/utils/__pycache__/ind2sub.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..db47ddfe3ceea17c03d977960dde9650da56d09c
Binary files /dev/null and b/src/models/utils/__pycache__/ind2sub.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/model_utils.cpython-310.pyc b/src/models/utils/__pycache__/model_utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ff68ffebe558dc2918b6c87e216effdfdd6e7ad2
Binary files /dev/null and b/src/models/utils/__pycache__/model_utils.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/model_utils.cpython-38.pyc b/src/models/utils/__pycache__/model_utils.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..39345e3032bda7e470984e634f3f071e396ce50b
Binary files /dev/null and b/src/models/utils/__pycache__/model_utils.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/show_tensor.cpython-310.pyc b/src/models/utils/__pycache__/show_tensor.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..893d45ad2caa553dfec0e16360e42464c70e5b5f
Binary files /dev/null and b/src/models/utils/__pycache__/show_tensor.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/show_tensor.cpython-38.pyc b/src/models/utils/__pycache__/show_tensor.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..048182b3f0fe2fe5dc872cc1430b6c01ee17d2f7
Binary files /dev/null and b/src/models/utils/__pycache__/show_tensor.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/transformer.cpython-310.pyc b/src/models/utils/__pycache__/transformer.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bd9d53022344a5520ba533b4409262cd5a9201d6
Binary files /dev/null and b/src/models/utils/__pycache__/transformer.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/transformer.cpython-38.pyc b/src/models/utils/__pycache__/transformer.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..fd7acdbd630bea9ea60446b5472626f30fe95ec5
Binary files /dev/null and b/src/models/utils/__pycache__/transformer.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/uper.cpython-310.pyc b/src/models/utils/__pycache__/uper.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c93f0b65302f89c0220ecd86bd2ee5169c7de87f
Binary files /dev/null and b/src/models/utils/__pycache__/uper.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/uper.cpython-38.pyc b/src/models/utils/__pycache__/uper.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..16bf2389ccfcefb58471399c79a6c60a730fb345
Binary files /dev/null and b/src/models/utils/__pycache__/uper.cpython-38.pyc differ
diff --git a/src/models/utils/__pycache__/utils.cpython-310.pyc b/src/models/utils/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..633391a0b315009cfa6a74a53fffe694ff13a361
Binary files /dev/null and b/src/models/utils/__pycache__/utils.cpython-310.pyc differ
diff --git a/src/models/utils/__pycache__/utils.cpython-38.pyc b/src/models/utils/__pycache__/utils.cpython-38.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1ce3831e616e0ce7983cd7515b5edd7c7651dc38
Binary files /dev/null and b/src/models/utils/__pycache__/utils.cpython-38.pyc differ
diff --git a/src/models/utils/compute_mae.py b/src/models/utils/compute_mae.py
new file mode 100644
index 0000000000000000000000000000000000000000..c61c743b4c9dfab7b99b952d152bff2abf642ffb
--- /dev/null
+++ b/src/models/utils/compute_mae.py
@@ -0,0 +1,37 @@
+import torch
+import numpy as np
+
+def compute_mae(x1, x2, mask = None): # tensor [B, 3, H, W]
+ if mask is not None:
+ dot = torch.sum(x1 * x2 * mask, dim=1, keepdim=True)
+ dot = torch.max(torch.min(dot, torch.Tensor([1.0-1.0e-12]).cuda().float()), torch.Tensor([-1.0+1.0e-12]).cuda().float())
+ emap = torch.abs(180 * torch.acos(dot)/np.pi) * mask
+ mae = torch.sum(emap) / torch.sum(mask)
+ return mae, emap
+ if mask is None:
+ '''
+ x1 : [B, H*W, 3]
+ x2 : [B, H*W, 3]
+ '''
+ dot = torch.sum(x1 * x2, dim=1, keepdim=True)
+ dot = torch.max(torch.min(dot, torch.Tensor([1.0-1.0e-12]).cuda().float()), torch.Tensor([-1.0+1.0e-12]).cuda().float())
+ error = torch.abs(180 * torch.acos(dot)/np.pi) / x1.shape[1] / x1.shape[0]
+ return error
+
+
+
+def compute_mae_np(x1, x2, mask = None): # numpy array [H, W, 3]
+ x1 = np.divide(x1, np.linalg.norm(x1, axis=2, keepdims=True) + 1.0e-12)
+ x2 = np.divide(x2, np.linalg.norm(x2, axis=2, keepdims=True) + 1.0e-12)
+
+ if mask is not None:
+ dot = np.sum(x1 * x2 * mask[:, :, np.newaxis], axis=-1, keepdims=True)
+ dot = np.maximum(np.minimum(dot, np.array([1.0-1.0e-12])), np.array([-1.0+1.0e-12]))
+ emap = np.abs(180 * np.arccos(dot)/np.pi) * mask[:, :, np.newaxis]
+ mae = np.sum(emap) / np.sum(mask)
+ return mae, emap
+ if mask is None:
+ dot = np.sum(x1 * x2, axis=-1, keepdims=True)
+ dot = np.maximum(np.minimum(dot, np.array([1.0-1.0e-12])), np.array([-1.0+1.0e-12]))
+ error = np.abs(180 * np.arccos(dot)/np.pi)
+ return error
\ No newline at end of file
diff --git a/src/models/utils/convnext.py b/src/models/utils/convnext.py
new file mode 100644
index 0000000000000000000000000000000000000000..f76072d41da11e64be677b5869566e6770523d5b
--- /dev/null
+++ b/src/models/utils/convnext.py
@@ -0,0 +1,148 @@
+# This file is modified version from the original convnext
+# 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.
+
+
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Block(nn.Module):
+ r""" ConvNeXt Block. There are two equivalent implementations:
+ (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W)
+ (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back
+ We use (2) as we find it slightly faster in PyTorch
+
+ Args:
+ dim (int): Number of input channels.
+ drop_path (float): Stochastic depth rate. Default: 0.0
+ layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+ """
+ def __init__(self, dim, drop_path=0., layer_scale_init_value=1e-6):
+ super().__init__()
+ self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
+ self.norm = LayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
+ self.act = nn.GELU()
+ self.pwconv2 = nn.Linear(4 * dim, dim)
+ self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((dim)),
+ requires_grad=True) if layer_scale_init_value > 0 else None
+ self.drop_path = nn.Identity()
+
+ def forward(self, x):
+ input = x
+ x = self.dwconv(x)
+ x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+ x = self.norm(x)
+ x = self.pwconv1(x)
+ x = self.act(x)
+ x = self.pwconv2(x)
+ if self.gamma is not None:
+ x = self.gamma * x
+ x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
+
+ x = input + self.drop_path(x)
+ return x
+
+# @BACKBONES.register_module()
+class ConvNeXt(nn.Module):
+ r""" ConvNeXt
+ A PyTorch impl of : `A ConvNet for the 2020s` -
+ https://arxiv.org/pdf/2201.03545.pdf
+
+ Args:
+ in_chans (int): Number of input image channels. Default: 3
+ num_classes (int): Number of classes for classification head. Default: 1000
+ depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+ dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+ drop_path_rate (float): Stochastic depth rate. Default: 0.
+ layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6.
+ head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+ """
+ def __init__(self, in_chans=3, depths=[3, 3, 9, 3], dims=[96, 192, 384, 768],
+ drop_path_rate=0., layer_scale_init_value=1e-6, out_indices=[0, 1, 2, 3], use_checkpoint=False
+ ):
+ super().__init__()
+ self.use_checkpoint = use_checkpoint
+ self.downsample_layers = nn.ModuleList() # stem and 3 intermediate downsampling conv layers
+ stem = nn.Sequential(
+ nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+ LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
+ )
+ self.downsample_layers.append(stem)
+ for i in range(3):
+ downsample_layer = nn.Sequential(
+ LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+ nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
+ )
+ self.downsample_layers.append(downsample_layer)
+
+ self.stages = nn.ModuleList() # 4 feature resolution stages, each consisting of multiple residual blocks
+ dp_rates=[x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+ cur = 0
+ for i in range(4):
+ stage = nn.Sequential(
+ *[Block(dim=dims[i], drop_path=dp_rates[cur + j],
+ layer_scale_init_value=layer_scale_init_value) for j in range(depths[i])]
+ )
+
+
+ self.stages.append(stage)
+ cur += depths[i]
+
+ self.out_indices = out_indices
+
+ norm_layer = partial(LayerNorm, eps=1e-6, data_format="channels_first")
+ for i_layer in range(4):
+ layer = norm_layer(dims[i_layer])
+ layer_name = f'norm{i_layer}'
+ self.add_module(layer_name, layer)
+
+
+
+ def forward_features(self, x):
+ outs = []
+ for i in range(4):
+ x = self.downsample_layers[i](x)
+ x = self.stages[i](x)
+
+ if i in self.out_indices:
+ norm_layer = getattr(self, f'norm{i}')
+ x_out = norm_layer(x)
+ outs.append(x_out)
+
+ return tuple(outs)
+
+ def forward(self, x):
+ x = self.forward_features(x)
+ return x
+
+class LayerNorm(nn.Module):
+ r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+ shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+ with shape (batch_size, channels, height, width).
+ """
+ def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(normalized_shape))
+ self.bias = nn.Parameter(torch.zeros(normalized_shape))
+ self.eps = eps
+ self.data_format = data_format
+ if self.data_format not in ["channels_last", "channels_first"]:
+ raise NotImplementedError
+ self.normalized_shape = (normalized_shape, )
+
+ def forward(self, x):
+ if self.data_format == "channels_last":
+ return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+ elif self.data_format == "channels_first":
+ u = x.mean(1, keepdim=True)
+ s = (x - u).pow(2).mean(1, keepdim=True)
+ x = (x - u) / torch.sqrt(s + self.eps)
+ x = self.weight[:, None, None] * x + self.bias[:, None, None]
+ return x
diff --git a/src/models/utils/decompose_tensors.py b/src/models/utils/decompose_tensors.py
new file mode 100644
index 0000000000000000000000000000000000000000..2dccd53e1fdb5336a251dce25c49a524ea724963
--- /dev/null
+++ b/src/models/utils/decompose_tensors.py
@@ -0,0 +1,82 @@
+"""
+Scalable, Detailed and Mask-free Universal Photometric Stereo Network (CVPR2023)
+# Copyright (c) 2023 Satoshi Ikehata
+# All rights reserved.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+def divide_tensor_spatial(x, block_size=256, method='tile_stride'):
+ assert x.dim() == 4, "Input tensor must have 4 dimensions [B, C, H, W]"
+ B, C, H, W = x.shape
+ assert H == W, "Height and Width must be equal"
+ assert H % block_size == 0 and W % block_size ==0, "The tensor size cannot be divided by the block size"
+ mosaic_scale = H // block_size
+
+ if method == 'tile_stride':
+ """ decomposing x into K x K of (Hc, Wc) non-overlapped blocks (grid)"""
+
+ K = mosaic_scale * mosaic_scale
+ fold_params_grid = dict(kernel_size=(mosaic_scale, mosaic_scale), stride=(mosaic_scale, mosaic_scale), padding=(0,0), dilation=(1,1))
+ unfold_grid = nn.Unfold(**fold_params_grid)
+ tensor_grids = unfold_grid(x) # (B, C * K, Hm * Hm)
+ tensor_grids = tensor_grids.reshape(B, C, K, block_size, block_size).permute(0, 2, 1, 3, 4) # (B, K, C, Hm, Hm)
+ return tensor_grids
+
+ if method == 'tile_block':
+ tensor_blocks = x.view(B, C, mosaic_scale, block_size, mosaic_scale, block_size)
+ tensor_blocks = tensor_blocks.permute(0, 2, 4, 1, 3, 5) # (B, mc, mc, C, Hm, Wm)
+ tensor_blocks = tensor_blocks.contiguous().view(B, mosaic_scale**2, C, block_size, block_size) ## (B, K, C, Hm, Hm)
+ return tensor_blocks
+
+ return -1
+
+def merge_tensor_spatial(x, method='tile_stride'):
+
+ K, N, feat_dim, Hm, Wm = x.shape
+ mosaic_scale = int(math.sqrt(K))
+
+ if method == 'tile_stride':
+ x = x.reshape(K, N, feat_dim, -1)
+ fold_params_grid = dict(kernel_size=(mosaic_scale, mosaic_scale), stride=(mosaic_scale, mosaic_scale), padding=(0,0), dilation=(1,1))
+ fold_grid = nn.Fold(output_size=(Hm * mosaic_scale, Wm * mosaic_scale), **fold_params_grid) # downsample based on the encoder
+ x = x.permute(1, 2, 0, 3).reshape(N, feat_dim * K, -1)
+ x = fold_grid(x)
+ return x
+
+ if method == 'tile_block':
+ x = x.permute(1, 0, 2, 3, 4).reshape(N, mosaic_scale, mosaic_scale, feat_dim, Hm, Wm)
+ x = x.permute(0, 3, 1, 4, 2, 5)
+ x = x.reshape(N, feat_dim, mosaic_scale * Hm, mosaic_scale * Wm)
+ return x
+
+def divide_overlapping_patches(input_tensor, patch_size, margin):
+ B, C, W, _ = input_tensor.shape
+ stride = patch_size - margin
+ padded_W = ((W - patch_size + stride - 1) // stride) * stride + patch_size
+ pad = padded_W - W
+
+ padded_tensor = F.pad(input_tensor, (0, pad, 0, pad), mode='constant', value=0)
+
+ patches = F.unfold(padded_tensor, kernel_size=patch_size, stride=stride)
+ patches = patches.view(B, C, patch_size, patch_size, -1).permute(0, 4, 1, 2, 3)
+
+ return patches
+
+def merge_overlappnig_patches(patches, patch_size, margin, original_size):
+ B, _, C, _, _ = patches.shape
+ stride = patch_size - margin
+ W = original_size[2]
+
+ patches = patches.permute(0, 2, 3, 4, 1).contiguous().view(B, C * patch_size * patch_size, -1)
+ output = F.fold(patches, (W, W), kernel_size=patch_size, stride=stride)
+
+ weight = torch.ones(patches.size()).to(patches.device)
+ weight = F.fold(weight, (W, W), kernel_size=patch_size, stride=stride)
+
+ output = output / weight
+ return output
+
diff --git a/src/models/utils/fpn.py b/src/models/utils/fpn.py
new file mode 100644
index 0000000000000000000000000000000000000000..4418d4cf2232ea89cff1419f4fb774c016125b8a
--- /dev/null
+++ b/src/models/utils/fpn.py
@@ -0,0 +1,52 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+class FeatureFusionModule(nn.Module):
+ def __init__(self, in_channels_list, out_channels, fusion_type='sum'):
+ super(FeatureFusionModule, self).__init__()
+ self.fusion_type = fusion_type
+ self.adjusted_features_convs = nn.ModuleList([nn.Conv2d(in_channels, out_channels, 1) for in_channels in in_channels_list])
+ self.output_conv = nn.Conv2d(out_channels * len(in_channels_list) if fusion_type == 'concat' else out_channels, out_channels, 1)
+
+ def forward(self, features):
+ h, w = features[0].shape[-2:]
+ adjusted_features = []
+
+ for feature, conv in zip(features, self.adjusted_features_convs):
+ # チャンネル数を調整
+ feature = conv(feature)
+ # アップサンプリング
+ adjusted_feature = F.interpolate(feature, size=(h, w), mode='bilinear', align_corners=True)
+ adjusted_features.append(adjusted_feature)
+
+ if self.fusion_type == 'concat':
+ fused_feature = torch.cat(adjusted_features, dim=1)
+ elif self.fusion_type == 'sum':
+ fused_feature = sum(adjusted_features)
+ else:
+ raise ValueError(f"Invalid fusion_type: {self.fusion_type}")
+
+ fused_feature = self.output_conv(fused_feature)
+ return fused_feature
+
+
+# # 特徴マップのリスト
+# feature_maps = [
+# torch.randn(10, 32, 64, 64),
+# torch.randn(10, 64, 32, 32),
+# torch.randn(10, 128, 16, 16),
+# torch.randn(10, 256, 8, 8),
+# ]
+
+# # モデルの構築
+# in_channels_list = [32, 64, 128, 256] # 各レベルのチャンネル数
+# out_channels = 64 # 融合後のチャンネル数を64に指定
+# fusion_type = 'sum' # 'sum'または'concat'を指定
+
+# feature_fusion_module = FeatureFusionModule(in_channels_list, out_channels, fusion_type=fusion_type)
+
+# # 融合
+# fused_feature = feature_fusion_module(feature_maps)
+
+# print(fused_feature.shape) # 融合後の特徴マップのサイズとチャンネル数を確認
diff --git a/src/models/utils/gauss_filter.py b/src/models/utils/gauss_filter.py
new file mode 100644
index 0000000000000000000000000000000000000000..52ceb7b8cb8dbbba6dbda9f1a021137efd74589c
--- /dev/null
+++ b/src/models/utils/gauss_filter.py
@@ -0,0 +1,77 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+import numbers
+
+class gauss_filter(nn.Module):
+ """
+ Apply gaussian smoothing on a
+ 1d, 2d or 3d tensor. Filtering is performed seperately for each channel
+ in the input using a depthwise convolution.
+ Arguments:
+ channels (int, sequence): Number of channels of the input tensors. Output will
+ have this number of channels as well.
+ kernel_size (int, sequence): Size of the gaussian kernel.
+ sigma (float, sequence): Standard deviation of the gaussian kernel.
+ dim (int, optional): The number of dimensions of the data.
+ Default value is 2 (spatial).
+ """
+ def __init__(self, channels, kernel_size, sigma, dim=2):
+ super(gauss_filter, self).__init__()
+ if isinstance(kernel_size, numbers.Number):
+ kernel_size = [kernel_size] * dim
+ if isinstance(sigma, numbers.Number):
+ sigma = [sigma] * dim
+
+ # The gaussian kernel is the product of the
+ # gaussian function of each dimension.
+ kernel = 1
+ self.kernel_size = kernel_size
+ meshgrids = torch.meshgrid(
+ [
+ torch.arange(size, dtype=torch.float32)
+ for size in kernel_size
+ ]
+ )
+ for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
+ mean = (size - 1) / 2
+ kernel *= 1 / (std * math.sqrt(2 * math.pi)) * \
+ torch.exp(-((mgrid - mean) / std) ** 2 / 2)
+
+ # Make sure sum of values in gaussian kernel equals 1.
+ kernel = kernel / torch.sum(kernel)
+
+ # Reshape to depthwise convolutional weight
+ kernel = kernel.view(1, 1, *kernel.size())
+ kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
+
+ self.register_buffer('weight', kernel)
+ self.groups = channels
+
+ if dim == 1:
+ self.conv = F.conv1d
+ elif dim == 2:
+ self.conv = F.conv2d
+ elif dim == 3:
+ self.conv = F.conv3d
+ else:
+ raise RuntimeError(
+ 'Only 1, 2 and 3 dimensions are supported. Received {}.'.format(dim)
+ )
+
+ def forward(self, input):
+ """
+ Apply gaussian filter to input.
+ Arguments:
+ input (torch.Tensor): Input to apply gaussian filter on.
+ Returns:
+ filtered (torch.Tensor): Filtered output.
+ """
+ return self.conv(input, weight=self.weight, groups=self.groups, padding=(self.kernel_size[0]//2, self.kernel_size[1]//2))
+
+
+# smoothing = GaussianSmoothing(3, 5, 1)
+# input = torch.rand(1, 3, 100, 100)
+# input = F.pad(input, (2, 2, 2, 2), mode='reflect')
+# output = smoothing(input)
\ No newline at end of file
diff --git a/src/models/utils/ind2sub.py b/src/models/utils/ind2sub.py
new file mode 100644
index 0000000000000000000000000000000000000000..03c2d390bb7a125a05267baa308366d7bfffa9d5
--- /dev/null
+++ b/src/models/utils/ind2sub.py
@@ -0,0 +1,12 @@
+import numpy as np
+import torch
+"""array_shape (rows, cols)
+ind = row * cols + col
+"""
+def ind2coords(array_shape, ind):
+ row = torch.div(ind, array_shape[1], rounding_mode='trunc')
+ col = ind % array_shape[1] # or numpy.mod(ind.astype('int'), array_shape[1])
+ coords = torch.zeros((1, 1, len(ind), 2), dtype=torch.float32)
+ coords[:, :, :, 1] = 2 * row.to(torch.float32) / array_shape[0] - 1
+ coords[:, :, :, 0] = 2 * col.to(torch.float32) / array_shape[1] - 1
+ return coords
\ No newline at end of file
diff --git a/src/models/utils/logger.py b/src/models/utils/logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..3425fd7b3e3ea7e076a6f131aed55163b37f9a79
--- /dev/null
+++ b/src/models/utils/logger.py
@@ -0,0 +1,28 @@
+from torch.utils.tensorboard import SummaryWriter
+import os, sys, glob
+
+
+class logger():
+ def __init__(self, args, mode, suffix = None):
+ self.outdir = f'{args.outdir}/{args.session_name}'
+ if args.test == False:
+ logdir = f'{self.outdir}/log'
+ for file in glob.glob(f'{logdir}/*'):
+ if os.path.isfile(file):
+ os.remove(file)
+ self.writer = SummaryWriter(log_dir=f'{logdir}', flush_secs=1)
+
+
+ def add(self, tag, item, step, itemtype):
+
+ if itemtype == 'Image':
+ if len(item.shape) == 3:
+ self.writer.add_image(tag, item, global_step=step, dataformats='CHW')
+ elif len(item.shape) == 4:
+ self.writer.add_images(tag, item, global_step=step, dataformats='NCHW')
+ else:
+ raise Exception("item.shape must be 3 or 4 (%d)" % len(item.shape))
+ elif itemtype == 'Scalar':
+ self.writer.add_scalar(tag, item, global_step=step)
+ else:
+ print('itemtype is not in "Image, Scalar"', file = sys.stderr)
diff --git a/src/models/utils/model_utils.py b/src/models/utils/model_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..60b5663220bf7876a9022aca358b1510b263eefc
--- /dev/null
+++ b/src/models/utils/model_utils.py
@@ -0,0 +1,199 @@
+import os
+import torch
+import numpy as np
+from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.nn.functional as F
+import matplotlib.pyplot as plt
+def pca_show(feats):
+ feats = feats.cpu().detach()
+ B, C, H, W = feats.shape
+ n_components = 3
+ from sklearn.decomposition import PCA
+ pca = PCA(n_components=n_components)
+ feats = F.interpolate(feats, size=(128, 128), mode='bilinear', align_corners=False)
+ features = feats[0, :, :, :].squeeze().reshape(C,-1).permute(1,0).cpu()
+ pca.fit(features)
+ pca_features = pca.transform(features)
+ pca_features = (pca_features - pca_features.min()) / (pca_features.max() - pca_features.min())
+ pca_features = pca_features * 255
+ pca_features = pca_features.reshape(128, 128, n_components).astype(np.uint8)
+ plt.imshow(pca_features)
+
+
+def loadmodel(model, filename, strict=True, remove_prefix=False):
+ if os.path.exists(filename):
+ params = torch.load('%s' % filename)
+ if remove_prefix:
+ new_params = {k.replace('module.', ''): v for k, v in params.items()}
+ model.load_state_dict(new_params,strict=strict)
+ print("prefix successfully removed")
+ else:
+ model.load_state_dict(params,strict=strict)
+ print('Load %s' % filename)
+ else:
+ print('Model Not Found')
+ return model
+
+def loadoptimizer(optimizer, filename):
+ if os.path.exists(filename):
+ params = torch.load('%s' % filename)
+ optimizer.load_state_dict(params)
+ print('Load %s' % filename)
+ return optimizer
+
+def loadscheduler(scheduler, filename):
+ if os.path.exists(filename):
+ params = torch.load('%s' % filename)
+ scheduler.load_state_dict(params)
+ print('Load %s' % filename)
+ else:
+ print('Scheduler Not Found')
+ return scheduler
+
+def savemodel(model, filename):
+ print('Save %s' % filename)
+ torch.save(model.state_dict(), filename)
+
+def saveoptimizer(optimizer, filename):
+ print('Save %s' % filename)
+ torch.save(optimizer.state_dict(), filename)
+
+def savescheduler(scheduler, filename):
+ print('Save %s' % filename)
+ torch.save(scheduler.state_dict(), filename)
+
+def optimizer_setup_Adam(net, lr = 0.0001, init=True, step_size=3, stype='step'):
+ print(f'Optimizer (Adam) lr={lr}')
+ if init==True:
+ net.init_weights()
+ net = torch.nn.DataParallel(net)
+ optim_params = [{'params': net.parameters(), 'lr': lr},] # confirmed
+ optimizer = torch.optim.Adam(optim_params, betas=(0.9, 0.999), weight_decay=0)
+ if stype == 'cos':
+ scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 30, eta_min=0, last_epoch=-1)
+ print('Cosine aneealing learning late scheduler')
+ if stype == 'step':
+ scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.8)
+ print(f'Step late scheduler x0.8 decay every {step_size}')
+ return net, optimizer, scheduler
+
+def optimizer_setup_SGD(net, lr = 0.01, momentum= 0.9, init=True):
+ print(f'Optimizer (SGD with momentum) lr={lr}')
+ if init==True:
+ net.init_weights()
+ net = torch.nn.DataParallel(net)
+ optim_params = [{'params': net.parameters(), 'lr': lr},] # confirmed
+ return net, torch.optim.SGD(optim_params, momentum=momentum, weight_decay=1e-4, nesterov=True)
+
+def optimizer_setup_AdamW(net, lr = 0.001, eps=1.0e-8, step_size = 20, init=True, stype='step', use_data_parallel=False):
+ print(f'Optimizer (AdamW) lr={lr}')
+ if init==True:
+ net.init_weights()
+ if use_data_parallel:
+ net = torch.nn.DataParallel(net)
+ optim_params = [{'params': net.parameters(), 'lr': lr},] # confirmed
+ optimizer = torch.optim.AdamW(optim_params, betas=(0.9, 0.999), eps=eps, weight_decay=0.01)
+
+ if stype=='cos':
+ scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 30, eta_min=0, last_epoch=-1)
+ print('Cosine aneealing learning late scheduler')
+ if stype == 'step':
+ scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=0.8)
+ print(f'Step late scheduler x0.8 decay every {step_size}')
+ return net, optimizer, scheduler
+
+def mode_change(net, Training):
+ if Training == True:
+ for param in net.parameters():
+ param.requires_grad = True
+ net.train()
+ if Training == False:
+ for param in net.parameters():
+ param.requires_grad = False
+ net.eval()
+
+
+def get_n_params(model):
+ pp=0
+ for p in list(model.parameters()):
+ nn=1
+ for s in list(p.size()):
+ nn = nn*s
+ pp += nn
+ return pp
+
+def loadCheckpoint(path, model, cuda=True):
+ if cuda:
+ checkpoint = torch.load(path)
+ else:
+ checkpoint = torch.load(path, map_location=lambda storage, loc: storage)
+ model.load_state_dict(checkpoint['state_dict'])
+
+def saveCheckpoint(save_path, epoch=-1, model=None, optimizer=None, records=None, args=None):
+ state = {'state_dict': model.state_dict(), 'model': args.model}
+ records = {'epoch': epoch, 'optimizer':optimizer.state_dict(), 'records': records,
+ 'args': args}
+ torch.save(state, os.path.join(save_path, 'checkp_%d.pth.tar' % (epoch)))
+ torch.save(records, os.path.join(save_path, 'checkp_%d_rec.pth.tar' % (epoch)))
+
+
+
+def masking(img, mask):
+ # img [B, C, H, W]
+ # mask [B, 1, H, W] [0,1]
+ img_masked = img * mask.expand((-1, img.shape[1], -1, -1))
+ return img_masked
+
+def print_model_parameters(model):
+ model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+ params = sum([np.prod(p.size()) for p in model_parameters])
+
+
+ print('# parameters: %d' % params)
+
+
+def angular_error(x1, x2, mask = None): # tensor [B, 3, H, W]
+
+ if mask is not None:
+ dot = torch.sum(x1 * x2 * mask, dim=1, keepdim=True)
+ dot = torch.max(torch.min(dot, torch.Tensor([1.0-1.0e-12])), torch.Tensor([-1.0+1.0e-12]))
+ emap = torch.abs(180 * torch.acos(dot)/np.pi) * mask
+ mae = torch.sum(emap) / torch.sum(mask)
+ return mae, emap
+ if mask is None:
+ dot = torch.sum(x1 * x2, dim=1, keepdim=True)
+ dot = torch.max(torch.min(dot, torch.Tensor([1.0-1.0e-12])), torch.Tensor([-1.0+1.0e-12]))
+ error = torch.abs(180 * torch.acos(dot)/np.pi)
+ return error
+
+def write_errors(filepath, error, trainid, numimg, objname = []):
+ dt_now = datetime.datetime.now()
+ print(filepath)
+
+ if len(objname) > 0:
+ with open(filepath, 'a') as f:
+ f.write('%s %03d %s %02d %.2f\n' % (dt_now, numimg, objname, trainid, error))
+ else:
+ with open(filepath, 'a') as f:
+ f.write('%s %03d %02d %.2f\n' % (dt_now, numimg, trainid, error))
+
+
+def save_nparray_as_hdf5(self, a, filename):
+ h5f = h5py.File(filename, 'w')
+ h5f.create_dataset('dataset_1', data=a)
+ h5f.close()
+
+def freeze_params(net):
+ for param in net.parameters():
+ param.requires_grad = False
+
+def unfreeze_params(net):
+ for param in net.parameters():
+ param.requires_grad = True
+
+def make_index_list(maxNumImages, numImageList):
+ index = np.zeros((len(numImageList) * maxNumImages), np.int32)
+ for k in range(len(numImageList)):
+ index[maxNumImages*k:maxNumImages*k+numImageList[k]] = 1
+ return index
+
diff --git a/src/models/utils/parser_utils.py b/src/models/utils/parser_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6b9dbcac7369bcca2302b6c81084f49e070868f
--- /dev/null
+++ b/src/models/utils/parser_utils.py
@@ -0,0 +1,31 @@
+import argparse
+import json
+import os
+
+def save_args(args, dst):
+ os.makedirs(dst, exist_ok=True)
+ if dst is None:
+ return
+ with open(f"{dst}args.json", 'w+') as f:
+ json.dump(vars(args), f)
+
+def load_args(args, src):
+ args_copy = args
+ if src is not None and os.path.exists(src):
+ with open(src, 'r') as f:
+ ns = json.load(f)
+ args_loaded = argparse.Namespace(**ns)
+ for arg in vars(args_copy):
+ if arg in ('training_dir', 'test_dir', 'pretrained',
+ 't_check_point', 't_test_dir', 't_out_dir', 't_num_img', 't_test_suffix', 't_test_ext', 't_test_data_type', 't_loader_imgsize', 't_normalize'):
+ setattr(args_loaded, arg, getattr(args_copy, arg))
+ return args_loaded
+ else:
+ print('no arg file found! args was not updated')
+ return args
+
+
+
+
+
+
diff --git a/src/models/utils/pytorch_ssim.py b/src/models/utils/pytorch_ssim.py
new file mode 100644
index 0000000000000000000000000000000000000000..dd6c7000ead23a0c0df245f508f29f5d295939bb
--- /dev/null
+++ b/src/models/utils/pytorch_ssim.py
@@ -0,0 +1,73 @@
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+import numpy as np
+from math import exp
+
+def gaussian(window_size, sigma):
+ gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+ return gauss/gauss.sum()
+
+def create_window(window_size, channel):
+ _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+ _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+ window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+ return window
+
+def _ssim(img1, img2, window, window_size, channel, size_average = True):
+ mu1 = F.conv2d(img1, window, padding = window_size//2, groups = channel)
+ mu2 = F.conv2d(img2, window, padding = window_size//2, groups = channel)
+
+ mu1_sq = mu1.pow(2)
+ mu2_sq = mu2.pow(2)
+ mu1_mu2 = mu1*mu2
+
+ sigma1_sq = F.conv2d(img1*img1, window, padding = window_size//2, groups = channel) - mu1_sq
+ sigma2_sq = F.conv2d(img2*img2, window, padding = window_size//2, groups = channel) - mu2_sq
+ sigma12 = F.conv2d(img1*img2, window, padding = window_size//2, groups = channel) - mu1_mu2
+
+ C1 = 0.01**2
+ C2 = 0.03**2
+
+ ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+
+ if size_average:
+ return ssim_map.mean()
+ else:
+ return ssim_map.mean(1).mean(1).mean(1)
+
+class SSIM(torch.nn.Module):
+ def __init__(self, window_size = 11, size_average = True):
+ super(SSIM, self).__init__()
+ self.window_size = window_size
+ self.size_average = size_average
+ self.channel = 1
+ self.window = create_window(window_size, self.channel)
+
+ def forward(self, img1, img2):
+ (_, channel, _, _) = img1.size()
+
+ if channel == self.channel and self.window.data.type() == img1.data.type():
+ window = self.window
+ else:
+ window = create_window(self.window_size, channel)
+
+ if img1.is_cuda:
+ window = window.cuda(img1.get_device())
+ window = window.type_as(img1)
+
+ self.window = window
+ self.channel = channel
+
+
+ return _ssim(img1, img2, window, self.window_size, channel, self.size_average)
+
+def ssim(img1, img2, window_size = 11, size_average = True):
+ (_, channel, _, _) = img1.size()
+ window = create_window(window_size, channel)
+
+ if img1.is_cuda:
+ window = window.cuda(img1.get_device())
+ window = window.type_as(img1)
+
+ return _ssim(img1, img2, window, window_size, channel, size_average)
diff --git a/src/models/utils/save_normal.py b/src/models/utils/save_normal.py
new file mode 100644
index 0000000000000000000000000000000000000000..27b28962a181d54c0960377498cd9bb79dff6ed7
--- /dev/null
+++ b/src/models/utils/save_normal.py
@@ -0,0 +1,17 @@
+ # """ Save normal and mask"""
+
+ # # temp = nout[0, :, :, :].permute(1,2,0).cpu().detach().numpy()
+ # # ### swap for surface normal integration
+ # # normal_map_est = np.zeros(temp.shape, np.float32)
+ # # normal_map_est[:, :, 0] = temp[:, :, 1]
+ # # normal_map_est[:, :, 1] = temp[:, :, 0]
+ # # normal_map_est[:, :, 2] = -temp[:, :, 2]
+ # # temp = nml[0, :, :, :].permute(1,2,0).cpu().detach().numpy()
+ # # normal_map_gt = np.zeros(temp.shape, np.float32)
+ # # normal_map_gt[:, :, 0] = temp[:, :, 1]
+ # # normal_map_gt[:, :, 1] = temp[:, :, 0]
+ # # normal_map_gt[:, :, 2] = -temp[:, :, 2]
+ # # mask = m[0, 0, :, :].cpu().detach().numpy().astype(np.bool)
+ # # result={'normal_map_est':normal_map_est, 'normal_map_gt': normal_map_gt ,'mask':mask}
+ # loss = 0
+ # return loss, output.cpu().detach().numpy(), input.cpu().detach().numpy(), result
\ No newline at end of file
diff --git a/src/models/utils/show_tensor.py b/src/models/utils/show_tensor.py
new file mode 100644
index 0000000000000000000000000000000000000000..46d60ce877cd6dd3e7f820bd133de323c9ce4193
--- /dev/null
+++ b/src/models/utils/show_tensor.py
@@ -0,0 +1,48 @@
+import torch.nn as nn
+import torch
+import math
+import numpy as np
+import matplotlib.pyplot as plt
+
+def masking(img, mask):
+ # img [B, C, H, W]
+ # mask [B, 1, H, W] [0,1]
+ img_masked = img * mask.expand((-1, img.shape[1], -1, -1))
+ return img_masked
+
+def imshow(img, mask = None, vmax = None, axis = None):
+ if mask is None:
+ img = img.data.cpu().numpy()[0,:,:,:]
+ else:
+ img = masking(img,mask).data.cpu().numpy()[0,:,:,:]
+ c = img.shape[0]
+ h = img.shape[1]
+ w = img.shape[2]
+ if c == 3:
+ img = np.reshape(img, (3,h,w)).transpose(1,2,0)
+ if c == 1:
+ img = img[0,:,:]
+ # plt.figure(figsize = (8,8))
+ if vmax is None:
+ if axis is None:
+ plt.imshow(img)
+ else:
+ axis.imshow(img)
+ else:
+ if axis is None:
+ plt.imshow(img, vmax = vmax)
+ else:
+ axis.imshow(img, vmax = vmax)
+
+def nmlshow(nml, mask = None, axis = None):
+ if mask is None:
+ nml = nml.data.cpu().numpy()[0,:,:,:]
+ else:
+ nml = masking(nml,mask).data.cpu().numpy()[0,:,:,:]
+ nml = np.transpose(nml, (1,2,0))
+ # plt.figure(figsize = (8,8))
+
+ if axis is None:
+ plt.imshow(-0.5 * nml + 0.5)
+ else:
+ axis.imshow(-0.5 * nml + 0.5)
diff --git a/src/models/utils/transformer.py b/src/models/utils/transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..f333a7758b2597db4b1e2b4aa0403ebb0eb715d3
--- /dev/null
+++ b/src/models/utils/transformer.py
@@ -0,0 +1,233 @@
+"""
+Scalable, Detailed and Mask-free Universal Photometric Stereo Network (CVPR2023)
+# Copyright (c) 2023 Satoshi Ikehata
+# All rights reserved.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+import math
+import torch.nn.init as init
+
+class MultiHeadAttentionBlock(nn.Module):
+ def __init__(self, dim_in, dim_out, num_heads, ln=False, attention_dropout = 0.1, dim_feedforward = 512,
+ q_bucket_size = 1024, k_bucket_size = 2048, attn_mode = 'Normal'):
+ super(MultiHeadAttentionBlock, self).__init__()
+ if attn_mode == 'Efficient':
+ self.q_bucket_size = q_bucket_size
+ self.k_bucket_size = k_bucket_size
+ self.attn_mode = attn_mode
+ self.dim_V = dim_out
+ self.dim_Q = dim_in
+ self.dim_K = dim_in
+ self.num_heads = num_heads
+ self.fc_q = nn.Linear(self.dim_Q, self.dim_V, bias=False) # dimin -> dimhidden
+ self.fc_k = nn.Linear(self.dim_K, self.dim_V, bias=False) # dimin -> dimhidden
+ self.fc_v = nn.Linear(self.dim_K, self.dim_V, bias=False) # dimhidden -> dim
+ if ln:
+ self.ln0 = nn.LayerNorm(self.dim_Q)
+ self.ln1 = nn.LayerNorm(self.dim_V)
+
+ self.dropout_attn = nn.Dropout(attention_dropout)
+ self.fc_o1 = nn.Linear(self.dim_V, dim_feedforward, bias=False)
+ self.fc_o2 = nn.Linear(dim_feedforward, self.dim_V, bias=False)
+ self.dropout1 = nn.Dropout(attention_dropout)
+ self.dropout2 = nn.Dropout(attention_dropout)
+
+ # memory efficient attention related parameters
+ # can be overriden on forward
+ self.q_bucket_size = q_bucket_size
+ self.k_bucket_size = k_bucket_size
+
+
+ # memory efficient attention
+ def summarize_qkv_chunk(self, q, k, v):
+ weight = torch.einsum('b h i d, b h j d -> b h i j', q, k)
+ weight_max = weight.amax(dim = -1, keepdim = True).detach()
+ weight = weight - weight_max
+
+ exp_weight = self.dropout_attn(weight.exp()) # attention_dropout
+ weighted_value = torch.einsum('b h i j, b h j d -> b h i d', exp_weight, v)
+ return exp_weight.sum(dim = -1), weighted_value, rearrange(weight_max, '... 1 -> ...')
+
+ def memory_efficient_attention(
+ self,
+ q, k, v,
+ q_bucket_size = 512,
+ k_bucket_size = 1024,
+ eps = 1e-8,
+ ):
+ scale = q.shape[-1] ** -0.5
+ q = q * scale
+
+ summarize_qkv_fn = self.summarize_qkv_chunk
+
+ # chunk all the inputs
+ q_chunks = q.split(q_bucket_size, dim = -2)
+ k_chunks = k.split(k_bucket_size, dim = -2)
+ v_chunks = v.split(k_bucket_size, dim = -2)
+
+ # loop through all chunks and accumulate
+ values = []
+ weights = []
+ for q_chunk in q_chunks:
+ exp_weights = []
+ weighted_values = []
+ weight_maxes = []
+ for (k_chunk, v_chunk) in zip(k_chunks, v_chunks):
+ exp_weight_chunk, weighted_value_chunk, weight_max_chunk = summarize_qkv_fn(
+ q_chunk,
+ k_chunk,
+ v_chunk
+ )
+ exp_weights.append(exp_weight_chunk)
+ weighted_values.append(weighted_value_chunk)
+ weight_maxes.append(weight_max_chunk)
+ weight_maxes = torch.stack(weight_maxes, dim = -1)
+ weighted_values = torch.stack(weighted_values, dim = -1)
+ exp_weights = torch.stack(exp_weights, dim = -1)
+
+ global_max = weight_maxes.amax(dim = -1, keepdim = True)
+ renorm_factor = (weight_maxes - global_max).exp().detach()
+
+ exp_weights = exp_weights * renorm_factor
+ weighted_values = weighted_values * rearrange(renorm_factor, '... c -> ... 1 c')
+
+ all_values = weighted_values.sum(dim = -1)
+ all_weights = exp_weights.sum(dim = -1)
+ values.append(all_values)
+ weights.append(all_weights)
+ values = torch.cat(values, dim=2)
+ weights = torch.cat(weights, dim=2)
+ # (rearrange(weights, '... -> ... 1')
+ normalized_values = values / (rearrange(weights, '... -> ... 1') + eps)
+ return normalized_values
+
+
+ def forward(
+ self,
+ x,y,
+ ):
+
+ x = x if getattr(self, 'ln0', None) is None else self.ln0(x) # pre-normalization
+ Q = self.fc_q(x) # input_dim -> embed dim
+ K, V = self.fc_k(y), self.fc_v(y) # input_dim -> embed dim
+ dim_split = self.dim_V // self.num_heads # multi-head attention
+ if self.attn_mode == 'Efficient':
+ q_bucket_size = self.q_bucket_size
+ k_bucket_size = self.k_bucket_size
+
+ Q_ = torch.stack(Q.split(int(dim_split), 2), 1)
+ K_ = torch.stack(K.split(int(dim_split), 2), 1)
+ V_ = torch.stack(V.split(int(dim_split), 2), 1)
+
+ A = self.memory_efficient_attention(Q_, K_, V_, q_bucket_size = q_bucket_size, k_bucket_size = k_bucket_size)
+ A = A.reshape(-1, A.shape[2], A.shape[3])
+ Q_ = Q_.reshape(-1, Q_.shape[2], Q_.shape[3])
+ O = torch.cat((Q_ + A).split(Q.size(0), 0), 2)
+ else: # Basic
+ Q_ = torch.cat(Q.split(int(dim_split), 2), 0)
+ K_ = torch.cat(K.split(int(dim_split), 2), 0)
+ V_ = torch.cat(V.split(int(dim_split), 2), 0)
+ A = self.dropout_attn(torch.softmax(Q_.bmm(K_.transpose(1,2))/math.sqrt(self.dim_V), 2)) # this may not be correct due to mult-head attention
+ A = A.bmm(V_) # A(Q, K, V) attention_output
+ O = torch.cat((Q_ + A).split(Q.size(0), 0), 2)
+
+
+ O_ = O if getattr(self, 'ln1', None) is None else self.ln1(O)
+ O = O + self.dropout2(self.fc_o2(self.dropout1(F.gelu(self.fc_o1(O_)))))
+ return O
+
+
+class SAB(nn.Module): # self attention block
+ def __init__(self, dim_in, dim_out, num_heads=4, ln=False, attention_dropout = 0.1, dim_feedforward = 512, attn_mode = 'Normal'):
+ super(SAB, self).__init__()
+ self.mab = MultiHeadAttentionBlock(dim_in, dim_out, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode)
+ def forward(self, X):
+ return self.mab(X, X)
+
+class CAB(nn.Module): # cross attention block
+ def __init__(self, dim_in, dim_out, num_heads=4, ln=False, attention_dropout = 0.1, dim_feedforward = 512, attn_mode = 'Normal'):
+ super(CAB, self).__init__()
+ self.mab = MultiHeadAttentionBlock(dim_in, dim_out, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode)
+ def forward(self, q, kv):
+ return self.mab(q, kv)
+
+class PMA(nn.Module):
+ def __init__(self, dim, num_heads, num_seeds, ln=False, attn_mode='Normal'):
+ super(PMA, self).__init__()
+ self.S = nn.Parameter(torch.Tensor(1, num_seeds, dim))
+ init.xavier_uniform_(self.S)
+ self.mab = MultiHeadAttentionBlock(dim, dim, num_heads, ln=ln, attn_mode=attn_mode)
+
+ def forward(self, X):
+ return self.mab(self.S.repeat(X.size(0), 1, 1), X)
+
+class CommunicationBlock(nn.Module):
+ def __init__(self, dim_input, num_enc_sab = 3, dim_hidden=384, dim_feedforward = 1024, num_heads=8, ln=False, attention_dropout=0.1, use_efficient_attention=False):
+ super(CommunicationBlock, self).__init__()
+
+ if use_efficient_attention:
+ attn_mode = 'Efficient'
+ else:
+ attn_mode = 'Normal'
+ self.dim_hidden = dim_hidden
+ modules_enc = []
+ modules_enc.append(SAB(dim_input, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
+ for k in range(num_enc_sab):
+ modules_enc.append(SAB(dim_hidden, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
+ self.enc = nn.Sequential(*modules_enc)
+
+ def forward(self, x):
+ x = self.enc(x)
+ return x
+
+class CrossAttentionBlock(nn.Module):
+ def __init__(self, dim_input, num_enc_sab = 3, dim_hidden=384, dim_feedforward = 1024, num_heads=8, ln=False, attention_dropout=0.1, use_efficient_attention=False):
+ super(CrossAttentionBlock, self).__init__()
+
+ if use_efficient_attention:
+ attn_mode = 'Efficient'
+ else:
+ attn_mode = 'Normal'
+ self.dim_hidden = dim_hidden
+ modules_enc = []
+ modules_enc.append(CAB(dim_input, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
+ for k in range(num_enc_sab):
+ modules_enc.append(CAB(dim_hidden, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
+ self.layers = nn.ModuleList(modules_enc)
+ def forward(self, q, kv):
+ for k in range(len(self.layers)):
+ q = self.layers[k](q, kv) # x[0] query, x[1] key value
+ return q # the output token length is len(x)
+
+class AggregationBlock(nn.Module):
+ def __init__(self, dim_input, num_enc_sab = 3, num_outputs = 1, dim_hidden=384, dim_feedforward = 1024, num_heads=8, ln=False, attention_dropout=0.1, use_efficient_attention=False):
+ super(AggregationBlock, self).__init__()
+
+ self.num_outputs = num_outputs
+ self.dim_hidden = dim_hidden
+
+ if use_efficient_attention:
+ attn_mode = 'Efficient'
+ else:
+ attn_mode = 'Normal'
+
+ modules_enc = []
+ modules_enc.append(SAB(dim_input, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
+ for k in range(num_enc_sab):
+ modules_enc.append(SAB(dim_hidden, dim_hidden, num_heads, ln=ln, attention_dropout = attention_dropout, dim_feedforward=dim_feedforward, attn_mode=attn_mode))
+ self.enc = nn.Sequential(*modules_enc)
+ modules_dec = []
+ modules_dec.append(PMA(dim_hidden, num_heads, num_outputs, attn_mode=attn_mode)) # after the PMA we should not put drop out
+ self.dec = nn.Sequential(*modules_dec)
+
+ def forward(self, x):
+ x = self.enc(x)
+ x = self.dec(x)
+ x = x.view(-1, self.num_outputs * self.dim_hidden)
+ return x
+
+
diff --git a/src/models/utils/uper.py b/src/models/utils/uper.py
new file mode 100644
index 0000000000000000000000000000000000000000..83ee48dadd35f2da2bd887053f44843797beaada
--- /dev/null
+++ b/src/models/utils/uper.py
@@ -0,0 +1,143 @@
+"""This file is modified version from mmsegmentation (https://github.com/open-mmlab/mmsegmentation)"""
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+class PPM(nn.ModuleList):
+ """Pooling Pyramid Module used in PSPNet.
+ Args:
+ pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
+ Module.
+ in_channels (int): Input channels.
+ channels (int): Channels after modules, before conv_seg.
+ conv_cfg (dict|None): Config of conv layers.
+ norm_cfg (dict|None): Config of norm layers.
+ act_cfg (dict): Config of activation layers.
+ align_corners (bool): align_corners argument of F.interpolate.
+ """
+
+ def __init__(self, pool_scales, in_channels, channels):
+ super(PPM, self).__init__()
+ self.pool_scales = pool_scales
+ self.in_channels = in_channels
+ self.channels = channels
+ for pool_scale in pool_scales:
+ self.append(
+ nn.Sequential(
+ nn.AdaptiveAvgPool2d(pool_scale),
+ nn.Conv2d(self.in_channels, self.channels, kernel_size=1),
+ nn.ReLU()
+ )
+ )
+
+ def forward(self, x):
+ """Forward function."""
+ ppm_outs = []
+ for ppm in self:
+ ppm_out = ppm(x)
+
+ upsampled_ppm_out = F.interpolate(
+ ppm_out.float(),
+ size=x.size()[2:],
+ mode='bilinear',
+ align_corners=False).to(torch.bfloat16)
+
+ ppm_outs.append(upsampled_ppm_out)
+ return ppm_outs
+
+class UPerHead(nn.Module):
+ """Unified Perceptual Parsing for Scene Understanding.
+ This head is the implementation of `UPerNet
+ `_.
+ Args:
+ pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
+ Module applied on the last feature. Default: (1, 2, 3, 6).
+ """
+
+ def __init__(self, in_channels = (96, 192, 384, 768), channels = 256, pool_scales=(1, 2, 3, 6),):
+ super(UPerHead, self).__init__()
+ # PSP Module
+ self.in_channels = in_channels
+ self.channels = channels
+ self.psp_modules = PPM(
+ pool_scales,
+ self.in_channels[-1],
+ self.channels
+ )
+
+ self.bottleneck = nn.Sequential(
+ nn.Conv2d(self.in_channels[-1] + len(pool_scales) * self.channels, self.channels, kernel_size=3, padding=1),
+ nn.ReLU())
+ # FPN Module
+ self.lateral_convs = nn.ModuleList()
+ self.fpn_convs = nn.ModuleList()
+ for in_channels in self.in_channels[:-1]: # skip the top layer
+ l_conv = nn.Sequential(
+ nn.Conv2d(in_channels, self.channels, kernel_size=1, padding=0),
+ nn.ReLU())
+
+
+ fpn_conv = nn.Sequential(
+ nn.Conv2d(self.channels, self.channels, kernel_size=3, padding=1),
+ nn.ReLU())
+
+ self.lateral_convs.append(l_conv)
+ self.fpn_convs.append(fpn_conv)
+
+ self.fpn_bottleneck = nn.Sequential(
+ nn.Conv2d(len(self.in_channels) * self.channels, self.channels, kernel_size=3, padding=1),
+ nn.ReLU())
+
+
+ def psp_forward(self, inputs):
+ """Forward function of PSP module."""
+
+ x = inputs[-1]
+ psp_outs = [x]
+ psp_outs.extend(self.psp_modules(x))
+ psp_outs = torch.cat(psp_outs, dim=1)
+ output = self.bottleneck(psp_outs)
+ return output
+
+ def forward(self, inputs):
+ """Forward function.
+ inputs = {x_96, x_192, x_384, x_768}
+ """
+
+ laterals = [
+ lateral_conv(inputs[i])
+ for i, lateral_conv in enumerate(self.lateral_convs)
+ ]
+
+ laterals.append(self.psp_forward(inputs))
+
+ # build top-down path
+ used_backbone_levels = len(laterals)
+ for i in range(used_backbone_levels - 1, 0, -1):
+ prev_shape = laterals[i - 1].shape[2:]
+ laterals[i - 1] = laterals[i - 1] + F.interpolate(
+ laterals[i].float(),
+ size = prev_shape,
+ mode='bilinear',
+ align_corners = False
+ ).to(torch.bfloat16)
+
+ # build outputs
+ fpn_outs = [
+ self.fpn_convs[i](laterals[i])
+ for i in range(used_backbone_levels - 1)
+ ]
+
+ # append psp feature
+ fpn_outs.append(laterals[-1])
+ for i in range(used_backbone_levels - 1, 0, -1):
+ fpn_outs[i] = F.interpolate(
+ fpn_outs[i].float(),
+ size=fpn_outs[0].shape[2:],
+ mode='bilinear',
+ align_corners=False).to(torch.bfloat16)
+ fpn_outs = torch.cat(fpn_outs, dim=1)
+ output = self.fpn_bottleneck(fpn_outs)
+
+ return output
\ No newline at end of file
diff --git a/src/models/utils/utils.py b/src/models/utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d7af1f68fa0ce0a48d11a708d53aa20aa8f78ba2
--- /dev/null
+++ b/src/models/utils/utils.py
@@ -0,0 +1,108 @@
+# 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.
+
+import torch
+import torch.nn as nn
+
+
+def position_grid_to_embed(pos_grid: torch.Tensor, embed_dim: int, omega_0: float = 100) -> torch.Tensor:
+ """
+ Convert 2D position grid (HxWx2) to sinusoidal embeddings (HxWxC)
+
+ Args:
+ pos_grid: Tensor of shape (H, W, 2) containing 2D coordinates
+ embed_dim: Output channel dimension for embeddings
+
+ Returns:
+ Tensor of shape (H, W, embed_dim) with positional embeddings
+ """
+ H, W, grid_dim = pos_grid.shape
+ assert grid_dim == 2
+ pos_flat = pos_grid.reshape(-1, grid_dim) # Flatten to (H*W, 2)
+
+ # Process x and y coordinates separately
+ emb_x = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 0], omega_0=omega_0) # [1, H*W, D/2]
+ emb_y = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 1], omega_0=omega_0) # [1, H*W, D/2]
+
+ # Combine and reshape
+ emb = torch.cat([emb_x, emb_y], dim=-1) # [1, H*W, D]
+
+ return emb.view(H, W, embed_dim) # [H, W, D]
+
+
+def make_sincos_pos_embed(embed_dim: int, pos: torch.Tensor, omega_0: float = 100) -> torch.Tensor:
+ """
+ This function generates a 1D positional embedding from a given grid using sine and cosine functions.
+
+ Args:
+ - embed_dim: The embedding dimension.
+ - pos: The position to generate the embedding from.
+
+ Returns:
+ - emb: The generated 1D positional embedding.
+ """
+ assert embed_dim % 2 == 0
+ omega = torch.arange(embed_dim // 2, dtype=torch.double, device=pos.device)
+ omega /= embed_dim / 2.0
+ omega = 1.0 / omega_0**omega # (D/2,)
+
+ pos = pos.reshape(-1) # (M,)
+ out = torch.einsum("m,d->md", pos, omega) # (M, D/2), outer product
+
+ emb_sin = torch.sin(out) # (M, D/2)
+ emb_cos = torch.cos(out) # (M, D/2)
+
+ emb = torch.cat([emb_sin, emb_cos], dim=1) # (M, D)
+ return emb.float()
+
+
+# Inspired by https://github.com/microsoft/moge
+
+
+def create_uv_grid(
+ width: int, height: int, aspect_ratio: float = None, dtype: torch.dtype = None, device: torch.device = None
+) -> torch.Tensor:
+ """
+ Create a normalized UV grid of shape (width, height, 2).
+
+ The grid spans horizontally and vertically according to an aspect ratio,
+ ensuring the top-left corner is at (-x_span, -y_span) and the bottom-right
+ corner is at (x_span, y_span), normalized by the diagonal of the plane.
+
+ Args:
+ width (int): Number of points horizontally.
+ height (int): Number of points vertically.
+ aspect_ratio (float, optional): Width-to-height ratio. Defaults to width/height.
+ dtype (torch.dtype, optional): Data type of the resulting tensor.
+ device (torch.device, optional): Device on which the tensor is created.
+
+ Returns:
+ torch.Tensor: A (width, height, 2) tensor of UV coordinates.
+ """
+ # Derive aspect ratio if not explicitly provided
+ if aspect_ratio is None:
+ aspect_ratio = float(width) / float(height)
+
+ # Compute normalized spans for X and Y
+ diag_factor = (aspect_ratio**2 + 1.0) ** 0.5
+ span_x = aspect_ratio / diag_factor
+ span_y = 1.0 / diag_factor
+
+ # Establish the linspace boundaries
+ left_x = -span_x * (width - 1) / width
+ right_x = span_x * (width - 1) / width
+ top_y = -span_y * (height - 1) / height
+ bottom_y = span_y * (height - 1) / height
+
+ # Generate 1D coordinates
+ x_coords = torch.linspace(left_x, right_x, steps=width, dtype=dtype, device=device)
+ y_coords = torch.linspace(top_y, bottom_y, steps=height, dtype=dtype, device=device)
+
+ # Create 2D meshgrid (width x height) and stack into UV
+ uu, vv = torch.meshgrid(x_coords, y_coords, indexing="xy")
+ uv_grid = torch.stack((uu, vv), dim=-1)
+
+ return uv_grid
diff --git a/src/models/utils/write_depth_image.py b/src/models/utils/write_depth_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..0d7b7e9a3ad8cba41a75407cfb5254b8dc74f2e5
--- /dev/null
+++ b/src/models/utils/write_depth_image.py
@@ -0,0 +1,5 @@
+import scipy.io
+
+def write_depth_image(data, filename): # kernel = 1st or 2nd
+ scipy.io.savemat(filename, {'name':data})
+ return 1
diff --git a/src/utils/__init__.py b/src/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b0707ca57ec89fc5f5cb1a023135eeb756a8e1e
--- /dev/null
+++ b/src/utils/__init__.py
@@ -0,0 +1,5 @@
+from src.utils.instantiators import instantiate_callbacks, instantiate_loggers
+from src.utils.logging_utils import log_hyperparameters
+from src.utils.pylogger import RankedLogger
+from src.utils.rich_utils import enforce_tags, print_config_tree
+from src.utils.utils import extras, get_metric_value, task_wrapper
diff --git a/src/utils/__pycache__/__init__.cpython-310.pyc b/src/utils/__pycache__/__init__.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ac3a3ea133c7a1ca57d09a952ff1a1f07a978ba7
Binary files /dev/null and b/src/utils/__pycache__/__init__.cpython-310.pyc differ
diff --git a/src/utils/__pycache__/instantiators.cpython-310.pyc b/src/utils/__pycache__/instantiators.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..80ce973eed017c532d25e0246e2e695c26f56f90
Binary files /dev/null and b/src/utils/__pycache__/instantiators.cpython-310.pyc differ
diff --git a/src/utils/__pycache__/logging_utils.cpython-310.pyc b/src/utils/__pycache__/logging_utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9113dedd12b2714384099756f38e457cbbd742ac
Binary files /dev/null and b/src/utils/__pycache__/logging_utils.cpython-310.pyc differ
diff --git a/src/utils/__pycache__/pylogger.cpython-310.pyc b/src/utils/__pycache__/pylogger.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cde314f222009cddf1e7c439891d7f5523457036
Binary files /dev/null and b/src/utils/__pycache__/pylogger.cpython-310.pyc differ
diff --git a/src/utils/__pycache__/rich_utils.cpython-310.pyc b/src/utils/__pycache__/rich_utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..0f8d7e69414c812a64baa28036d128860794a5b0
Binary files /dev/null and b/src/utils/__pycache__/rich_utils.cpython-310.pyc differ
diff --git a/src/utils/__pycache__/utils.cpython-310.pyc b/src/utils/__pycache__/utils.cpython-310.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9cb2906b829fbb280f25134d00e5dc1c38ddb160
Binary files /dev/null and b/src/utils/__pycache__/utils.cpython-310.pyc differ
diff --git a/src/utils/instantiators.py b/src/utils/instantiators.py
new file mode 100644
index 0000000000000000000000000000000000000000..82b9278a465d39565942f862442ebe79549825d7
--- /dev/null
+++ b/src/utils/instantiators.py
@@ -0,0 +1,56 @@
+from typing import List
+
+import hydra
+from lightning import Callback
+from lightning.pytorch.loggers import Logger
+from omegaconf import DictConfig
+
+from src.utils import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+def instantiate_callbacks(callbacks_cfg: DictConfig) -> List[Callback]:
+ """Instantiates callbacks from config.
+
+ :param callbacks_cfg: A DictConfig object containing callback configurations.
+ :return: A list of instantiated callbacks.
+ """
+ callbacks: List[Callback] = []
+
+ if not callbacks_cfg:
+ log.warning("No callback configs found! Skipping..")
+ return callbacks
+
+ if not isinstance(callbacks_cfg, DictConfig):
+ raise TypeError("Callbacks config must be a DictConfig!")
+
+ for _, cb_conf in callbacks_cfg.items():
+ if isinstance(cb_conf, DictConfig) and "_target_" in cb_conf:
+ log.info(f"Instantiating callback <{cb_conf._target_}>")
+ callbacks.append(hydra.utils.instantiate(cb_conf))
+
+ return callbacks
+
+
+def instantiate_loggers(logger_cfg: DictConfig) -> List[Logger]:
+ """Instantiates loggers from config.
+
+ :param logger_cfg: A DictConfig object containing logger configurations.
+ :return: A list of instantiated loggers.
+ """
+ logger: List[Logger] = []
+
+ if not logger_cfg:
+ log.warning("No logger configs found! Skipping...")
+ return logger
+
+ if not isinstance(logger_cfg, DictConfig):
+ raise TypeError("Logger config must be a DictConfig!")
+
+ for _, lg_conf in logger_cfg.items():
+ if isinstance(lg_conf, DictConfig) and "_target_" in lg_conf:
+ log.info(f"Instantiating logger <{lg_conf._target_}>")
+ logger.append(hydra.utils.instantiate(lg_conf))
+
+ return logger
diff --git a/src/utils/logging_utils.py b/src/utils/logging_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..360abcdceec82e551995f756ce6ec3b2d06ae641
--- /dev/null
+++ b/src/utils/logging_utils.py
@@ -0,0 +1,57 @@
+from typing import Any, Dict
+
+from lightning_utilities.core.rank_zero import rank_zero_only
+from omegaconf import OmegaConf
+
+from src.utils import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+@rank_zero_only
+def log_hyperparameters(object_dict: Dict[str, Any]) -> None:
+ """Controls which config parts are saved by Lightning loggers.
+
+ Additionally saves:
+ - Number of model parameters
+
+ :param object_dict: A dictionary containing the following objects:
+ - `"cfg"`: A DictConfig object containing the main config.
+ - `"model"`: The Lightning model.
+ - `"trainer"`: The Lightning trainer.
+ """
+ hparams = {}
+
+ cfg = OmegaConf.to_container(object_dict["cfg"])
+ model = object_dict["model"]
+ trainer = object_dict["trainer"]
+
+ if not trainer.logger:
+ log.warning("Logger not found! Skipping hyperparameter logging...")
+ return
+
+ hparams["model"] = cfg["model"]
+
+ # save number of model parameters
+ hparams["model/params/total"] = sum(p.numel() for p in model.parameters())
+ hparams["model/params/trainable"] = sum(
+ p.numel() for p in model.parameters() if p.requires_grad
+ )
+ hparams["model/params/non_trainable"] = sum(
+ p.numel() for p in model.parameters() if not p.requires_grad
+ )
+
+ hparams["data"] = cfg["data"]
+ hparams["trainer"] = cfg["trainer"]
+
+ hparams["callbacks"] = cfg.get("callbacks")
+ hparams["extras"] = cfg.get("extras")
+
+ hparams["task_name"] = cfg.get("task_name")
+ hparams["tags"] = cfg.get("tags")
+ hparams["ckpt_path"] = cfg.get("ckpt_path")
+ hparams["seed"] = cfg.get("seed")
+
+ # send hparams to all loggers
+ for logger in trainer.loggers:
+ logger.log_hyperparams(hparams)
diff --git a/src/utils/loss.py b/src/utils/loss.py
new file mode 100644
index 0000000000000000000000000000000000000000..87f4a548db5de5374294a75ff7ab86bc423e6912
--- /dev/null
+++ b/src/utils/loss.py
@@ -0,0 +1,298 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+import numpy as np
+from math import exp
+
+class MultiScaleDerivativeLoss(nn.Module):
+ def __init__(self, operator='scharr', p=1, reduction='mean', normalize_input=False, num_scales=4):
+ """
+ operator: 'scharr' (一阶) or 'laplace' (二阶)
+ p: 1 for L1, 2 for L2
+ reduction: 'mean' or 'sum'
+ normalize_input: whether to normalize input vectors (for normals)
+ num_scales: number of scales in the pyramid (e.g., 4 = 原图, 1/2, 1/4, 1/8)
+ """
+ super().__init__()
+ assert operator in ['scharr', 'laplace']
+ assert p in [1, 2]
+ assert reduction in ['mean', 'sum']
+ assert num_scales >= 1
+
+ self.operator = operator
+ self.p = p
+ self.reduction = reduction
+ self.normalize_input = normalize_input
+ self.num_scales = num_scales
+
+ def forward(self, pred, gt):
+ """
+ pred, gt: [B, C, H, W] tensors
+ """
+ pred_pyramid = self._build_pyramid(pred)
+ gt_pyramid = self._build_pyramid(gt)
+
+ total_loss = 0.0
+
+ for pred_i, gt_i in zip(pred_pyramid, gt_pyramid):
+ if self.normalize_input:
+ pred_i = F.normalize(pred_i, dim=1)
+ gt_i = F.normalize(gt_i, dim=1)
+
+ grad_pred = self._compute_gradient(pred_i)
+ grad_gt = self._compute_gradient(gt_i)
+
+ diff = grad_pred - grad_gt
+ if self.p == 1:
+ diff = torch.abs(diff)
+ else:
+ diff = diff ** 2
+
+ if self.reduction == 'mean':
+ total_loss += diff.mean()
+ else:
+ total_loss += diff.sum()
+
+ return total_loss / self.num_scales
+
+ def _build_pyramid(self, img):
+ """Construct a multi-scale pyramid from input image"""
+ pyramid = [img]
+ for i in range(1, self.num_scales):
+ scale = 0.5 ** i
+ img = F.interpolate(img, scale_factor=scale, mode='bicubic', align_corners=False, recompute_scale_factor=True,antialias=True)
+ pyramid.append(img)
+ return pyramid
+
+ def _compute_gradient(self, img):
+ B, C, H, W = img.shape
+ device = img.device
+
+ if self.operator == 'scharr':
+ kernel_x = torch.tensor([[[-3., 0., 3.],
+ [-10., 0., 10.],
+ [-3., 0., 3.]]], device=device) / 16.0
+ kernel_y = torch.tensor([[[-3., -10., -3.],
+ [0., 0., 0.],
+ [3., 10., 3.]]], device=device) / 16.0
+ kernel_x = kernel_x.unsqueeze(0).expand(C, 1, 3, 3)
+ kernel_y = kernel_y.unsqueeze(0).expand(C, 1, 3, 3)
+
+ grad_x = F.conv2d(img, kernel_x, padding=1, groups=C)
+ grad_y = F.conv2d(img, kernel_y, padding=1, groups=C)
+ return torch.cat([grad_x, grad_y], dim=1) # [B, 2C, H, W]
+
+ elif self.operator == 'laplace':
+ kernel = torch.tensor([[[0., 1., 0.],
+ [1., -4., 1.],
+ [0., 1., 0.]]], device=device)
+ kernel = kernel.unsqueeze(0).expand(C, 1, 3, 3)
+ return F.conv2d(img, kernel, padding=1, groups=C) # [B, C, H, W]
+
+class CosineLoss(torch.nn.Module):
+ def __init__(self):
+ super(CosineLoss, self).__init__()
+
+ def forward(self, N, N_hat):
+ """
+ N: 真实法向量, 形状 (B, C, H, W)
+ N_hat: 预测法向量, 形状应与 N 相同
+ """
+ # 创建非零 mask（按像素维度求L2范数）
+ _,_,H,W = N.shape
+ mask = (N.norm(p=2, dim=1, keepdim=True) > 0) # shape: (B, 1, H, W)，True表示N非零
+ mse = F.mse_loss(N, N_hat, reduction='mean') * H * W /2048
+ dot_product = torch.sum(N * N_hat, dim=1, keepdim=True) # shape: (B, 1, H, W)
+
+ # 仅在非零区域计算 loss
+ loss = 1 - dot_product
+ loss = loss[mask] # 只取非零像素位置
+ return loss.mean(), mse
+
+
+
+
+
+
+def gaussian(window_size, sigma):
+ gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])
+ return gauss/gauss.sum()
+
+def create_window(window_size, channel):
+ _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+ _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+ window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+ return window
+
+def _ssim(img1, img2, window, window_size, channel, size_average = True, stride=None):
+ mu1 = F.conv2d(img1, window, padding = (window_size-1)//2, groups = channel, stride=stride)
+ mu2 = F.conv2d(img2, window, padding = (window_size-1)//2, groups = channel, stride=stride)
+
+ mu1_sq = mu1.pow(2)
+ mu2_sq = mu2.pow(2)
+ mu1_mu2 = mu1*mu2
+
+ sigma1_sq = F.conv2d(img1*img1, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu1_sq
+ sigma2_sq = F.conv2d(img2*img2, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu2_sq
+ sigma12 = F.conv2d(img1*img2, window, padding = (window_size-1)//2, groups = channel, stride=stride) - mu1_mu2
+
+ C1 = 0.01**2
+ C2 = 0.03**2
+
+ ssim_map = ((2*mu1_mu2 + C1)*(2*sigma12 + C2))/((mu1_sq + mu2_sq + C1)*(sigma1_sq + sigma2_sq + C2))
+
+ if size_average:
+ return ssim_map.mean()
+ else:
+ return ssim_map.mean(1).mean(1).mean(1)
+
+class SSIM(torch.nn.Module):
+ def __init__(self, window_size = 3, size_average = True, stride=3):
+ super(SSIM, self).__init__()
+ self.window_size = window_size
+ self.size_average = size_average
+ self.channel = 1
+ self.stride = stride
+ self.window = create_window(window_size, self.channel)
+
+ def forward(self, img1, img2):
+ """
+ img1, img2: torch.Tensor([b,c,h,w])
+ """
+ (_, channel, _, _) = img1.size()
+
+ if channel == self.channel and self.window.data.type() == img1.data.type():
+ window = self.window
+ else:
+ window = create_window(self.window_size, channel)
+
+ if img1.is_cuda:
+ window = window.cuda(img1.get_device())
+ window = window.type_as(img1)
+
+ self.window = window
+ self.channel = channel
+
+
+ return _ssim(img1, img2, window, self.window_size, channel, self.size_average, stride=self.stride)
+
+
+def ssim(img1, img2, window_size = 11, size_average = True):
+ (_, channel, _, _) = img1.size()
+ window = create_window(window_size, channel)
+
+ if img1.is_cuda:
+ window = window.cuda(img1.get_device())
+ window = window.type_as(img1)
+
+ return _ssim(img1, img2, window, window_size, channel, size_average)
+
+
+
+
+class S3IM(torch.nn.Module):
+ def __init__(self, kernel_size=4, stride=4, repeat_time=10, patch_height=64, patch_width=32):
+ super(S3IM, self).__init__()
+ self.kernel_size = kernel_size
+ self.stride = stride
+ self.repeat_time = repeat_time
+ self.patch_height = patch_height
+ self.patch_width = patch_width
+ self.ssim_loss = SSIM(window_size=self.kernel_size, stride=self.stride)
+
+ def forward(self, src_vec, tar_vec):
+ """
+ Args:
+ src_vec: [B, N, C] e.g., [batch, pixels, channels]
+ tar_vec: [B, N, C]
+ Returns:
+ loss: scalar tensor
+ """
+ B, N, C = src_vec.shape
+ device = src_vec.device
+ patch_list_src, patch_list_tar = [], []
+
+ for b in range(B):
+ index_list = []
+ for i in range(self.repeat_time):
+ if i == 0:
+ tmp_index = torch.arange(N, device=device)
+ else:
+ tmp_index = torch.randperm(N, device=device)
+ index_list.append(tmp_index)
+
+ res_index = torch.cat(index_list) # [M * N]
+ tar_all = tar_vec[b][res_index] # [M*N, C]
+ src_all = src_vec[b][res_index] # [M*N, C]
+
+ # reshape into [1, C, H, W]
+ tar_patch = tar_all.permute(1, 0).reshape(1, C, self.patch_height, self.patch_width * self.repeat_time)
+ src_patch = src_all.permute(1, 0).reshape(1, C, self.patch_height, self.patch_width * self.repeat_time)
+
+ patch_list_tar.append(tar_patch)
+ patch_list_src.append(src_patch)
+
+ # Stack all batches: [B, C, H, W]
+ tar_tensor = torch.cat(patch_list_tar, dim=0)
+ src_tensor = torch.cat(patch_list_src, dim=0)
+
+ # 计算 batch-wise SSIM，输出为 [B]
+ ssim_scores = self.ssim_loss(src_tensor, tar_tensor)
+
+ # 损失为 1 - mean SSIM
+ loss = 1.0 - ssim_scores
+ return loss
+
+
+
+torch.manual_seed(0)
+
+# 假设每张图片提取出 64 x 64 个像素，每个像素 3 通道
+# H, W, C = 64, 32, 3
+# N = H * W
+# B = 4
+# # 随机生成两个图像特征向量：[N, C]
+# src_vec = torch.rand(B, N, C) # 模拟重建图像
+# tar_vec = torch.rand(B, N, C) # 模拟 ground truth 图像
+
+# # 初始化 S3IM 模块
+# s3im_loss_fn = S3IM(kernel_size=4, stride=4, repeat_time=10, patch_height=64, patch_width=32)
+
+# # 计算损失
+# loss = s3im_loss_fn(src_vec, tar_vec)
+
+def weighted_huber_loss(
+ input: torch.Tensor,
+ target: torch.Tensor,
+ weight: torch.Tensor, # 新增的置信度权重张量
+ reduction: str = 'mean',
+ delta: float = 1.0,
+) -> torch.Tensor:
+ # 广播对齐所有张量
+ expanded_input, expanded_target = torch.broadcast_tensors(input, target)
+ expanded_weight, _ = torch.broadcast_tensors(weight, input) # 确保权重可广播
+
+ # 计算逐元素误差
+ diff = expanded_input - expanded_target
+ abs_diff = torch.abs(diff)
+
+ # Huber损失分段计算
+ loss = torch.where(
+ abs_diff <= delta,
+ 0.5 * (diff ** 2),
+ delta * (abs_diff - 0.5 * delta)
+ )
+
+ # 应用权重
+ weighted_loss = expanded_weight * loss
+
+ # 汇总方式
+ if reduction == 'mean':
+ return torch.mean(weighted_loss)
+ elif reduction == 'sum':
+ return torch.sum(weighted_loss)
+ elif reduction == 'none':
+ return weighted_loss
+ else:
+ raise ValueError(f"Unsupported reduction: {reduction}")
\ No newline at end of file
diff --git a/src/utils/normal_process.py b/src/utils/normal_process.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc0bd636f11c4f73866e9d85aafd822db5f62987
--- /dev/null
+++ b/src/utils/normal_process.py
@@ -0,0 +1,45 @@
+import numpy as np
+import json
+import matplotlib.pyplot as plt
+
+
+
+
+def get_transform_matrix(json_file_path):
+ """
+ 根据 file_path 从 JSON 文件中查找对应的 transform_matrix。
+
+ :param json_file_path: JSON 文件路径
+ :param target_file_path: 目标 file_path，用于匹配
+ :return: 4×4 transform_matrix (NumPy 数组) 或 None（如果未找到）
+ """
+ # 读取 JSON 文件
+ # name_prefix = target_file_prefix + '_hdri' #光照不同不会影响相机参数，所以只取0
+ with open(json_file_path, "r") as f:
+ data = json.load(f)
+
+ # 遍历 JSON，查找目标 file_path 的 transform_matrix
+ for frame_list in data["frames"]:
+ return np.array(frame_list.get('transform_matrix')) # 遍历所有 HDRI 组
+
+def blender_world_normal_2_opengl_camera(normals_world: np.ndarray, c2w: np.ndarray, visualization = False) -> np.ndarray:
+ H, W, C = normals_world.shape
+ if C == 4:
+ normals_world = normals_world[..., :3]
+
+ R_c2w = c2w[:3, :3]
+ R_opencv = R_c2w.T
+
+ transformed_normals = normals_world.reshape(-1, 3).T
+ transformed_normals = R_opencv @ transformed_normals
+ transformed_normals = transformed_normals.T
+ transformed_normals = transformed_normals.reshape(H, W, 3)
+ if visualization:
+ plt.imshow(transformed_normals)
+ plt.axis('off')
+ plt.savefig('N_world.png',bbox_inches='tight', pad_inches=0)
+ transformed_normals = transformed_normals * 0.5 + 0.5
+ plt.imshow(transformed_normals)
+ plt.axis('off')
+ plt.savefig('N_camera.png',bbox_inches='tight', pad_inches=0)
+ return transformed_normals
\ No newline at end of file
diff --git a/src/utils/pylogger.py b/src/utils/pylogger.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4ee8675ebde11b2a43b0679a03cd88d9268bc71
--- /dev/null
+++ b/src/utils/pylogger.py
@@ -0,0 +1,51 @@
+import logging
+from typing import Mapping, Optional
+
+from lightning_utilities.core.rank_zero import rank_prefixed_message, rank_zero_only
+
+
+class RankedLogger(logging.LoggerAdapter):
+ """A multi-GPU-friendly python command line logger."""
+
+ def __init__(
+ self,
+ name: str = __name__,
+ rank_zero_only: bool = False,
+ extra: Optional[Mapping[str, object]] = None,
+ ) -> None:
+ """Initializes a multi-GPU-friendly python command line logger that logs on all processes
+ with their rank prefixed in the log message.
+
+ :param name: The name of the logger. Default is ``__name__``.
+ :param rank_zero_only: Whether to force all logs to only occur on the rank zero process. Default is `False`.
+ :param extra: (Optional) A dict-like object which provides contextual information. See `logging.LoggerAdapter`.
+ """
+ logger = logging.getLogger(name)
+ super().__init__(logger=logger, extra=extra)
+ self.rank_zero_only = rank_zero_only
+
+ def log(self, level: int, msg: str, rank: Optional[int] = None, *args, **kwargs) -> None:
+ """Delegate a log call to the underlying logger, after prefixing its message with the rank
+ of the process it's being logged from. If `'rank'` is provided, then the log will only
+ occur on that rank/process.
+
+ :param level: The level to log at. Look at `logging.__init__.py` for more information.
+ :param msg: The message to log.
+ :param rank: The rank to log at.
+ :param args: Additional args to pass to the underlying logging function.
+ :param kwargs: Any additional keyword args to pass to the underlying logging function.
+ """
+ if self.isEnabledFor(level):
+ msg, kwargs = self.process(msg, kwargs)
+ current_rank = getattr(rank_zero_only, "rank", None)
+ if current_rank is None:
+ raise RuntimeError("The `rank_zero_only.rank` needs to be set before use")
+ msg = rank_prefixed_message(msg, current_rank)
+ if self.rank_zero_only:
+ if current_rank == 0:
+ self.logger.log(level, msg, *args, **kwargs)
+ else:
+ if rank is None:
+ self.logger.log(level, msg, *args, **kwargs)
+ elif current_rank == rank:
+ self.logger.log(level, msg, *args, **kwargs)
diff --git a/src/utils/rich_utils.py b/src/utils/rich_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..aeec6806bb1e4a15a04b91b710a546231590ab14
--- /dev/null
+++ b/src/utils/rich_utils.py
@@ -0,0 +1,99 @@
+from pathlib import Path
+from typing import Sequence
+
+import rich
+import rich.syntax
+import rich.tree
+from hydra.core.hydra_config import HydraConfig
+from lightning_utilities.core.rank_zero import rank_zero_only
+from omegaconf import DictConfig, OmegaConf, open_dict
+from rich.prompt import Prompt
+
+from src.utils import pylogger
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+@rank_zero_only
+def print_config_tree(
+ cfg: DictConfig,
+ print_order: Sequence[str] = (
+ "data",
+ "model",
+ "callbacks",
+ "logger",
+ "trainer",
+ "paths",
+ "extras",
+ ),
+ resolve: bool = False,
+ save_to_file: bool = False,
+) -> None:
+ """Prints the contents of a DictConfig as a tree structure using the Rich library.
+
+ :param cfg: A DictConfig composed by Hydra.
+ :param print_order: Determines in what order config components are printed. Default is ``("data", "model",
+ "callbacks", "logger", "trainer", "paths", "extras")``.
+ :param resolve: Whether to resolve reference fields of DictConfig. Default is ``False``.
+ :param save_to_file: Whether to export config to the hydra output folder. Default is ``False``.
+ """
+ style = "dim"
+ tree = rich.tree.Tree("CONFIG", style=style, guide_style=style)
+
+ queue = []
+
+ # add fields from `print_order` to queue
+ for field in print_order:
+ queue.append(field) if field in cfg else log.warning(
+ f"Field '{field}' not found in config. Skipping '{field}' config printing..."
+ )
+
+ # add all the other fields to queue (not specified in `print_order`)
+ for field in cfg:
+ if field not in queue:
+ queue.append(field)
+
+ # generate config tree from queue
+ for field in queue:
+ branch = tree.add(field, style=style, guide_style=style)
+
+ config_group = cfg[field]
+ if isinstance(config_group, DictConfig):
+ branch_content = OmegaConf.to_yaml(config_group, resolve=resolve)
+ else:
+ branch_content = str(config_group)
+
+ branch.add(rich.syntax.Syntax(branch_content, "yaml"))
+
+ # print config tree
+ rich.print(tree)
+
+ # save config tree to file
+ if save_to_file:
+ with open(Path(cfg.paths.output_dir, "config_tree.log"), "w") as file:
+ rich.print(tree, file=file)
+
+
+@rank_zero_only
+def enforce_tags(cfg: DictConfig, save_to_file: bool = False) -> None:
+ """Prompts user to input tags from command line if no tags are provided in config.
+
+ :param cfg: A DictConfig composed by Hydra.
+ :param save_to_file: Whether to export tags to the hydra output folder. Default is ``False``.
+ """
+ if not cfg.get("tags"):
+ if "id" in HydraConfig().cfg.hydra.job:
+ raise ValueError("Specify tags before launching a multirun!")
+
+ log.warning("No tags provided in config. Prompting user to input tags...")
+ tags = Prompt.ask("Enter a list of comma separated tags", default="dev")
+ tags = [t.strip() for t in tags.split(",") if t != ""]
+
+ with open_dict(cfg):
+ cfg.tags = tags
+
+ log.info(f"Tags: {cfg.tags}")
+
+ if save_to_file:
+ with open(Path(cfg.paths.output_dir, "tags.log"), "w") as file:
+ rich.print(cfg.tags, file=file)
diff --git a/src/utils/utils.py b/src/utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..fc0a2502f37b777015afd8ce4214ddf86cdfac3b
--- /dev/null
+++ b/src/utils/utils.py
@@ -0,0 +1,121 @@
+import warnings
+from importlib.util import find_spec
+from typing import Any, Callable, Dict, Optional, Tuple
+
+from omegaconf import DictConfig
+
+from src.utils import pylogger, rich_utils
+
+log = pylogger.RankedLogger(__name__, rank_zero_only=True)
+
+
+def extras(cfg: DictConfig) -> None:
+ """Applies optional utilities before the task is started.
+
+ Utilities:
+ - Ignoring python warnings
+ - Setting tags from command line
+ - Rich config printing
+
+ :param cfg: A DictConfig object containing the config tree.
+ """
+ # return if no `extras` config
+ if not cfg.get("extras"):
+ log.warning("Extras config not found! ")
+ return
+
+ # disable python warnings
+ if cfg.extras.get("ignore_warnings"):
+ log.info("Disabling python warnings! ")
+ warnings.filterwarnings("ignore")
+
+ # prompt user to input tags from command line if none are provided in the config
+ if cfg.extras.get("enforce_tags"):
+ log.info("Enforcing tags! ")
+ rich_utils.enforce_tags(cfg, save_to_file=True)
+
+ # pretty print config tree using Rich library
+ if cfg.extras.get("print_config"):
+ log.info("Printing config tree with Rich! ")
+ rich_utils.print_config_tree(cfg, resolve=True, save_to_file=True)
+
+
+def task_wrapper(task_func: Callable) -> Callable:
+ """Optional decorator that controls the failure behavior when executing the task function.
+
+ This wrapper can be used to:
+ - make sure loggers are closed even if the task function raises an exception (prevents multirun failure)
+ - save the exception to a `.log` file
+ - mark the run as failed with a dedicated file in the `logs/` folder (so we can find and rerun it later)
+ - etc. (adjust depending on your needs)
+
+ Example:
+ ```
+ @utils.task_wrapper
+ def train(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+ ...
+ return metric_dict, object_dict
+ ```
+
+ :param task_func: The task function to be wrapped.
+
+ :return: The wrapped task function.
+ """
+
+ def wrap(cfg: DictConfig) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+ # execute the task
+ try:
+ metric_dict, object_dict = task_func(cfg=cfg)
+
+ # things to do if exception occurs
+ except Exception as ex:
+ # save exception to `.log` file
+ log.exception("")
+
+ # some hyperparameter combinations might be invalid or cause out-of-memory errors
+ # so when using hparam search plugins like Optuna, you might want to disable
+ # raising the below exception to avoid multirun failure
+ raise ex
+
+ # things to always do after either success or exception
+ finally:
+ # display output dir path in terminal
+ log.info(f"Output dir: {cfg.paths.output_dir}")
+
+ # always close wandb run (even if exception occurs so multirun won't fail)
+ if find_spec("wandb"): # check if wandb is installed
+ import wandb
+
+ if wandb.run:
+ log.info("Closing wandb!")
+ wandb.finish()
+
+ return metric_dict, object_dict
+
+ return wrap
+
+
+def get_metric_value(metric_dict: Dict[str, Any], metric_name: Optional[str]) -> Optional[float]:
+ """Safely retrieves value of the metric logged in LightningModule.
+
+ :param metric_dict: A dict containing metric values.
+ :param metric_name: If provided, the name of the metric to retrieve.
+ :return: If a metric name was provided, the value of the metric.
+ """
+ if not metric_name:
+ log.info("Metric name is None! Skipping metric value retrieval...")
+ return None
+
+ if metric_name not in metric_dict:
+ raise Exception(
+ f"Metric value not found! \n"
+ "Make sure metric name logged in LightningModule is correct!\n"
+ "Make sure `optimized_metric` name in `hparams_search` config is correct!"
+ )
+
+ metric_value = metric_dict[metric_name].item()
+ log.info(f"Retrieved metric value! <{metric_name}={metric_value}>")
+
+ return metric_value
+
+