[Type Hint] VAE models (#365)

* [Type Hint] VAE models

* Update src/diffusers/models/vae.py

* apply suggestions from code review
Co-authored-by: Anton Lozhkov <aglozhkov@gmail.com>

src/diffusers/models/vae.py

@@ -411,7 +411,7 @@ class VQModel(ModelMixin, ConfigMixin):
             act_fn=act_fn,
         )
 
-    def encode(self, x, return_dict: bool = True):
+    def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> VQEncoderOutput:
         h = self.encoder(x)
         h = self.quant_conv(h)