(
self,
in_channels: int = 16,
hidden_size: int = 2560,
patch_size: int = 2,
text_hidden_size: int = 4096,
pos_embed_max_size: int = 128,
)
| 773 | |
| 774 | class CogView3PlusPatchEmbed(nn.Module): |
| 775 | def __init__( |
| 776 | self, |
| 777 | in_channels: int = 16, |
| 778 | hidden_size: int = 2560, |
| 779 | patch_size: int = 2, |
| 780 | text_hidden_size: int = 4096, |
| 781 | pos_embed_max_size: int = 128, |
| 782 | ): |
| 783 | super().__init__() |
| 784 | self.in_channels = in_channels |
| 785 | self.hidden_size = hidden_size |
| 786 | self.patch_size = patch_size |
| 787 | self.text_hidden_size = text_hidden_size |
| 788 | self.pos_embed_max_size = pos_embed_max_size |
| 789 | # Linear projection for image patches |
| 790 | self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) |
| 791 | |
| 792 | # Linear projection for text embeddings |
| 793 | self.text_proj = nn.Linear(text_hidden_size, hidden_size) |
| 794 | |
| 795 | pos_embed = get_2d_sincos_pos_embed( |
| 796 | hidden_size, pos_embed_max_size, base_size=pos_embed_max_size, output_type="pt" |
| 797 | ) |
| 798 | pos_embed = pos_embed.reshape(pos_embed_max_size, pos_embed_max_size, hidden_size) |
| 799 | self.register_buffer("pos_embed", pos_embed.float(), persistent=False) |
| 800 | |
| 801 | def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor) -> torch.Tensor: |
| 802 | batch_size, channel, height, width = hidden_states.shape |
nothing calls this directly
no test coverage detected