r""" A module to blend spatial and temporal features. Parameters: alpha (`float`): The initial value of the blending factor. merge_strategy (`str`, *optional*, defaults to `learned_with_images`): The merge strategy to use for the temporal mixing. switch_s
| 721 | |
| 722 | |
| 723 | class AlphaBlender(nn.Module): |
| 724 | r""" |
| 725 | A module to blend spatial and temporal features. |
| 726 | |
| 727 | Parameters: |
| 728 | alpha (`float`): The initial value of the blending factor. |
| 729 | merge_strategy (`str`, *optional*, defaults to `learned_with_images`): |
| 730 | The merge strategy to use for the temporal mixing. |
| 731 | switch_spatial_to_temporal_mix (`bool`, *optional*, defaults to `False`): |
| 732 | If `True`, switch the spatial and temporal mixing. |
| 733 | """ |
| 734 | |
| 735 | strategies = ["learned", "fixed", "learned_with_images"] |
| 736 | |
| 737 | def __init__( |
| 738 | self, |
| 739 | alpha: float, |
| 740 | merge_strategy: str = "learned_with_images", |
| 741 | switch_spatial_to_temporal_mix: bool = False, |
| 742 | ): |
| 743 | super().__init__() |
| 744 | self.merge_strategy = merge_strategy |
| 745 | self.switch_spatial_to_temporal_mix = switch_spatial_to_temporal_mix # For TemporalVAE |
| 746 | |
| 747 | if merge_strategy not in self.strategies: |
| 748 | raise ValueError(f"merge_strategy needs to be in {self.strategies}") |
| 749 | |
| 750 | if self.merge_strategy == "fixed": |
| 751 | self.register_buffer("mix_factor", torch.Tensor([alpha])) |
| 752 | elif self.merge_strategy == "learned" or self.merge_strategy == "learned_with_images": |
| 753 | self.register_parameter("mix_factor", torch.nn.Parameter(torch.Tensor([alpha]))) |
| 754 | else: |
| 755 | raise ValueError(f"Unknown merge strategy {self.merge_strategy}") |
| 756 | |
| 757 | def get_alpha(self, image_only_indicator: torch.Tensor, ndims: int) -> torch.Tensor: |
| 758 | if self.merge_strategy == "fixed": |
| 759 | alpha = self.mix_factor |
| 760 | |
| 761 | elif self.merge_strategy == "learned": |
| 762 | alpha = torch.sigmoid(self.mix_factor) |
| 763 | |
| 764 | elif self.merge_strategy == "learned_with_images": |
| 765 | if image_only_indicator is None: |
| 766 | raise ValueError("Please provide image_only_indicator to use learned_with_images merge strategy") |
| 767 | |
| 768 | alpha = torch.where( |
| 769 | image_only_indicator.bool(), |
| 770 | torch.ones(1, 1, device=image_only_indicator.device), |
| 771 | torch.sigmoid(self.mix_factor)[..., None], |
| 772 | ) |
| 773 | |
| 774 | # (batch, channel, frames, height, width) |
| 775 | if ndims == 5: |
| 776 | alpha = alpha[:, None, :, None, None] |
| 777 | # (batch*frames, height*width, channels) |
| 778 | elif ndims == 3: |
| 779 | alpha = alpha.reshape(-1)[:, None, None] |
| 780 | else: |