2D Image to Patch Embedding
| 305 | |
| 306 | |
| 307 | class PatchEmbedMR(nn.Module): |
| 308 | """ 2D Image to Patch Embedding |
| 309 | """ |
| 310 | def __init__( |
| 311 | self, |
| 312 | patch_size: int = 2, |
| 313 | in_chans: int = 4, |
| 314 | embed_dim: int = 768, |
| 315 | bias: bool = True, |
| 316 | ): |
| 317 | super().__init__() |
| 318 | self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias) |
| 319 | |
| 320 | def forward(self, x): |
| 321 | x = self.proj(x) |
| 322 | x = x.flatten(2).transpose(1, 2) # NCHW -> NLC |
| 323 | return x |
| 324 | |
| 325 | |
| 326 | class OmniGenOriginalModel(nn.Module): |