| 8 | from diffusers.models.modeling_utils import ModelMixin |
| 9 | |
| 10 | class PoseNet(ModelMixin): |
| 11 | def __init__(self, noise_latent_channels=320): |
| 12 | super().__init__() |
| 13 | # multiple convolution layers |
| 14 | self.conv_layers = nn.Sequential( |
| 15 | nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1), |
| 16 | nn.SiLU(), |
| 17 | nn.Conv2d(in_channels=3, out_channels=16, kernel_size=4, stride=2, padding=1), |
| 18 | nn.SiLU(), |
| 19 | |
| 20 | nn.Conv2d(in_channels=16, out_channels=16, kernel_size=3, padding=1), |
| 21 | nn.SiLU(), |
| 22 | nn.Conv2d(in_channels=16, out_channels=32, kernel_size=4, stride=2, padding=1), |
| 23 | nn.SiLU(), |
| 24 | |
| 25 | nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1), |
| 26 | nn.SiLU(), |
| 27 | nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, stride=2, padding=1), |
| 28 | nn.SiLU(), |
| 29 | |
| 30 | nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1), |
| 31 | nn.SiLU(), |
| 32 | nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1), |
| 33 | nn.SiLU() |
| 34 | ) |
| 35 | |
| 36 | # Final projection layer |
| 37 | self.final_proj = nn.Conv2d(in_channels=128, out_channels=noise_latent_channels, kernel_size=1) |
| 38 | |
| 39 | # Initialize layers |
| 40 | self._initialize_weights() |
| 41 | |
| 42 | self.scale = nn.Parameter(torch.ones(1) * 2) |
| 43 | |
| 44 | def _initialize_weights(self): |
| 45 | """Initialize weights with He. initialization and zero out the biases |
| 46 | """ |
| 47 | for m in self.conv_layers: |
| 48 | if isinstance(m, nn.Conv2d): |
| 49 | n = m.kernel_size[0] * m.kernel_size[1] * m.in_channels |
| 50 | init.normal_(m.weight, mean=0.0, std=np.sqrt(2. / n)) |
| 51 | if m.bias is not None: |
| 52 | init.zeros_(m.bias) |
| 53 | init.zeros_(self.final_proj.weight) |
| 54 | if self.final_proj.bias is not None: |
| 55 | init.zeros_(self.final_proj.bias) |
| 56 | |
| 57 | def forward(self, x): |
| 58 | if x.ndim == 5: |
| 59 | x = einops.rearrange(x, "b f c h w -> (b f) c h w") |
| 60 | x = self.conv_layers(x) |
| 61 | x = self.final_proj(x) |
| 62 | |
| 63 | return x * self.scale |
| 64 | |
| 65 | @classmethod |
| 66 | def from_pretrained(cls, pretrained_model_path): |
| 67 | """load pretrained pose-net weights |
no outgoing calls
no test coverage detected