| 50 | |
| 51 | |
| 52 | class ControlNet(nn.Module): |
| 53 | def __init__( |
| 54 | self, |
| 55 | image_size, |
| 56 | in_channels, |
| 57 | model_channels, |
| 58 | hint_channels, |
| 59 | num_res_blocks, |
| 60 | attention_resolutions, |
| 61 | dropout=0, |
| 62 | channel_mult=(1, 2, 4, 8), |
| 63 | conv_resample=True, |
| 64 | dims=2, |
| 65 | use_checkpoint=False, |
| 66 | use_fp16=False, |
| 67 | num_heads=-1, |
| 68 | num_head_channels=-1, |
| 69 | num_heads_upsample=-1, |
| 70 | use_scale_shift_norm=False, |
| 71 | resblock_updown=False, |
| 72 | use_new_attention_order=False, |
| 73 | use_spatial_transformer=False, # custom transformer support |
| 74 | transformer_depth=1, # custom transformer support |
| 75 | context_dim=None, # custom transformer support |
| 76 | n_embed=None, # custom support for prediction of discrete ids into codebook of first stage vq model |
| 77 | legacy=True, |
| 78 | disable_self_attentions=None, |
| 79 | num_attention_blocks=None, |
| 80 | disable_middle_self_attn=False, |
| 81 | use_linear_in_transformer=False, |
| 82 | ): |
| 83 | super().__init__() |
| 84 | if use_spatial_transformer: |
| 85 | assert context_dim is not None, 'Fool!! You forgot to include the dimension of your cross-attention conditioning...' |
| 86 | |
| 87 | if context_dim is not None: |
| 88 | assert use_spatial_transformer, 'Fool!! You forgot to use the spatial transformer for your cross-attention conditioning...' |
| 89 | from omegaconf.listconfig import ListConfig |
| 90 | if type(context_dim) == ListConfig: |
| 91 | context_dim = list(context_dim) |
| 92 | |
| 93 | if num_heads_upsample == -1: |
| 94 | num_heads_upsample = num_heads |
| 95 | |
| 96 | if num_heads == -1: |
| 97 | assert num_head_channels != -1, 'Either num_heads or num_head_channels has to be set' |
| 98 | |
| 99 | if num_head_channels == -1: |
| 100 | assert num_heads != -1, 'Either num_heads or num_head_channels has to be set' |
| 101 | |
| 102 | self.dims = dims |
| 103 | self.image_size = image_size |
| 104 | self.in_channels = in_channels |
| 105 | self.model_channels = model_channels |
| 106 | if isinstance(num_res_blocks, int): |
| 107 | self.num_res_blocks = len(channel_mult) * [num_res_blocks] |
| 108 | else: |
| 109 | if len(num_res_blocks) != len(channel_mult): |
nothing calls this directly
no outgoing calls
no test coverage detected