(
self,
height=224,
width=224,
patch_size=16,
in_channels=3,
embed_dim=768,
layer_norm=False,
flatten=True,
bias=True,
interpolation_scale=1,
pos_embed_type="sincos",
pos_embed_max_size=None, # For SD3 cropping
)
| 474 | """ |
| 475 | |
| 476 | def __init__( |
| 477 | self, |
| 478 | height=224, |
| 479 | width=224, |
| 480 | patch_size=16, |
| 481 | in_channels=3, |
| 482 | embed_dim=768, |
| 483 | layer_norm=False, |
| 484 | flatten=True, |
| 485 | bias=True, |
| 486 | interpolation_scale=1, |
| 487 | pos_embed_type="sincos", |
| 488 | pos_embed_max_size=None, # For SD3 cropping |
| 489 | ): |
| 490 | super().__init__() |
| 491 | |
| 492 | num_patches = (height // patch_size) * (width // patch_size) |
| 493 | self.flatten = flatten |
| 494 | self.layer_norm = layer_norm |
| 495 | self.pos_embed_max_size = pos_embed_max_size |
| 496 | |
| 497 | self.proj = nn.Conv2d( |
| 498 | in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias |
| 499 | ) |
| 500 | if layer_norm: |
| 501 | self.norm = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6) |
| 502 | else: |
| 503 | self.norm = None |
| 504 | |
| 505 | self.patch_size = patch_size |
| 506 | self.height, self.width = height // patch_size, width // patch_size |
| 507 | self.base_size = height // patch_size |
| 508 | self.interpolation_scale = interpolation_scale |
| 509 | |
| 510 | # Calculate positional embeddings based on max size or default |
| 511 | if pos_embed_max_size: |
| 512 | grid_size = pos_embed_max_size |
| 513 | else: |
| 514 | grid_size = int(num_patches**0.5) |
| 515 | |
| 516 | if pos_embed_type is None: |
| 517 | self.pos_embed = None |
| 518 | elif pos_embed_type == "sincos": |
| 519 | pos_embed = get_2d_sincos_pos_embed( |
| 520 | embed_dim, |
| 521 | grid_size, |
| 522 | base_size=self.base_size, |
| 523 | interpolation_scale=self.interpolation_scale, |
| 524 | output_type="pt", |
| 525 | ) |
| 526 | persistent = True if pos_embed_max_size else False |
| 527 | self.register_buffer("pos_embed", pos_embed.float().unsqueeze(0), persistent=persistent) |
| 528 | else: |
| 529 | raise ValueError(f"Unsupported pos_embed_type: {pos_embed_type}") |
| 530 | |
| 531 | def cropped_pos_embed(self, height, width): |
| 532 | """Crops positional embeddings for SD3 compatibility.""" |
nothing calls this directly
no test coverage detected