| 1758 | |
| 1759 | class MidBlockTemporalDecoder(nn.Module): |
| 1760 | def __init__( |
| 1761 | self, |
| 1762 | in_channels: int, |
| 1763 | out_channels: int, |
| 1764 | attention_head_dim: int = 512, |
| 1765 | num_layers: int = 1, |
| 1766 | upcast_attention: bool = False, |
| 1767 | ): |
| 1768 | super().__init__() |
| 1769 | |
| 1770 | resnets = [] |
| 1771 | attentions = [] |
| 1772 | for i in range(num_layers): |
| 1773 | input_channels = in_channels if i == 0 else out_channels |
| 1774 | resnets.append( |
| 1775 | SpatioTemporalResBlock( |
| 1776 | in_channels=input_channels, |
| 1777 | out_channels=out_channels, |
| 1778 | temb_channels=None, |
| 1779 | eps=1e-6, |
| 1780 | temporal_eps=1e-5, |
| 1781 | merge_factor=0.0, |
| 1782 | merge_strategy="learned", |
| 1783 | switch_spatial_to_temporal_mix=True, |
| 1784 | ) |
| 1785 | ) |
| 1786 | |
| 1787 | attentions.append( |
| 1788 | Attention( |
| 1789 | query_dim=in_channels, |
| 1790 | heads=in_channels // attention_head_dim, |
| 1791 | dim_head=attention_head_dim, |
| 1792 | eps=1e-6, |
| 1793 | upcast_attention=upcast_attention, |
| 1794 | norm_num_groups=32, |
| 1795 | bias=True, |
| 1796 | residual_connection=True, |
| 1797 | ) |
| 1798 | ) |
| 1799 | |
| 1800 | self.attentions = nn.ModuleList(attentions) |
| 1801 | self.resnets = nn.ModuleList(resnets) |
| 1802 | |
| 1803 | def forward( |
| 1804 | self, |