(
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,
num_classes=None,
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,
use_spatial_transformer=False, # custom transformer support
transformer_depth=1, # custom transformer support
context_dim=None, # custom transformer support
n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model
legacy=True,
)
| 655 | """ |
| 656 | |
| 657 | def __init__( |
| 658 | self, |
| 659 | image_size, |
| 660 | in_channels, |
| 661 | model_channels, |
| 662 | out_channels, |
| 663 | num_res_blocks, |
| 664 | attention_resolutions, |
| 665 | dropout=0, |
| 666 | channel_mult=(1, 2, 4, 8), |
| 667 | conv_resample=True, |
| 668 | dims=2, |
| 669 | num_classes=None, |
| 670 | use_checkpoint=False, |
| 671 | use_fp16=False, |
| 672 | num_heads=-1, |
| 673 | num_head_channels=-1, |
| 674 | num_heads_upsample=-1, |
| 675 | use_scale_shift_norm=False, |
| 676 | resblock_updown=False, |
| 677 | use_new_attention_order=False, |
| 678 | use_spatial_transformer=False, # custom transformer support |
| 679 | transformer_depth=1, # custom transformer support |
| 680 | context_dim=None, # custom transformer support |
| 681 | n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model |
| 682 | legacy=True, |
| 683 | ): |
| 684 | super().__init__() |
| 685 | if use_spatial_transformer: |
| 686 | assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...' |
| 687 | |
| 688 | if context_dim is not None: |
| 689 | assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...' |
| 690 | from omegaconf.listconfig import ListConfig |
| 691 | if type(context_dim) == ListConfig: |
| 692 | context_dim = list(context_dim) |
| 693 | |
| 694 | if num_heads_upsample == -1: |
| 695 | num_heads_upsample = num_heads |
| 696 | |
| 697 | if num_heads == -1: |
| 698 | assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' |
| 699 | |
| 700 | if num_head_channels == -1: |
| 701 | assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' |
| 702 | |
| 703 | self.image_size = image_size |
| 704 | self.in_channels = in_channels |
| 705 | self.model_channels = model_channels |
| 706 | self.out_channels = out_channels |
| 707 | self.num_res_blocks = num_res_blocks |
| 708 | self.attention_resolutions = attention_resolutions |
| 709 | self.dropout = dropout |
| 710 | self.channel_mult = channel_mult |
| 711 | self.conv_resample = conv_resample |
| 712 | self.num_classes = num_classes |
| 713 | self.use_checkpoint = use_checkpoint |
| 714 | self.dtype = th.float16 if use_fp16 else th.float32 |
nothing calls this directly
no test coverage detected