Fix encoder

convert_to_openvino.py

@@ -11,12 +11,12 @@ taesd.save_config("./")
 
 # TAESD Decoder
 
-taesd.forward = lambda latent_sample: taesd.decode(latent_sample)
+taesd.forward = lambda latent_sample: taesd.decode(x=latent_sample)
 export(model = taesd, config = VaeDecoderOnnxConfig( config = taesd.config, task = "semantic-segmentation"), output = "./vae_decoder/openvino_model.xml")
 taesd.save_config("./vae_decoder")
 
 # TAESD Encoder
 
-taesd.forward = lambda sample: taesd.encode(sample)
+taesd.forward = lambda sample: {"latent_sample": taesd.encode(x=sample)["latents"]}
 export(model = taesd, config = VaeEncoderOnnxConfig( config = taesd.config, task = "semantic-segmentation"), output = "./vae_encoder/openvino_model.xml")
 taesd.save_config("./vae_encoder")