[Type Hint] VAE models (huggingface#365)

daspartho · anton-l · web-flow · commit 38fff9e2da60 · 2022-09-05T19:09:48.000+02:00
* [Type Hint] VAE models

* Update src/diffusers/models/vae.py

* apply suggestions from code review

Co-authored-by: Anton Lozhkov &lt;aglozhkov@gmail.com&gt;
diff --git a/models/vae.py b/models/vae.py
@@ -411,7 +411,7 @@ def __init__(
             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)
 

Original file line number	Diff line number	Diff line change
`@@ -411,7 +411,7 @@ def __init__(`
`411`	`411`	`act_fn=act_fn,`
`412`	`412`	`)`
`413`	`413`
`414`		`- def encode(self, x, return_dict: bool = True):`
	`414`	`+ def encode(self, x: torch.FloatTensor, return_dict: bool = True) -> VQEncoderOutput:`
`415`	`415`	`h = self.encoder(x)`
`416`	`416`	`h = self.quant_conv(h)`
`417`	`417`