add the option of upsample function for tiny vae (#7604)

Co-authored-by: Sayak Paul <[email protected]>
Co-authored-by: YiYi Xu <[email protected]>

Author: 3 people

src/diffusers/models/autoencoders/autoencoder_tiny.py

@@ -102,6 +102,7 @@ def __init__(
         encoder_block_out_channels: Tuple[int, ...] = (64, 64, 64, 64),
         decoder_block_out_channels: Tuple[int, ...] = (64, 64, 64, 64),
         act_fn: str = "relu",
+        upsample_fn: str = "nearest",
         latent_channels: int = 4,
         upsampling_scaling_factor: int = 2,
         num_encoder_blocks: Tuple[int, ...] = (1, 3, 3, 3),
@@ -133,6 +134,7 @@ def __init__(
             block_out_channels=decoder_block_out_channels,
             upsampling_scaling_factor=upsampling_scaling_factor,
             act_fn=act_fn,
+            upsample_fn=upsample_fn,
         )
 
         self.latent_magnitude = latent_magnitude

src/diffusers/models/autoencoders/vae.py

@@ -926,6 +926,7 @@ def __init__(
         block_out_channels: Tuple[int, ...],
         upsampling_scaling_factor: int,
         act_fn: str,
+        upsample_fn: str,
     ):
         super().__init__()
 
@@ -942,7 +943,7 @@ def __init__(
                 layers.append(AutoencoderTinyBlock(num_channels, num_channels, act_fn))
 
             if not is_final_block:
-                layers.append(nn.Upsample(scale_factor=upsampling_scaling_factor))
+                layers.append(nn.Upsample(scale_factor=upsampling_scaling_factor, mode=upsample_fn))
 
             conv_out_channel = num_channels if not is_final_block else out_channels
             layers.append(