File size: 856 Bytes
f753ab8 06aec4b f753ab8 06aec4b f753ab8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
from diffusers import AutoencoderTiny
from optimum.exporters.openvino import export
from optimum.exporters.onnx.model_configs import VaeDecoderOnnxConfig, VaeEncoderOnnxConfig
taesd = AutoencoderTiny.from_pretrained("madebyollin/taesd")
# Config in root of repo
taesd.save_config("./")
# TAESD Decoder
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: {"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")
|