# Copyright (c) Meta Platforms, Inc. and affiliates. # Modified for AI Toolkit by Ostris # All rights reserved. # # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. # https://raw.githubusercontent.com/facebookresearch/sapiens2/refs/heads/main/sapiens/backbones/standalone/sapiens2.py import math from typing import Any, List, Literal, Optional, Sequence, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor from torch.nn.init import trunc_normal_ from torch.utils.checkpoint import checkpoint from toolkit.paths import MODELS_PATH import os # ---------------------------------------------------------------------------- def to_2tuple(x): if isinstance(x, (str, bytes)): return (x, x) if isinstance(x, Sequence): x = tuple(x) if len(x) == 2: return x raise ValueError("Expected scalar or length-2 iterable") return (x, x) class RopePositionEmbedding(nn.Module): def __init__( self, embed_dim: int, *, num_heads: int, base: float | None = 100.0, min_period: float | None = None, max_period: float | None = None, normalize_coords: Literal["min", "max", "separate"] = "separate", shift_coords: float | None = None, jitter_coords: float | None = None, rescale_coords: float | None = None, dtype: torch.dtype | None = None, device: torch.device | None = None, ): super().__init__() assert embed_dim % (4 * num_heads) == 0 both_periods = min_period is not None and max_period is not None if (base is None and not both_periods) or (base is not None and both_periods): raise ValueError( "Either `base` or `min_period`+`max_period` must be provided." ) D_head = embed_dim // num_heads self.base = base self.min_period = min_period self.max_period = max_period self.D_head = D_head self.normalize_coords = normalize_coords self.shift_coords = shift_coords self.jitter_coords = jitter_coords self.rescale_coords = rescale_coords # Needs persistent=True because we do teacher.load_state_dict(student.state_dict()) to initialize the teacher self.dtype = dtype or torch.float32 # Don't rely on self.periods.dtype self.register_buffer( "periods", torch.empty(D_head // 4, device=device, dtype=self.dtype), persistent=True, ) self._init_weights() def forward(self, *, H: int, W: int) -> tuple[Tensor, Tensor]: device = self.periods.device dtype = self.dtype dd = {"device": device, "dtype": dtype} # Prepare coords in range [-1, +1] if self.normalize_coords == "max": max_HW = max(H, W) coords_h = torch.arange(0.5, H, **dd) / max_HW # [H] coords_w = torch.arange(0.5, W, **dd) / max_HW # [W] elif self.normalize_coords == "min": min_HW = min(H, W) coords_h = torch.arange(0.5, H, **dd) / min_HW # [H] coords_w = torch.arange(0.5, W, **dd) / min_HW # [W] elif self.normalize_coords == "separate": coords_h = torch.arange(0.5, H, **dd) / H # [H] coords_w = torch.arange(0.5, W, **dd) / W # [W] else: raise ValueError(f"Unknown normalize_coords: {self.normalize_coords}") coords = torch.stack( torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1 ) # [H, W, 2] coords = coords.flatten(0, 1) # [HW, 2] coords = 2.0 * coords - 1.0 # Shift range [0, 1] to [-1, +1] # Shift coords by adding a uniform value in [-shift, shift] if self.training and self.shift_coords is not None: shift_hw = torch.empty(2, **dd).uniform_( -self.shift_coords, self.shift_coords ) coords += shift_hw[None, :] # Jitter coords by multiplying the range [-1, 1] by a log-uniform value in [1/jitter, jitter] if self.training and self.jitter_coords is not None: jitter_max = np.log(self.jitter_coords) jitter_min = -jitter_max jitter_hw = torch.empty(2, **dd).uniform_(jitter_min, jitter_max).exp() coords *= jitter_hw[None, :] # Rescale coords by multiplying the range [-1, 1] by a log-uniform value in [1/rescale, rescale] if self.training and self.rescale_coords is not None: rescale_max = np.log(self.rescale_coords) rescale_min = -rescale_max rescale_hw = torch.empty(1, **dd).uniform_(rescale_min, rescale_max).exp() coords *= rescale_hw # Prepare angles and sin/cos angles = ( 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] ) # [HW, 2, D//4] angles = angles.flatten(1, 2) # [HW, D//2] angles = angles.tile(2) # [HW, D] cos = torch.cos(angles) # [HW, D] sin = torch.sin(angles) # [HW, D] return (sin, cos) # 2 * [HW, D] def _init_weights(self): device = self.periods.device dtype = self.dtype if self.base is not None: periods = self.base ** ( 2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2) ) # [D//4] else: base = self.max_period / self.min_period exponents = torch.linspace( 0, 1, self.D_head // 4, device=device, dtype=dtype ) # [D//4] range [0, 1] periods = base**exponents # range [1, max_period / min_period] periods = periods / base # range [min_period / max_period, 1] periods = periods * self.max_period # range [min_period, max_period] self.periods.data = periods # ------------------------------------------------------------------------------- class Tokenizer(nn.Module): """Stacked window self‑attention that emits one token per window by re‑using TransformerEncoderLayer blocks.""" def __init__( self, embed_dims: int, window_size: int = 4, num_heads: int = 4, num_tokenizer_layers: int = 1, qkv_bias: bool = True, use_qk_norm: bool = False, chunk_size: int = 1024, # max windows per chunk ): super().__init__() self.ws = window_size self.chunk_size = chunk_size # local absolute positional embeddings for [CLS] + patch tokens self.local_pos_embed = nn.Parameter( torch.zeros(1, 1 + window_size * window_size, embed_dims) ) trunc_normal_(self.local_pos_embed, std=0.02) # build N identical TransformerEncoderLayer blocks self.blocks = nn.ModuleList( [ TransformerEncoderLayer2( embed_dims=embed_dims, num_heads=num_heads, feedforward_channels=embed_dims * 4, # standard FFN size qkv_bias=qkv_bias, use_qk_norm=use_qk_norm, ) for _ in range(num_tokenizer_layers) ] ) # shared CLS token for pooling self.w_cls = nn.Parameter(torch.zeros(1, 1, embed_dims)) trunc_normal_(self.w_cls, std=0.02) self.gradient_checkpointing = False def forward( self, x: torch.Tensor, hw: Tuple[int, int], ) -> Tuple[torch.Tensor, Tuple[int, int]]: """Args: x : B, N, C (N = H*W) hw : (H, W) before reduction Returns: x_ : B, (H/ws)*(W/ws), C hw_: (H/ws, W/ws) """ B, N, C = x.shape H, W = hw ws = self.ws assert H % ws == 0 and W % ws == 0, ( f"Image size {H}×{W} must be divisible by window {ws}." ) # reshape tokens → non‑overlapping windows x = x.view(B, H, W, C) ph, pw = H // ws, W // ws ## ints in eager mode ph, pw = int(ph), int(pw) ## ints in scripting mode x = x.view(B, ph, ws, pw, ws, C) # B, H/ws, ws, W/ws, ws, C x = x.permute(0, 1, 3, 2, 4, 5) # B, H/ws, W/ws, ws, ws, C x = x.contiguous().view(B * ph * pw, ws * ws, C) # (B*H/ws*W/ws), ws², C)) total_windows = x.size(0) chunk_size = int(min(self.chunk_size, total_windows)) token_out = x.new_empty(total_windows, C) use_ckpt = torch.is_grad_enabled() and self.gradient_checkpointing def _run_blocks(t: torch.Tensor) -> torch.Tensor: for blk in self.blocks: t = blk(t) return t for i in range(0, total_windows, chunk_size): chunk = x[i : i + chunk_size] # (m, ws², C) m = chunk.size(0) cls = self.w_cls.expand(m, -1, -1) # (m, 1, C) chunk = torch.cat([cls, chunk], dim=1) # (m, 1+ws², C) chunk = chunk + self.local_pos_embed # add local PE if use_ckpt: chunk = checkpoint(_run_blocks, chunk, use_reentrant=False) else: chunk = _run_blocks(chunk) token_out[i : i + m] = chunk[:, 0] # take CLS out token = token_out.view(B, ph * pw, C) # (B, (H/ws)*(W return token, (ph, pw) # ------------------------------------------------------------------------------- class GroupedQueryAttention(nn.Module): def __init__( self, embed_dims, num_heads, num_kv_heads=None, input_dims=None, attn_drop=0.0, proj_drop=0.0, qkv_bias=True, qk_scale=None, proj_bias=True, use_qk_norm=True, v_shortcut=False, layer_scale_init_value=0.0, ): super().__init__() # Core dims self.embed_dims = embed_dims self.num_heads = num_heads self.num_kv_heads = num_kv_heads or num_heads assert self.num_heads % self.num_kv_heads == 0, ( "num_kv_heads must divide num_heads" ) self.head_dim = embed_dims // num_heads self.input_dims = input_dims or embed_dims # Features self.attn_drop = attn_drop self.v_shortcut = v_shortcut self.use_qk_norm = use_qk_norm # Attention operation selection if qk_scale is not None: scale = qk_scale else: scale = self.head_dim**-0.5 assert qk_scale is None, "qk_scale is not supported" self.attn_op = F.scaled_dot_product_attention # Q/K/V projections self.wq = nn.Linear(self.input_dims, embed_dims, bias=qkv_bias) self.wk = nn.Linear( self.input_dims, self.num_kv_heads * self.head_dim, bias=qkv_bias ) self.wv = nn.Linear( self.input_dims, self.num_kv_heads * self.head_dim, bias=qkv_bias ) if self.use_qk_norm: self.q_norm = nn.RMSNorm(self.head_dim, eps=1e-6) self.k_norm = nn.RMSNorm(self.head_dim, eps=1e-6) # Output projection + dropout self.proj = nn.Linear(embed_dims, embed_dims, bias=proj_bias) self.proj_drop = nn.Dropout(proj_drop) # Optional LayerScale if layer_scale_init_value > 0: self.gamma = LayerScale(embed_dims, scale=layer_scale_init_value) else: self.gamma = nn.Identity() def apply_rope( self, q: Tensor, k: Tensor, rope: Tensor | Tuple[Tensor, Tensor] ) -> Tuple[Tensor, Tensor]: # All operations will use the dtype of rope, the output is cast back to the dtype of q and k q_dtype = q.dtype k_dtype = k.dtype sin, cos = rope rope_dtype = sin.dtype q = q.to(dtype=rope_dtype) k = k.to(dtype=rope_dtype) N = q.shape[-2] prefix = N - sin.shape[-2] ## extra tokens assert prefix >= 0 q_prefix = q[:, :, :prefix, :] q = self._rope_apply(q[:, :, prefix:, :], sin, cos) # [B, head, hw, D//head] q = torch.cat((q_prefix, q), dim=-2) # [B, head, N, D//head] k_prefix = k[:, :, :prefix, :] k = self._rope_apply(k[:, :, prefix:, :], sin, cos) # [B, head, hw, D//head] k = torch.cat((k_prefix, k), dim=-2) # [B, head, N, D//head] q = q.to(dtype=q_dtype) k = k.to(dtype=k_dtype) return q, k def _rope_rotate_half(self, x: Tensor) -> Tensor: # x: [ x0 x1 x2 x3 x4 x5] # out: [-x3 -x4 -x5 x0 x1 x2] x1, x2 = x.chunk(2, dim=-1) return torch.cat([-x2, x1], dim=-1) def _rope_apply(self, x: Tensor, sin: Tensor, cos: Tensor) -> Tensor: # x: [..., D], eg [x0, x1, x2, x3, x4, x5] # sin: [..., D], eg [sin0, sin1, sin2, sin0, sin1, sin2] # cos: [..., D], eg [cos0, cos1, cos2, cos0, cos1, cos2] return (x * cos) + (self._rope_rotate_half(x) * sin) def forward(self, x, rope=None): B, N, _ = x.shape # Q: (B, N, num_heads, head_dim) q = self.wq(x).view(B, N, self.num_heads, self.head_dim) # K/V: (B, N, num_kv_heads, head_dim) k = self.wk(x).view(B, N, self.num_kv_heads, self.head_dim) v = self.wv(x).view(B, N, self.num_kv_heads, self.head_dim) # (B, heads, N, head_dim) q = q.permute(0, 2, 1, 3) k = k.permute(0, 2, 1, 3) v = v.permute(0, 2, 1, 3) if self.use_qk_norm: q = self.q_norm(q) k = self.k_norm(k) # Repeat KV heads if group ratio >1 if self.num_kv_heads != self.num_heads: factor = self.num_heads // self.num_kv_heads k = k.repeat_interleave(factor, dim=1) v = v.repeat_interleave(factor, dim=1) if rope is not None: q, k = self.apply_rope(q, k, rope) # Scaled dot-product attention attn_out = self.attn_op( q, k, v, dropout_p=self.attn_drop if self.training else 0.0 ) # (B, num_heads, N, head_dim) # Merge heads -> (B, N, embed_dims) out = attn_out.permute(0, 2, 1, 3).reshape(B, N, self.embed_dims) # Output projection + drop + layer scale out = self.proj(out) out = self.gamma(self.proj_drop(out)) # Optional V-shortcut (only when MQA) if self.v_shortcut and self.num_kv_heads == 1: raise NotImplementedError return out # ------------------------------------------------------------------------------- class TransformerEncoderLayer2(nn.Module): def __init__( self, embed_dims, num_heads, num_kv_heads=None, feedforward_channels=None, drop_rate=0.0, attn_drop_rate=0.0, layer_scale_init_value=0.0, use_qk_norm=True, qkv_bias=True, ): super(TransformerEncoderLayer2, self).__init__() self.embed_dims = embed_dims self.ln1 = nn.RMSNorm(self.embed_dims, eps=1e-6) self.attn = GroupedQueryAttention( embed_dims=embed_dims, num_heads=num_heads, num_kv_heads=num_kv_heads, attn_drop=attn_drop_rate, proj_drop=drop_rate, qkv_bias=qkv_bias, layer_scale_init_value=layer_scale_init_value, use_qk_norm=use_qk_norm, ) self.ln2 = nn.RMSNorm(self.embed_dims, eps=1e-6) self.ffn = SwiGLUFFN( embed_dims=embed_dims, feedforward_channels=feedforward_channels, ) @property def norm1(self): return self.ln1 @property def norm2(self): return self.ln2 def forward(self, x, rope=None): x = x + self.attn(self.ln1(x), rope=rope) x = self.ffn(self.ln2(x), identity=x) return x ##----------------------------------- class Sapiens2(nn.Module): arch_zoo = { **dict.fromkeys( ["sapiens2_0.1b"], { "embed_dims": 768, "num_layers": 12, "num_heads": 12, "feedforward_channels": 768 * 4, "num_tokenizer_layers": 2, }, ), **dict.fromkeys( ["sapiens2_0.4b"], { "embed_dims": 1024, "num_layers": 24, "num_heads": 16, "feedforward_channels": 1024 * 4, "num_tokenizer_layers": 2, }, ), **dict.fromkeys( ["sapiens2_0.8b"], { "embed_dims": 1280, "num_layers": 32, "num_heads": 16, "feedforward_channels": 1280 * 4, "num_tokenizer_layers": 3, }, ), **dict.fromkeys( ["sapiens2_1b"], { "embed_dims": 1536, "num_layers": 40, "num_heads": 24, "feedforward_channels": 1536 * 4, "num_tokenizer_layers": 4, }, ), **dict.fromkeys( ["sapiens2_5b"], { "embed_dims": 2432, "num_layers": 56, "num_heads": 32, "feedforward_channels": 2432 * 4, "num_tokenizer_layers": 6, }, ), } num_extra_tokens = 1 # class token OUT_TYPES = {"raw", "cls_token", "featmap"} _supports_gradient_checkpointing = True def __init__( self, arch="sapiens2_1b", img_size=(1024, 768), patch_size=16, in_channels=3, out_indices=-1, drop_rate=0.0, window_size=4, use_tokenizer=False, ## 4k resolution use_qk_norm=True, qkv_bias=True, final_norm=True, out_type="raw", with_cls_token=True, layer_scale_init_value=1e-4, ## non zero init to activate layerscale frozen_stages=-1, patch_cfg=dict(), layer_cfgs=dict(), pos_embed_rope_base: float = 100.0, pos_embed_rope_min_period: float | None = None, pos_embed_rope_max_period: float | None = None, pos_embed_rope_normalize_coords: Literal["min", "max", "separate"] = "separate", pos_embed_rope_shift_coords: float | None = None, pos_embed_rope_jitter_coords: float | None = None, pos_embed_rope_rescale_coords: float | None = None, pos_embed_rope_dtype: str = "bf16", n_storage_tokens: int = 8, ): super().__init__() arch = arch.lower() assert arch in set(self.arch_zoo), ( f"Arch {arch} is not in default archs {set(self.arch_zoo)}" ) self.arch_settings = self.arch_zoo[arch] self.embed_dims = self.arch_settings["embed_dims"] self.num_layers = self.arch_settings["num_layers"] self.patch_size = patch_size self.window_size = window_size img_size = to_2tuple(img_size) encoder_img_size = ( (img_size[0] // window_size, img_size[1] // window_size) if use_tokenizer else img_size ) self.img_size = to_2tuple(encoder_img_size) # Set patch embedding _patch_cfg = dict( in_channels=in_channels, input_size=self.img_size, embed_dims=self.embed_dims, kernel_size=patch_size, stride=patch_size, bias=True, ) _patch_cfg.update(patch_cfg) self.patch_embed = PatchEmbed(**_patch_cfg) self.patch_resolution = self.patch_embed.init_out_size num_patches = self.patch_resolution[0] * self.patch_resolution[1] self.rope_embed = RopePositionEmbedding( embed_dim=self.embed_dims, num_heads=self.arch_settings["num_heads"], base=pos_embed_rope_base, min_period=pos_embed_rope_min_period, max_period=pos_embed_rope_max_period, normalize_coords=pos_embed_rope_normalize_coords, shift_coords=pos_embed_rope_shift_coords, jitter_coords=pos_embed_rope_jitter_coords, rescale_coords=pos_embed_rope_rescale_coords, dtype=torch.bfloat16 if pos_embed_rope_dtype == "bf16" else torch.float32, ) # Set out type if out_type not in self.OUT_TYPES: raise ValueError( f"Unsupported `out_type` {out_type}, please " f"choose from {self.OUT_TYPES}" ) self.out_type = out_type if use_tokenizer == True: self.tokenizer = Tokenizer( embed_dims=self.embed_dims, window_size=self.window_size, num_heads=self.arch_settings["num_heads"], num_tokenizer_layers=self.arch_settings["num_tokenizer_layers"], qkv_bias=True, use_qk_norm=False, ) else: self.tokenizer = None # Set cls + storage tokens self.with_cls_token = with_cls_token if with_cls_token: self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dims)) elif out_type != "cls_token": self.cls_token = None self.num_extra_tokens = 0 else: raise ValueError('with_cls_token must be True when `out_type="cls_token"`.') ## registers self.n_storage_tokens = int(n_storage_tokens) self.storage_tokens = ( nn.Parameter(torch.zeros(1, self.n_storage_tokens, self.embed_dims)) if self.n_storage_tokens > 0 else None ) # how many non-patch tokens are at the front self.num_extra_tokens = ( 1 if self.cls_token is not None else 0 ) + self.n_storage_tokens if isinstance(out_indices, int): out_indices = [out_indices] assert isinstance(out_indices, Sequence), ( f'"out_indices" must by a sequence or int, get {type(out_indices)} instead.' ) for i, index in enumerate(out_indices): if index < 0: out_indices[i] = self.num_layers + index assert 0 <= out_indices[i] <= self.num_layers, ( f"Invalid out_indices {index}" ) self.out_indices = out_indices self.blocks = nn.Sequential() if isinstance(layer_cfgs, dict): layer_cfgs = [layer_cfgs] * self.num_layers mhsa_early, mhsa_late = 8, 8 for i in range(self.num_layers): if i < mhsa_early or i >= self.num_layers - mhsa_late: num_kv_heads = None ## use MHSA else: num_kv_heads = self.arch_settings["num_heads"] // 2 # Use GQA _layer_cfg = dict( embed_dims=self.embed_dims, num_heads=self.arch_settings["num_heads"], num_kv_heads=num_kv_heads, feedforward_channels=self.arch_settings["feedforward_channels"], use_qk_norm=use_qk_norm, layer_scale_init_value=layer_scale_init_value, drop_rate=drop_rate, qkv_bias=qkv_bias, ) _layer_cfg.update(layer_cfgs[i]) self.blocks.append(TransformerEncoderLayer2(**_layer_cfg)) self.frozen_stages = frozen_stages self.final_norm = final_norm if final_norm: self.ln1 = nn.RMSNorm(self.embed_dims, eps=1e-6) # freeze stages only when self.frozen_stages > 0 if self.frozen_stages > 0: self._freeze_stages() ## load init weights self.init_weights() self.gradient_checkpointing = False def enable_gradient_checkpointing(self, enable=True): self.gradient_checkpointing = enable if self.tokenizer is not None: self.tokenizer.gradient_checkpointing = enable @property def device(self): return next(self.parameters()).device @property def dtype(self): return next(self.parameters()).dtype def init_weights(self): # Initialize class token and storagr token embeddings if self.with_cls_token: trunc_normal_(self.cls_token, std=0.02) if self.storage_tokens is not None: trunc_normal_(self.storage_tokens, std=0.02) # Apply custom initialization to all submodules self.apply(self._init_weights) def _init_weights(self, m): if isinstance(m, nn.Linear): # Use a truncated normal distribution for linear layer weights trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, (nn.LayerNorm, nn.RMSNorm)): # Initialize normalization layers to act as an identity function if hasattr(m, "bias") and m.bias is not None: nn.init.constant_(m.bias, 0) if hasattr(m, "weight") and m.weight is not None: nn.init.constant_(m.weight, 1.0) elif isinstance(m, nn.Conv2d): # Initialize conv layer weights like linear layers trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.constant_(m.bias, 0) def _freeze_stages(self): ## freeze tokenizer if self.frozen_stages >= 1 and self.tokenizer is not None: self.tokenizer.eval() for param in self.tokenizer.parameters(): param.requires_grad = False # freeze patch embedding self.patch_embed.eval() for param in self.patch_embed.parameters(): param.requires_grad = False # freeze cls_token if self.cls_token is not None: self.cls_token.requires_grad = False if self.storage_tokens is not None: self.storage_tokens.requires_grad = False # freeze layers for i in range(1, self.frozen_stages + 1): m = self.blocks[i - 1] m.eval() for param in m.parameters(): param.requires_grad = False # freeze the last layer norm if self.frozen_stages == len(self.blocks): if self.final_norm: self.ln1.eval() for param in self.ln1.parameters(): param.requires_grad = False def forward(self, x): B = x.shape[0] x, patch_resolution = self.patch_embed(x) # (B, 256*256, C) if self.tokenizer is not None: x, patch_resolution = self.tokenizer(x, patch_resolution) # prepend [CLS] and storage tokens prepend = [] if self.cls_token is not None: prepend.append(self.cls_token.expand(B, -1, -1)) if self.storage_tokens is not None: prepend.append(self.storage_tokens.expand(B, -1, -1)) if len(prepend) > 0: x = torch.cat(prepend + [x], dim=1) rope_sincos = self.rope_embed(H=patch_resolution[0], W=patch_resolution[1]) outs = [] for i, layer in enumerate(self.blocks): if self.gradient_checkpointing and self.training: x = checkpoint(layer, x, rope_sincos, use_reentrant=False) else: x = layer(x, rope=rope_sincos) if i == len(self.blocks) - 1 and self.final_norm: x = self.ln1(x) if i in self.out_indices: outs.append(self._format_output(x, patch_resolution)) return tuple(outs) def _format_output(self, x, hw): if self.out_type == "raw": return x if self.out_type == "cls_token": return x[:, 0] patch_token = x[:, self.num_extra_tokens :] if self.out_type == "featmap": B = x.size(0) # (B, N, C) -> (B, H, W, C) -> (B, C, H, W) return patch_token.reshape(B, *hw, -1).permute(0, 3, 1, 2) @property def norm1(self): return self.ln1 # ---------------------------------------------------------------------------- class LayerScale(nn.Module): def __init__( self, dim: int, inplace: bool = False, data_format: str = "channels_last", scale: float = 1e-5, ): super().__init__() assert data_format in ( "channels_last", "channels_first", ), "'data_format' could only be channels_last or channels_first." self.inplace = inplace self.data_format = data_format self.weight = nn.Parameter(torch.ones(dim) * scale) def forward(self, x) -> torch.Tensor: if self.data_format == "channels_first": shape = tuple((1, -1, *(1 for _ in range(x.dim() - 2)))) else: shape = tuple((*(1 for _ in range(x.dim() - 1)), -1)) if self.inplace: return x.mul_(self.weight.view(*shape)) else: return x * self.weight.view(*shape) # ---------------------------------------------------------------------------- class PatchEmbed(nn.Module): def __init__( self, in_channels=3, embed_dims=768, kernel_size=16, stride=16, padding="corner", dilation=1, bias=True, input_size=None, ): super().__init__() self.embed_dims = embed_dims if stride is None: stride = kernel_size kernel_size = to_2tuple(kernel_size) stride = to_2tuple(stride) dilation = to_2tuple(dilation) padding = 0 padding = to_2tuple(padding) self.projection = nn.Conv2d( in_channels=in_channels, out_channels=embed_dims, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=bias, ) if input_size: input_size = to_2tuple(input_size) self.init_input_size = input_size h_out = ( input_size[0] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1 ) // stride[0] + 1 w_out = ( input_size[1] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1 ) // stride[1] + 1 self.init_out_size = (h_out, w_out) else: self.init_input_size = None self.init_out_size = None def forward(self, x): x = self.projection(x) out_size = (x.shape[2], x.shape[3]) x = x.flatten(2).transpose(1, 2) return x, out_size # ---------------------------------------------------------------------------- class SwiGLUFFN(nn.Module): """SwiGLU FFN layer. https://github.com/facebookresearch/dinov2/blob/main/dinov2/layers/swiglu_ffn.py """ # noqa def __init__( self, embed_dims: int, feedforward_channels: Optional[int] = None, out_dims: Optional[int] = None, layer_scale_init_value: float = 0.0, bias: bool = True, add_identity: bool = True, ) -> None: super().__init__() self.embed_dims = embed_dims self.out_dims = out_dims or embed_dims hidden_dims = feedforward_channels or embed_dims self.w12 = nn.Linear(self.embed_dims, 2 * hidden_dims, bias=bias) self.w3 = nn.Linear(hidden_dims, self.out_dims, bias=bias) if layer_scale_init_value > 0: self.gamma2 = LayerScale(dim=embed_dims, scale=layer_scale_init_value) else: self.gamma2 = nn.Identity() self.add_identity = add_identity def forward( self, x: torch.Tensor, identity: Optional[torch.Tensor] = None ) -> torch.Tensor: x12 = self.w12(x) x1, x2 = x12.chunk(2, dim=-1) hidden = F.silu(x1) * x2 out = self.w3(hidden) out = self.gamma2(out) if self.out_dims != self.embed_dims or not self.add_identity: # due to the dimension inconsistence or user setting # not to apply residual operation return out if identity is None: identity = x return identity + out # ---------------------------------------------------------------------------- _IMAGENET_MEAN = (0.485, 0.456, 0.406) _IMAGENET_STD = (0.229, 0.224, 0.225) def imagenet_normalize(tensors_0_1: torch.Tensor) -> torch.Tensor: """Apply ImageNet normalization to a (B, C, H, W) RGB tensor in [0, 1].""" mean = torch.as_tensor( _IMAGENET_MEAN, dtype=tensors_0_1.dtype, device=tensors_0_1.device ).view(1, 3, 1, 1) std = torch.as_tensor( _IMAGENET_STD, dtype=tensors_0_1.dtype, device=tensors_0_1.device ).view(1, 3, 1, 1) return (tensors_0_1 - mean) / std # ---------------------------------------------------------------------------- class MattingHead(nn.Module): """Matting decode head from https://github.com/facebookresearch/sapiens2/blob/main/sapiens/dense/src/models/heads/matting_head.py Predicts a 4-channel output: pre-multiplied foreground RGB (channels 0-2) and soft alpha matte (channel 3), all in [0, 1] after sigmoid. """ def __init__( self, in_channels: int = 1536, upsample_channels: Sequence[int] = (768, 512, 256, 128), conv_out_channels: Optional[Sequence[int]] = (64, 32, 16), conv_kernel_sizes: Optional[Sequence[int]] = (3, 3, 3), out_channels: int = 4, ): super().__init__() self.in_channels = in_channels self.input_conv = nn.Sequential( nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1), nn.InstanceNorm2d(in_channels), nn.SiLU(inplace=True), ) up_blocks = [] cur_ch = in_channels for out_ch in upsample_channels: up_blocks.append( nn.Sequential( nn.Conv2d(cur_ch, out_ch * 4, kernel_size=3, padding=1), nn.PixelShuffle(2), nn.InstanceNorm2d(out_ch), nn.SiLU(inplace=True), ) ) cur_ch = out_ch self.upsample_blocks = nn.Sequential(*up_blocks) conv_layers = [] if conv_out_channels and conv_kernel_sizes: for out_ch, k in zip(conv_out_channels, conv_kernel_sizes): conv_layers.extend( [ nn.Conv2d(cur_ch, out_ch, k, padding=(k - 1) // 2), nn.InstanceNorm2d(out_ch), nn.SiLU(inplace=True), ] ) cur_ch = out_ch self.conv_layers = nn.Sequential(*conv_layers) self.conv_matting = nn.Conv2d(cur_ch, out_channels, kernel_size=1) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.input_conv(x) x = self.upsample_blocks(x) x = self.conv_layers(x) return self.conv_matting(x).sigmoid() # ---------------------------------------------------------------------------- class Sapiens2Matting(nn.Module): """Sapiens2 backbone + MattingHead for human image matting. Reference: https://github.com/facebookresearch/sapiens2/blob/main/docs/MATTING.md """ _ARCH_TO_EMBED_DIM = { "sapiens2_0.1b": 768, "sapiens2_0.4b": 1024, "sapiens2_0.8b": 1280, "sapiens2_1b": 1536, "sapiens2_5b": 2432, } def __init__( self, arch: str = "sapiens2_1b", img_size: Tuple[int, int] = (1024, 768), patch_size: int = 16, ): super().__init__() arch = arch.lower() if arch not in self._ARCH_TO_EMBED_DIM: raise ValueError(f"Unsupported arch {arch}") self.arch = arch self.img_size = to_2tuple(img_size) self.patch_size = patch_size self.backbone = Sapiens2( arch=arch, img_size=img_size, patch_size=patch_size, final_norm=True, use_tokenizer=False, with_cls_token=True, out_type="featmap", ) self.decode_head = MattingHead( in_channels=self._ARCH_TO_EMBED_DIM[arch], upsample_channels=(768, 512, 256, 128), conv_out_channels=(64, 32, 16), conv_kernel_sizes=(3, 3, 3), out_channels=4, ) @classmethod def from_pretrained( cls, repo_id: str = "facebook/sapiens2-matting-1b", filename: str = "sapiens2_1b_matting.safetensors", arch: str = "sapiens2_1b", img_size: Tuple[int, int] = (1024, 768), patch_size: int = 16, device: Optional[torch.device] = None, dtype: Optional[torch.dtype] = None, ) -> "Sapiens2Matting": import huggingface_hub from safetensors.torch import load_file safetensors_path = os.path.join(MODELS_PATH, "sapiens2", filename) if not os.path.exists(safetensors_path): print(f"Downloading pretrained weights from HuggingFace Hub: {repo_id}/{filename}...") os.makedirs(os.path.dirname(safetensors_path), exist_ok=True) safetensors_path = huggingface_hub.hf_hub_download( repo_id=repo_id, filename=filename, local_dir=os.path.join(MODELS_PATH, "sapiens2"), ) model = cls(arch=arch, img_size=img_size, patch_size=patch_size) state_dict = load_file(safetensors_path) model.load_state_dict(state_dict) model.eval() if device is not None or dtype is not None: model.to(device=device, dtype=dtype) return model @property def device(self): return next(self.parameters()).device @property def dtype(self): return next(self.parameters()).dtype @torch.no_grad() def forward(self, image, max_res: int = 1024): """Take a PIL image and return a PIL alpha-matte mask in RGB mode at the original input size. The image is run through the model at its native aspect ratio, snapped to a multiple of patch_size and capped at max_res*max_res pixels.""" from torchvision import transforms p = self.patch_size w, h = image.size target_h, target_w = h, w if target_h * target_w > max_res * max_res: scale = math.sqrt((max_res * max_res) / (target_h * target_w)) target_h = int(target_h * scale) target_w = int(target_w * scale) target_h = max(p, (target_h // p) * p) target_w = max(p, (target_w // p) * p) transform_image = transforms.Compose( [ transforms.Resize((target_h, target_w)), transforms.ToTensor(), transforms.Normalize(_IMAGENET_MEAN, _IMAGENET_STD), ] ) input_images = ( transform_image(image).unsqueeze(0).to(self.device, dtype=self.dtype) ) feat = self.backbone(input_images)[0] out = self.decode_head(feat) # (1, 4, H, W) in [0, 1] alpha = out[0, 3].float().cpu() mask = transforms.ToPILImage()(alpha) mask = mask.resize(image.size).convert("RGB") return mask