add forward for vq model

0b7daa6d · patil-suraj · 99568c5a · 0b7daa6d
Commit 0b7daa6d authored Jun 29, 2022 by patil-suraj
Show whitespace changes
Inline Side-by-side

Showing with 5 additions and 0 deletions

src/diffusers/models/vae.py src/diffusers/models/vae.py +5 -0

No files found.
--- a/src/diffusers/models/vae.py
+++ b/src/diffusers/models/vae.py
@@ -535,6 +535,11 @@ class VQModel(ModelMixin, ConfigMixin):
        dec = self.decoder(quant)
        return dec
    
+    def forward(self, x):
+        h = self.encode(x)
+        dec = self.decode(h)
+        return dec
+

 class AutoencoderKL(ModelMixin, ConfigMixin):
    def __init__(