(
self,
*,
ch,
out_ch,
ch_mult=(1, 2, 4, 8),
num_res_blocks,
attn_resolutions,
dropout=0.0,
resamp_with_conv=True,
in_channels,
resolution,
use_timestep=True,
use_linear_attn=False,
attn_type="vanilla",
)
| 262 | |
| 263 | class Model(nn.Module): |
| 264 | def __init__( |
| 265 | self, |
| 266 | *, |
| 267 | ch, |
| 268 | out_ch, |
| 269 | ch_mult=(1, 2, 4, 8), |
| 270 | num_res_blocks, |
| 271 | attn_resolutions, |
| 272 | dropout=0.0, |
| 273 | resamp_with_conv=True, |
| 274 | in_channels, |
| 275 | resolution, |
| 276 | use_timestep=True, |
| 277 | use_linear_attn=False, |
| 278 | attn_type="vanilla", |
| 279 | ): |
| 280 | super().__init__() |
| 281 | if use_linear_attn: |
| 282 | attn_type = "linear" |
| 283 | self.ch = ch |
| 284 | self.temb_ch = self.ch * 4 |
| 285 | self.num_resolutions = len(ch_mult) |
| 286 | self.num_res_blocks = num_res_blocks |
| 287 | self.resolution = resolution |
| 288 | self.in_channels = in_channels |
| 289 | |
| 290 | self.use_timestep = use_timestep |
| 291 | if self.use_timestep: |
| 292 | # timestep embedding |
| 293 | self.temb = nn.Module() |
| 294 | self.temb.dense = nn.ModuleList( |
| 295 | [ |
| 296 | torch.nn.Linear(self.ch, self.temb_ch), |
| 297 | torch.nn.Linear(self.temb_ch, self.temb_ch), |
| 298 | ] |
| 299 | ) |
| 300 | |
| 301 | # downsampling |
| 302 | self.conv_in = torch.nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) |
| 303 | |
| 304 | curr_res = resolution |
| 305 | in_ch_mult = (1,) + tuple(ch_mult) |
| 306 | self.down = nn.ModuleList() |
| 307 | for i_level in range(self.num_resolutions): |
| 308 | block = nn.ModuleList() |
| 309 | attn = nn.ModuleList() |
| 310 | block_in = ch * in_ch_mult[i_level] |
| 311 | block_out = ch * ch_mult[i_level] |
| 312 | for i_block in range(self.num_res_blocks): |
| 313 | block.append( |
| 314 | ResnetBlock( |
| 315 | in_channels=block_in, |
| 316 | out_channels=block_out, |
| 317 | temb_channels=self.temb_ch, |
| 318 | dropout=dropout, |
| 319 | ) |
| 320 | ) |
| 321 | block_in = block_out |
no test coverage detected