|
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") |
|
|
|
|
|
|
|
taesd.save_config("./") |
|
|
|
|
|
|
|
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.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") |
|
|
