Spaces:
Running on Zero
Running on Zero
File size: 4,581 Bytes
dc09dcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | # Modified for SplatAnything: inference-only release.
from typing import Callable, Optional, Tuple, Union
import torch
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
from itertools import repeat
import collections.abc
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)
(H, W) = (x.size(2), x.size(3))
x = x.flatten(2).transpose(1, 2)
x = self.norm(x)
if not self.flatten_embedding:
x = x.reshape(-1, H, W, self.embed_dim)
return x
class PatchEmbed_Mlp(PatchEmbed):
def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, norm_layer=None, flatten_embedding=True):
super().__init__(img_size, patch_size, in_chans, embed_dim, norm_layer, flatten_embedding)
self.proj = nn.Sequential(PixelUnshuffle(patch_size), Permute((0, 2, 3, 1)), Mlp(in_chans * patch_size ** 2, 4 * embed_dim, embed_dim), Permute((0, 3, 1, 2)))
class PixelUnshuffle(nn.Module):
def __init__(self, downscale_factor):
super().__init__()
self.downscale_factor = downscale_factor
def forward(self, input):
if input.numel() == 0:
(C, H, W) = input.shape[-3:]
assert H and W and (H % self.downscale_factor == W % self.downscale_factor == 0)
return input.view(*input.shape[:-3], C * self.downscale_factor ** 2, H // self.downscale_factor, W // self.downscale_factor)
else:
return F.pixel_unshuffle(input, self.downscale_factor)
class Permute(torch.nn.Module):
dims: tuple[int, ...]
def __init__(self, dims: tuple[int, ...]) -> None:
super().__init__()
self.dims = tuple(dims)
def __repr__(self):
return f'Permute{self.dims}'
def forward(self, input: torch.Tensor) -> torch.Tensor:
return input.permute(*self.dims)
def _ntuple(n):
def parse(x):
if isinstance(x, collections.abc.Iterable) and (not isinstance(x, str)):
return x
return tuple(repeat(x, n))
return parse
to_2tuple = _ntuple(2)
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, bias=True, drop=0.0):
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)
self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0])
self.act = act_layer()
self.drop1 = nn.Dropout(drop_probs[0])
self.fc2 = nn.Linear(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
|