r""" A SpatioTemporal Resnet block. Parameters: in_channels (`int`): The number of channels in the input. out_channels (`int`, *optional*, default to be `None`): The number of output channels for the first conv2d layer. If None, same as `in_channels`. tem
| 633 | |
| 634 | # VideoResBlock |
| 635 | class SpatioTemporalResBlock(nn.Module): |
| 636 | r""" |
| 637 | A SpatioTemporal Resnet block. |
| 638 | |
| 639 | Parameters: |
| 640 | in_channels (`int`): The number of channels in the input. |
| 641 | out_channels (`int`, *optional*, default to be `None`): |
| 642 | The number of output channels for the first conv2d layer. If None, same as `in_channels`. |
| 643 | temb_channels (`int`, *optional*, default to `512`): the number of channels in timestep embedding. |
| 644 | eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the spatial resenet. |
| 645 | temporal_eps (`float`, *optional*, defaults to `eps`): The epsilon to use for the temporal resnet. |
| 646 | merge_factor (`float`, *optional*, defaults to `0.5`): The merge factor to use for the temporal mixing. |
| 647 | merge_strategy (`str`, *optional*, defaults to `learned_with_images`): |
| 648 | The merge strategy to use for the temporal mixing. |
| 649 | switch_spatial_to_temporal_mix (`bool`, *optional*, defaults to `False`): |
| 650 | If `True`, switch the spatial and temporal mixing. |
| 651 | """ |
| 652 | |
| 653 | def __init__( |
| 654 | self, |
| 655 | in_channels: int, |
| 656 | out_channels: Optional[int] = None, |
| 657 | temb_channels: int = 512, |
| 658 | eps: float = 1e-6, |
| 659 | temporal_eps: Optional[float] = None, |
| 660 | merge_factor: float = 0.5, |
| 661 | merge_strategy="learned_with_images", |
| 662 | switch_spatial_to_temporal_mix: bool = False, |
| 663 | ): |
| 664 | super().__init__() |
| 665 | |
| 666 | self.spatial_res_block = ResnetBlock2D( |
| 667 | in_channels=in_channels, |
| 668 | out_channels=out_channels, |
| 669 | temb_channels=temb_channels, |
| 670 | eps=eps, |
| 671 | ) |
| 672 | |
| 673 | self.temporal_res_block = TemporalResnetBlock( |
| 674 | in_channels=out_channels if out_channels is not None else in_channels, |
| 675 | out_channels=out_channels if out_channels is not None else in_channels, |
| 676 | temb_channels=temb_channels, |
| 677 | eps=temporal_eps if temporal_eps is not None else eps, |
| 678 | ) |
| 679 | |
| 680 | self.time_mixer = AlphaBlender( |
| 681 | alpha=merge_factor, |
| 682 | merge_strategy=merge_strategy, |
| 683 | switch_spatial_to_temporal_mix=switch_spatial_to_temporal_mix, |
| 684 | ) |
| 685 | |
| 686 | def forward( |
| 687 | self, |
| 688 | hidden_states: torch.Tensor, |
| 689 | temb: Optional[torch.Tensor] = None, |
| 690 | image_only_indicator: Optional[torch.Tensor] = None, |
| 691 | ): |
| 692 | num_frames = image_only_indicator.shape[-1] |