| | 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
|
| |
|
