|
|
from dataclasses import dataclass, field
|
|
|
from typing import Callable, List, Optional
|
|
|
|
|
|
import torch
|
|
|
import torch.nn as nn
|
|
|
import torch.nn.functional as F
|
|
|
from einops import rearrange
|
|
|
from jaxtyping import Float
|
|
|
from torch import Tensor
|
|
|
from torch.amp import custom_bwd, custom_fwd
|
|
|
from torch.autograd import Function
|
|
|
|
|
|
from spar3d.models.utils import BaseModule, normalize
|
|
|
from spar3d.utils import get_device
|
|
|
|
|
|
|
|
|
def conditional_decorator(decorator_with_args, condition, *args, **kwargs):
|
|
|
def wrapper(fn):
|
|
|
if condition:
|
|
|
if len(kwargs) == 0:
|
|
|
return decorator_with_args
|
|
|
return decorator_with_args(*args, **kwargs)(fn)
|
|
|
else:
|
|
|
return fn
|
|
|
|
|
|
return wrapper
|
|
|
|
|
|
|
|
|
class PixelShuffleUpsampleNetwork(BaseModule):
|
|
|
@dataclass
|
|
|
class Config(BaseModule.Config):
|
|
|
in_channels: int = 1024
|
|
|
out_channels: int = 40
|
|
|
scale_factor: int = 4
|
|
|
|
|
|
conv_layers: int = 4
|
|
|
conv_kernel_size: int = 3
|
|
|
|
|
|
cfg: Config
|
|
|
|
|
|
def configure(self) -> None:
|
|
|
layers = []
|
|
|
output_channels = self.cfg.out_channels * self.cfg.scale_factor**2
|
|
|
|
|
|
in_channels = self.cfg.in_channels
|
|
|
for i in range(self.cfg.conv_layers):
|
|
|
cur_out_channels = (
|
|
|
in_channels if i != self.cfg.conv_layers - 1 else output_channels
|
|
|
)
|
|
|
layers.append(
|
|
|
nn.Conv2d(
|
|
|
in_channels,
|
|
|
cur_out_channels,
|
|
|
self.cfg.conv_kernel_size,
|
|
|
padding=(self.cfg.conv_kernel_size - 1) // 2,
|
|
|
)
|
|
|
)
|
|
|
if i != self.cfg.conv_layers - 1:
|
|
|
layers.append(nn.ReLU(inplace=True))
|
|
|
|
|
|
layers.append(nn.PixelShuffle(self.cfg.scale_factor))
|
|
|
|
|
|
self.upsample = nn.Sequential(*layers)
|
|
|
|
|
|
def forward(
|
|
|
self, triplanes: Float[Tensor, "B 3 Ci Hp Wp"]
|
|
|
) -> Float[Tensor, "B 3 Co Hp2 Wp2"]:
|
|
|
return rearrange(
|
|
|
self.upsample(
|
|
|
rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
|
|
|
),
|
|
|
"(B Np) Co Hp Wp -> B Np Co Hp Wp",
|
|
|
Np=3,
|
|
|
)
|
|
|
|
|
|
|
|
|
class _TruncExp(Function):
|
|
|
|
|
|
|
|
|
@staticmethod
|
|
|
@conditional_decorator(
|
|
|
custom_fwd,
|
|
|
"cuda" in get_device(),
|
|
|
cast_inputs=torch.float32,
|
|
|
device_type="cuda",
|
|
|
)
|
|
|
def forward(ctx, x):
|
|
|
ctx.save_for_backward(x)
|
|
|
return torch.exp(x)
|
|
|
|
|
|
@staticmethod
|
|
|
@conditional_decorator(custom_bwd, "cuda" in get_device())
|
|
|
def backward(ctx, g):
|
|
|
x = ctx.saved_tensors[0]
|
|
|
return g * torch.exp(torch.clamp(x, max=15))
|
|
|
|
|
|
|
|
|
trunc_exp = _TruncExp.apply
|
|
|
|
|
|
|
|
|
def get_activation(name) -> Callable:
|
|
|
if name is None:
|
|
|
return lambda x: x
|
|
|
name = name.lower()
|
|
|
if name == "none" or name == "linear" or name == "identity":
|
|
|
return lambda x: x
|
|
|
elif name == "lin2srgb":
|
|
|
return lambda x: torch.where(
|
|
|
x > 0.0031308,
|
|
|
torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) * 1.055 - 0.055,
|
|
|
12.92 * x,
|
|
|
).clamp(0.0, 1.0)
|
|
|
elif name == "exp":
|
|
|
return lambda x: torch.exp(x)
|
|
|
elif name == "shifted_exp":
|
|
|
return lambda x: torch.exp(x - 1.0)
|
|
|
elif name == "trunc_exp":
|
|
|
return trunc_exp
|
|
|
elif name == "shifted_trunc_exp":
|
|
|
return lambda x: trunc_exp(x - 1.0)
|
|
|
elif name == "sigmoid":
|
|
|
return lambda x: torch.sigmoid(x)
|
|
|
elif name == "tanh":
|
|
|
return lambda x: torch.tanh(x)
|
|
|
elif name == "shifted_softplus":
|
|
|
return lambda x: F.softplus(x - 1.0)
|
|
|
elif name == "scale_-11_01":
|
|
|
return lambda x: x * 0.5 + 0.5
|
|
|
elif name == "negative":
|
|
|
return lambda x: -x
|
|
|
elif name == "normalize_channel_last":
|
|
|
return lambda x: normalize(x)
|
|
|
elif name == "normalize_channel_first":
|
|
|
return lambda x: normalize(x, dim=1)
|
|
|
else:
|
|
|
try:
|
|
|
return getattr(F, name)
|
|
|
except AttributeError:
|
|
|
raise ValueError(f"Unknown activation function: {name}")
|
|
|
|
|
|
|
|
|
class LambdaModule(torch.nn.Module):
|
|
|
def __init__(self, lambd: Callable[[torch.Tensor], torch.Tensor]):
|
|
|
super().__init__()
|
|
|
self.lambd = lambd
|
|
|
|
|
|
def forward(self, x):
|
|
|
return self.lambd(x)
|
|
|
|
|
|
|
|
|
def get_activation_module(name) -> torch.nn.Module:
|
|
|
return LambdaModule(get_activation(name))
|
|
|
|
|
|
|
|
|
@dataclass
|
|
|
class HeadSpec:
|
|
|
name: str
|
|
|
out_channels: int
|
|
|
n_hidden_layers: int
|
|
|
output_activation: Optional[str] = None
|
|
|
out_bias: float = 0.0
|
|
|
|
|
|
|
|
|
class MaterialMLP(BaseModule):
|
|
|
@dataclass
|
|
|
class Config(BaseModule.Config):
|
|
|
in_channels: int = 120
|
|
|
n_neurons: int = 64
|
|
|
activation: str = "silu"
|
|
|
heads: List[HeadSpec] = field(default_factory=lambda: [])
|
|
|
|
|
|
cfg: Config
|
|
|
|
|
|
def configure(self) -> None:
|
|
|
assert len(self.cfg.heads) > 0
|
|
|
heads = {}
|
|
|
for head in self.cfg.heads:
|
|
|
head_layers = []
|
|
|
for i in range(head.n_hidden_layers):
|
|
|
head_layers += [
|
|
|
nn.Linear(
|
|
|
self.cfg.in_channels if i == 0 else self.cfg.n_neurons,
|
|
|
self.cfg.n_neurons,
|
|
|
),
|
|
|
self.make_activation(self.cfg.activation),
|
|
|
]
|
|
|
head_layers += [
|
|
|
nn.Linear(
|
|
|
self.cfg.n_neurons,
|
|
|
head.out_channels,
|
|
|
),
|
|
|
]
|
|
|
heads[head.name] = nn.Sequential(*head_layers)
|
|
|
self.heads = nn.ModuleDict(heads)
|
|
|
|
|
|
def make_activation(self, activation):
|
|
|
if activation == "relu":
|
|
|
return nn.ReLU(inplace=True)
|
|
|
elif activation == "silu":
|
|
|
return nn.SiLU(inplace=True)
|
|
|
else:
|
|
|
raise NotImplementedError
|
|
|
|
|
|
def keys(self):
|
|
|
return self.heads.keys()
|
|
|
|
|
|
def forward(
|
|
|
self, x, include: Optional[List] = None, exclude: Optional[List] = None
|
|
|
):
|
|
|
if include is not None and exclude is not None:
|
|
|
raise ValueError("Cannot specify both include and exclude.")
|
|
|
if include is not None:
|
|
|
heads = [h for h in self.cfg.heads if h.name in include]
|
|
|
elif exclude is not None:
|
|
|
heads = [h for h in self.cfg.heads if h.name not in exclude]
|
|
|
else:
|
|
|
heads = self.cfg.heads
|
|
|
|
|
|
out = {
|
|
|
head.name: get_activation(head.output_activation)(
|
|
|
self.heads[head.name](x) + head.out_bias
|
|
|
)
|
|
|
for head in heads
|
|
|
}
|
|
|
|
|
|
return out
|
|
|
|
