Addition of non_quantized flag to enable non-quantised encoding for L… (#474)

* Addition of non_quantized flag to enable non-quantised encoding for LDM training on the VQVAE.

* non_quantized > quantized
Default true.

---------

Co-authored-by: virginiafdez <virginia.fernandez@kcl.ac.uk>

Author: virginiafdez

generative/networks/nets/vqvae.py

@@ -442,10 +442,12 @@ def forward(self, images: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
 
         return reconstruction, quantization_losses
 
-    def encode_stage_2_inputs(self, x: torch.Tensor) -> torch.Tensor:
+    def encode_stage_2_inputs(self, x: torch.Tensor, quantized: bool = True) -> torch.Tensor:
         z = self.encode(x)
         e, _ = self.quantize(z)
-        return e
+        if quantized:
+            return e
+        return z
 
     def decode_stage_2_outputs(self, z: torch.Tensor) -> torch.Tensor:
         e, _ = self.quantize(z)