(
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,
)
| 221 | use_fp16 (bool): Whether to use FP16. |
| 222 | """ |
| 223 | def __init__( |
| 224 | self, |
| 225 | out_channels: int, |
| 226 | latent_channels: int, |
| 227 | num_res_blocks: int, |
| 228 | channels: List[int], |
| 229 | num_res_blocks_middle: int = 2, |
| 230 | norm_type: Literal["group", "layer"] = "layer", |
| 231 | use_fp16: bool = False, |
| 232 | ): |
| 233 | super().__init__() |
| 234 | self.out_channels = out_channels |
| 235 | self.latent_channels = latent_channels |
| 236 | self.num_res_blocks = num_res_blocks |
| 237 | self.channels = channels |
| 238 | self.num_res_blocks_middle = num_res_blocks_middle |
| 239 | self.norm_type = norm_type |
| 240 | self.use_fp16 = use_fp16 |
| 241 | self.dtype = torch.float16 if use_fp16 else torch.float32 |
| 242 | |
| 243 | self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1) |
| 244 | |
| 245 | self.middle_block = nn.Sequential(*[ |
| 246 | ResBlock3d(channels[0], channels[0]) |
| 247 | for _ in range(num_res_blocks_middle) |
| 248 | ]) |
| 249 | |
| 250 | self.blocks = nn.ModuleList([]) |
| 251 | for i, ch in enumerate(channels): |
| 252 | self.blocks.extend([ |
| 253 | ResBlock3d(ch, ch) |
| 254 | for _ in range(num_res_blocks) |
| 255 | ]) |
| 256 | if i < len(channels) - 1: |
| 257 | self.blocks.append( |
| 258 | UpsampleBlock3d(ch, channels[i+1]) |
| 259 | ) |
| 260 | |
| 261 | self.out_layer = nn.Sequential( |
| 262 | norm_layer(norm_type, channels[-1]), |
| 263 | nn.SiLU(), |
| 264 | nn.Conv3d(channels[-1], out_channels, 3, padding=1) |
| 265 | ) |
| 266 | |
| 267 | if use_fp16: |
| 268 | self.convert_to_fp16() |
| 269 | |
| 270 | @property |
| 271 | def device(self) -> torch.device: |
nothing calls this directly
no test coverage detected