(
self,
patch_size: int = 2,
patch_size_t: Optional[int] = 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,
extra_encoder_cond_channels: int = -1,
use_FrameIn: bool = False,
)
| 625 | |
| 626 | class CogVideoXPatchEmbed(nn.Module): |
| 627 | def __init__( |
| 628 | self, |
| 629 | patch_size: int = 2, |
| 630 | patch_size_t: Optional[int] = None, |
| 631 | in_channels: int = 16, |
| 632 | embed_dim: int = 1920, |
| 633 | text_embed_dim: int = 4096, |
| 634 | bias: bool = True, |
| 635 | sample_width: int = 90, |
| 636 | sample_height: int = 60, |
| 637 | sample_frames: int = 49, |
| 638 | temporal_compression_ratio: int = 4, |
| 639 | max_text_seq_length: int = 226, |
| 640 | spatial_interpolation_scale: float = 1.875, |
| 641 | temporal_interpolation_scale: float = 1.0, |
| 642 | use_positional_embeddings: bool = True, |
| 643 | use_learned_positional_embeddings: bool = True, |
| 644 | extra_encoder_cond_channels: int = -1, |
| 645 | use_FrameIn: bool = False, |
| 646 | ) -> None: |
| 647 | super().__init__() |
| 648 | |
| 649 | self.patch_size = patch_size |
| 650 | self.patch_size_t = patch_size_t |
| 651 | self.embed_dim = embed_dim |
| 652 | self.sample_height = sample_height |
| 653 | self.sample_width = sample_width |
| 654 | self.sample_frames = sample_frames |
| 655 | self.temporal_compression_ratio = temporal_compression_ratio |
| 656 | self.max_text_seq_length = max_text_seq_length |
| 657 | self.spatial_interpolation_scale = spatial_interpolation_scale |
| 658 | self.temporal_interpolation_scale = temporal_interpolation_scale |
| 659 | self.use_positional_embeddings = use_positional_embeddings |
| 660 | self.use_learned_positional_embeddings = use_learned_positional_embeddings |
| 661 | self.use_FrameIn = use_FrameIn |
| 662 | |
| 663 | |
| 664 | if patch_size_t is None: |
| 665 | # CogVideoX 1.0 checkpoints |
| 666 | self.proj = nn.Conv2d( |
| 667 | in_channels, embed_dim, kernel_size=(patch_size, patch_size), stride=patch_size, bias=bias |
| 668 | ) |
| 669 | else: |
| 670 | # CogVideoX 1.5 checkpoints |
| 671 | self.proj = nn.Linear(in_channels * patch_size * patch_size * patch_size_t, embed_dim) |
| 672 | |
| 673 | # Extra channels for the motion abjustment |
| 674 | |
| 675 | self.text_proj = nn.Linear(text_embed_dim, embed_dim) |
| 676 | |
| 677 | |
| 678 | if use_positional_embeddings or use_learned_positional_embeddings: |
| 679 | persistent = use_learned_positional_embeddings |
| 680 | pos_embedding = self._get_positional_embeddings(sample_height, sample_width, sample_frames) |
| 681 | self.register_buffer("pos_embedding", pos_embedding, persistent=persistent) |
| 682 | # print("mean value of pos_embedding is ", torch.mean(pos_embedding)) |
| 683 | |
| 684 | # if use_FrameIn: # Position Embedding for the reference |
nothing calls this directly
no test coverage detected