(
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,
)
| 112 | use_fp16 (bool): Whether to use FP16. |
| 113 | """ |
| 114 | def __init__( |
| 115 | self, |
| 116 | in_channels: int, |
| 117 | latent_channels: int, |
| 118 | num_res_blocks: int, |
| 119 | channels: List[int], |
| 120 | num_res_blocks_middle: int = 2, |
| 121 | norm_type: Literal["group", "layer"] = "layer", |
| 122 | use_fp16: bool = False, |
| 123 | ): |
| 124 | super().__init__() |
| 125 | self.in_channels = in_channels |
| 126 | self.latent_channels = latent_channels |
| 127 | self.num_res_blocks = num_res_blocks |
| 128 | self.channels = channels |
| 129 | self.num_res_blocks_middle = num_res_blocks_middle |
| 130 | self.norm_type = norm_type |
| 131 | self.use_fp16 = use_fp16 |
| 132 | self.dtype = torch.float16 if use_fp16 else torch.float32 |
| 133 | |
| 134 | self.input_layer = nn.Conv3d(in_channels, channels[0], 3, padding=1) |
| 135 | |
| 136 | self.blocks = nn.ModuleList([]) |
| 137 | for i, ch in enumerate(channels): |
| 138 | self.blocks.extend([ |
| 139 | ResBlock3d(ch, ch) |
| 140 | for _ in range(num_res_blocks) |
| 141 | ]) |
| 142 | if i < len(channels) - 1: |
| 143 | self.blocks.append( |
| 144 | DownsampleBlock3d(ch, channels[i+1]) |
| 145 | ) |
| 146 | |
| 147 | self.middle_block = nn.Sequential(*[ |
| 148 | ResBlock3d(channels[-1], channels[-1]) |
| 149 | for _ in range(num_res_blocks_middle) |
| 150 | ]) |
| 151 | |
| 152 | self.out_layer = nn.Sequential( |
| 153 | norm_layer(norm_type, channels[-1]), |
| 154 | nn.SiLU(), |
| 155 | nn.Conv3d(channels[-1], latent_channels*2, 3, padding=1) |
| 156 | ) |
| 157 | |
| 158 | if use_fp16: |
| 159 | self.convert_to_fp16() |
| 160 | |
| 161 | @property |
| 162 | def device(self) -> torch.device: |
nothing calls this directly
no test coverage detected