(
self,
patch_size: int = 2,
patch_size_t: int | None = None,
in_channels: int = 16,
embed_dim: int = 1920,
text_embed_dim: int = 4096,
bias: bool = True,
sample_width: int = 90,
sample_height: int = 60,
sample_frames: int = 49,
temporal_compression_ratio: int = 4,
max_text_seq_length: int = 226,
spatial_interpolation_scale: float = 1.875,
temporal_interpolation_scale: float = 1.0,
use_positional_embeddings: bool = True,
use_learned_positional_embeddings: bool = True,
)
| 639 | |
| 640 | class CogVideoXPatchEmbed(nn.Module): |
| 641 | def __init__( |
| 642 | self, |
| 643 | patch_size: int = 2, |
| 644 | patch_size_t: int | None = None, |
| 645 | in_channels: int = 16, |
| 646 | embed_dim: int = 1920, |
| 647 | text_embed_dim: int = 4096, |
| 648 | bias: bool = True, |
| 649 | sample_width: int = 90, |
| 650 | sample_height: int = 60, |
| 651 | sample_frames: int = 49, |
| 652 | temporal_compression_ratio: int = 4, |
| 653 | max_text_seq_length: int = 226, |
| 654 | spatial_interpolation_scale: float = 1.875, |
| 655 | temporal_interpolation_scale: float = 1.0, |
| 656 | use_positional_embeddings: bool = True, |
| 657 | use_learned_positional_embeddings: bool = True, |
| 658 | ) -> None: |
| 659 | super().__init__() |
| 660 | |
| 661 | self.patch_size = patch_size |
| 662 | self.patch_size_t = patch_size_t |
| 663 | self.embed_dim = embed_dim |
| 664 | self.sample_height = sample_height |
| 665 | self.sample_width = sample_width |
| 666 | self.sample_frames = sample_frames |
| 667 | self.temporal_compression_ratio = temporal_compression_ratio |
| 668 | self.max_text_seq_length = max_text_seq_length |
| 669 | self.spatial_interpolation_scale = spatial_interpolation_scale |
| 670 | self.temporal_interpolation_scale = temporal_interpolation_scale |
| 671 | self.use_positional_embeddings = use_positional_embeddings |
| 672 | self.use_learned_positional_embeddings = use_learned_positional_embeddings |
| 673 | |
| 674 | if patch_size_t is None: |
| 675 | # CogVideoX 1.0 checkpoints |
| 676 | self.proj = nn.Conv2d( |
| 677 | in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias |
| 678 | ) |
| 679 | else: |
| 680 | # CogVideoX 1.5 checkpoints |
| 681 | self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim) |
| 682 | |
| 683 | self.text_proj = nn.Linear(text_embed_dim, embed_dim) |
| 684 | |
| 685 | if use_positional_embeddings or use_learned_positional_embeddings: |
| 686 | persistent = use_learned_positional_embeddings |
| 687 | pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames) |
| 688 | self.register_buffer("pos_embedding", pos_embedding, persistent=persistent) |
| 689 | |
| 690 | def _get_positional_embeddings( |
| 691 | self, sample_height: int, sample_width: int, sample_frames: int, device: torch.device | None = None |
nothing calls this directly
no test coverage detected