(self, input_res, patch_size, width, layers, heads, window_size, encoder_dim=None, condition_len=576, condition_dim=None, drop_path_rate=0.1)
| 443 | |
| 444 | class TransformerConditionalDecoder(TransformerBase): |
| 445 | def __init__(self, input_res, patch_size, width, layers, heads, window_size, encoder_dim=None, condition_len=576, condition_dim=None, drop_path_rate=0.1): |
| 446 | self.input_res = input_res |
| 447 | self.patch_size = patch_size |
| 448 | self.width = width |
| 449 | token_len = (input_res[0] // patch_size) * (input_res[1] // patch_size) |
| 450 | super().__init__(width, layers, heads, window_size, token_len, ConditionalResAttBlock, drop_path_rate=drop_path_rate) |
| 451 | self.positional_embedding = nn.Parameter(torch.zeros(1, token_len*2, width)) |
| 452 | nn.init.trunc_normal_(self.positional_embedding, std=0.02) |
| 453 | self.cls_embedding = nn.Parameter(torch.zeros(1, 2, width)) |
| 454 | nn.init.trunc_normal_(self.cls_embedding, std=0.02) |
| 455 | self.positional_encoding = SinusoidalPositionalEncoding(max_len=condition_len, d_model=width) |
| 456 | |
| 457 | if condition_dim is not None: |
| 458 | self.condition_proj = nn.Linear(condition_dim, width, bias=False) |
| 459 | else: |
| 460 | self.condition_proj = nn.Identity() |
| 461 | |
| 462 | self.out_proj = nn.Identity() |
| 463 | |
| 464 | self.dropout = nn.Dropout(drop_path_rate) |
| 465 | self.condition_ln = LayerNorm(width) |
| 466 | |
| 467 | def forward(self, latent, condition): |
| 468 | b, v = latent.shape[:2] |
nothing calls this directly
no test coverage detected