| 1821 | |
| 1822 | class UpBlockTemporalDecoder(nn.Module): |
| 1823 | def __init__( |
| 1824 | self, |
| 1825 | in_channels: int, |
| 1826 | out_channels: int, |
| 1827 | num_layers: int = 1, |
| 1828 | add_upsample: bool = True, |
| 1829 | ): |
| 1830 | super().__init__() |
| 1831 | resnets = [] |
| 1832 | for i in range(num_layers): |
| 1833 | input_channels = in_channels if i == 0 else out_channels |
| 1834 | |
| 1835 | resnets.append( |
| 1836 | SpatioTemporalResBlock( |
| 1837 | in_channels=input_channels, |
| 1838 | out_channels=out_channels, |
| 1839 | temb_channels=None, |
| 1840 | eps=1e-6, |
| 1841 | temporal_eps=1e-5, |
| 1842 | merge_factor=0.0, |
| 1843 | merge_strategy="learned", |
| 1844 | switch_spatial_to_temporal_mix=True, |
| 1845 | ) |
| 1846 | ) |
| 1847 | self.resnets = nn.ModuleList(resnets) |
| 1848 | |
| 1849 | if add_upsample: |
| 1850 | self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)]) |
| 1851 | else: |
| 1852 | self.upsamplers = None |
| 1853 | |
| 1854 | def forward( |
| 1855 | self, |