Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized training. This requires ControlNets to convert image-based condit
| 42 | |
| 43 | |
| 44 | class ControlNetConditioningEmbedding(nn.Module): |
| 45 | """ |
| 46 | Quoting from https://arxiv.org/abs/2302.05543: "Stable Diffusion uses a pre-processing method similar to VQ-GAN |
| 47 | [11] to convert the entire dataset of 512 × 512 images into smaller 64 × 64 “latent images” for stabilized |
| 48 | training. This requires ControlNets to convert image-based conditions to 64 × 64 feature space to match the |
| 49 | convolution size. We use a tiny network E(·) of four convolution layers with 4 × 4 kernels and 2 × 2 strides |
| 50 | (activated by ReLU, channels are 16, 32, 64, 128, initialized with Gaussian weights, trained jointly with the full |
| 51 | model) to encode image-space conditions ... into feature maps ..." |
| 52 | """ |
| 53 | |
| 54 | def __init__( |
| 55 | self, |
| 56 | conditioning_embedding_channels: int, |
| 57 | conditioning_channels: int = 3, |
| 58 | block_out_channels: Tuple[int] = (16, 32, 96, 256), |
| 59 | ): |
| 60 | super().__init__() |
| 61 | |
| 62 | self.conv_in = nn.Conv2d(conditioning_channels, block_out_channels[0], kernel_size=3, padding=1) |
| 63 | |
| 64 | self.blocks = nn.ModuleList([]) |
| 65 | |
| 66 | for i in range(len(block_out_channels) - 1): |
| 67 | channel_in = block_out_channels[i] |
| 68 | channel_out = block_out_channels[i + 1] |
| 69 | self.blocks.append(nn.Conv2d(channel_in, channel_in, kernel_size=3, padding=1)) |
| 70 | self.blocks.append(nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1, stride=2)) |
| 71 | |
| 72 | self.conv_out = zero_module( |
| 73 | nn.Conv2d(block_out_channels[-1], conditioning_embedding_channels, kernel_size=3, padding=1) |
| 74 | ) |
| 75 | |
| 76 | def forward(self, conditioning): |
| 77 | embedding = self.conv_in(conditioning) |
| 78 | embedding = F.silu(embedding) |
| 79 | |
| 80 | for block in self.blocks: |
| 81 | embedding = block(embedding) |
| 82 | embedding = F.silu(embedding) |
| 83 | |
| 84 | embedding = self.conv_out(embedding) |
| 85 | |
| 86 | return embedding |
| 87 | |
| 88 | |
| 89 | class ControlNetModel(ModelMixin, ConfigMixin): |