| 963 | |
| 964 | |
| 965 | class UpBlockTemporalDecoder(nn.Module): |
| 966 | def __init__( |
| 967 | self, |
| 968 | in_channels: int, |
| 969 | out_channels: int, |
| 970 | num_layers: int = 1, |
| 971 | add_upsample: bool = True, |
| 972 | ): |
| 973 | super().__init__() |
| 974 | resnets = [] |
| 975 | for i in range(num_layers): |
| 976 | input_channels = in_channels if i == 0 else out_channels |
| 977 | |
| 978 | resnets.append( |
| 979 | SpatioTemporalResBlock( |
| 980 | in_channels=input_channels, |
| 981 | out_channels=out_channels, |
| 982 | temb_channels=None, |
| 983 | eps=1e-6, |
| 984 | temporal_eps=1e-5, |
| 985 | merge_factor=0.0, |
| 986 | merge_strategy="learned", |
| 987 | switch_spatial_to_temporal_mix=True, |
| 988 | ) |
| 989 | ) |
| 990 | self.resnets = nn.ModuleList(resnets) |
| 991 | |
| 992 | if add_upsample: |
| 993 | self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) |
| 994 | else: |
| 995 | self.upsamplers = None |
| 996 | |
| 997 | def forward( |
| 998 | self, |
| 999 | hidden_states: torch.Tensor, |
| 1000 | image_only_indicator: torch.Tensor, |
| 1001 | ) -> torch.Tensor: |
| 1002 | for resnet in self.resnets: |
| 1003 | hidden_states = resnet( |
| 1004 | hidden_states, |
| 1005 | image_only_indicator=image_only_indicator, |
| 1006 | ) |
| 1007 | |
| 1008 | if self.upsamplers is not None: |
| 1009 | for upsampler in self.upsamplers: |
| 1010 | hidden_states = upsampler(hidden_states) |
| 1011 | |
| 1012 | return hidden_states |
| 1013 | |
| 1014 | |
| 1015 | class UNetMidBlockSpatioTemporal(nn.Module): |