(
self,
image_size,
in_channels,
model_channels,
out_channels,
num_res_blocks,
attention_resolutions,
dropout=0,
channel_mult=(1, 2, 4, 8),
conv_resample=True,
dims=2,
use_checkpoint=False,
use_fp16=False,
num_heads=1,
num_head_channels=-1,
num_heads_upsample=-1,
use_scale_shift_norm=False,
resblock_updown=False,
use_new_attention_order=False,
*args,
**kwargs
)
| 1345 | """ |
| 1346 | |
| 1347 | def __init__( |
| 1348 | self, |
| 1349 | image_size, |
| 1350 | in_channels, |
| 1351 | model_channels, |
| 1352 | out_channels, |
| 1353 | num_res_blocks, |
| 1354 | attention_resolutions, |
| 1355 | dropout=0, |
| 1356 | channel_mult=(1, 2, 4, 8), |
| 1357 | conv_resample=True, |
| 1358 | dims=2, |
| 1359 | use_checkpoint=False, |
| 1360 | use_fp16=False, |
| 1361 | num_heads=1, |
| 1362 | num_head_channels=-1, |
| 1363 | num_heads_upsample=-1, |
| 1364 | use_scale_shift_norm=False, |
| 1365 | resblock_updown=False, |
| 1366 | use_new_attention_order=False, |
| 1367 | *args, |
| 1368 | **kwargs |
| 1369 | ): |
| 1370 | super().__init__() |
| 1371 | |
| 1372 | if num_heads_upsample == -1: |
| 1373 | num_heads_upsample = num_heads |
| 1374 | |
| 1375 | self.in_channels = in_channels |
| 1376 | self.model_channels = model_channels |
| 1377 | self.out_channels = out_channels |
| 1378 | self.num_res_blocks = num_res_blocks |
| 1379 | self.attention_resolutions = attention_resolutions |
| 1380 | self.dropout = dropout |
| 1381 | self.channel_mult = channel_mult |
| 1382 | self.conv_resample = conv_resample |
| 1383 | self.use_checkpoint = use_checkpoint |
| 1384 | self.dtype = th.float16 if use_fp16 else th.float32 |
| 1385 | self.num_heads = num_heads |
| 1386 | self.num_head_channels = num_head_channels |
| 1387 | self.num_heads_upsample = num_heads_upsample |
| 1388 | |
| 1389 | time_embed_dim = model_channels * 4 |
| 1390 | self.time_embed = nn.Sequential( |
| 1391 | linear(model_channels, time_embed_dim), |
| 1392 | nn.SiLU(), |
| 1393 | linear(time_embed_dim, time_embed_dim), |
| 1394 | ) |
| 1395 | |
| 1396 | self.input_blocks = nn.ModuleList( |
| 1397 | [ |
| 1398 | TimestepEmbedSequential( |
| 1399 | conv_nd(dims, in_channels, model_channels, 3, padding=1) |
| 1400 | ) |
| 1401 | ] |
| 1402 | ) |
| 1403 | self._feature_size = model_channels |
| 1404 | input_block_chans = [] |
nothing calls this directly
no test coverage detected