Fix mypy errors in diffusion labs/models/vae (#508)

Summary:
This is causing CI failures on unrelated PRs.

Pull Request resolved: #508

Reviewed By: kartikayk

Differential Revision: D50935831

Pulled By: ebsmothers

fbshipit-source-id: 93f6df64a8ee8795473dff023825a637a34f7727

torchmultimodal/diffusion_labs/models/vae/res_block.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 from torch import nn, Tensor
@@ -83,15 +83,18 @@ def __init__(
 
         if use_downsample and use_upsample:
             raise ValueError("Cannot use both upsample and downsample in res block")
-        elif use_downsample:
-            hidden_updownsample_layer = nn.AvgPool2d(kernel_size=2, stride=2)
-            skip_updownsample_layer = nn.AvgPool2d(kernel_size=2, stride=2)
-        elif use_upsample:
-            hidden_updownsample_layer = nn.Upsample(scale_factor=2, mode="nearest")
-            skip_updownsample_layer = nn.Upsample(scale_factor=2, mode="nearest")
         else:
-            hidden_updownsample_layer = nn.Identity()
-            skip_updownsample_layer = nn.Identity()
+            hidden_updownsample_layer: Union[nn.AvgPool2d, nn.Upsample, nn.Identity]
+            skip_updownsample_layer: Union[nn.AvgPool2d, nn.Upsample, nn.Identity]
+            if use_downsample:
+                hidden_updownsample_layer = nn.AvgPool2d(kernel_size=2, stride=2)
+                skip_updownsample_layer = nn.AvgPool2d(kernel_size=2, stride=2)
+            if use_upsample:
+                hidden_updownsample_layer = nn.Upsample(scale_factor=2, mode="nearest")
+                skip_updownsample_layer = nn.Upsample(scale_factor=2, mode="nearest")
+            else:
+                hidden_updownsample_layer = nn.Identity()
+                skip_updownsample_layer = nn.Identity()
 
         self.cond_proj = cond_proj
         self.in_block = nn.Sequential(

torchmultimodal/diffusion_labs/models/vae/residual_sampling.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from typing import Tuple, Union
+
 from torch import nn as nn, Tensor
 from torch.nn import functional as F
 
@@ -52,6 +54,7 @@ def __init__(
         asymmetric_padding: bool = True,
     ):
         super().__init__()
+        padding: Union[int, Tuple[int, int, int, int]]
         if asymmetric_padding:
             padding = (0, 1, 0, 1)
         else:

torchmultimodal/diffusion_labs/models/vae/vae.py

@@ -78,15 +78,15 @@ def ldm_variational_autoencoder(
     embedding_channels: int,
     in_channels: int,
     out_channels: int,
-    z_channels,
+    z_channels: int,
     channels: int,
     num_res_blocks: int,
     channel_multipliers: Sequence[int] = (1, 2, 4, 8),
     dropout: float = 0.0,
     norm_groups: int = 32,
     norm_eps: float = 1e-6,
     output_alpha_channel: bool = False,
-):
+) -> VariationalAutoencoder:
     encoder = nn.Sequential(
         # pyre-ignore
         OrderedDict(