Add upsample_size to AttnUpBlock2D, AttnDownBlock2D (huggingface#3275)

The argument `upsample_size` needs to be added to these modules to allow compatibility with other blocks that require this argument.

Author: will-rice

models/unet_2d_blocks.py

@@ -734,7 +734,7 @@ def __init__(
         else:
             self.downsamplers = None
 
-    def forward(self, hidden_states, temb=None):
+    def forward(self, hidden_states, temb=None, upsample_size=None):
         output_states = ()
 
         for resnet, attn in zip(self.resnets, self.attentions):
@@ -1720,7 +1720,7 @@ def __init__(
         else:
             self.upsamplers = None
 
-    def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
+    def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None):
         for resnet, attn in zip(self.resnets, self.attentions):
             # pop res hidden states
             res_hidden_states = res_hidden_states_tuple[-1]