Use TAESD

app.py

@@ -9,11 +9,14 @@ import gradio as gr
 import numpy as np
 import PIL.Image
 
+from huggingface_hub import snapshot_download
 from diffusers import DiffusionPipeline
 
 from lcm_scheduler import LCMScheduler
 from lcm_ov_pipeline import OVLatentConsistencyModelPipeline
 
+from optimum.intel.openvino.modeling_diffusion import OVModelVaeDecoder, OVBaseModel
+
 import os
 from tqdm import tqdm
 import gradio_user_history as gr_user_history
@@ -38,8 +41,20 @@ width = int(os.getenv("IMAGE_WIDTH", "512"))
 height = int(os.getenv("IMAGE_HEIGHT", "512"))
 num_images = int(os.getenv("NUM_IMAGES", "1"))
 
+class CustomOVModelVaeDecoder(OVModelVaeDecoder):
+    def __init__(
+        self, model: openvino.runtime.Model, parent_model: OVBaseModel, ov_config: Optional[Dict[str, str]] = None, model_dir: str = None,
+    ):
+        super(OVModelVaeDecoder, self).__init__(model, parent_model, ov_config, "vae_decoder", model_dir)
+
 scheduler = LCMScheduler.from_pretrained(model_id, subfolder="scheduler")
 pipe = OVLatentConsistencyModelPipeline.from_pretrained(model_id, scheduler = scheduler, compile = False, ov_config = {"CACHE_DIR":""})
+
+# Inject TAESD
+
+taesd_dir = snapshot_download(repo_id="deinferno/taesd-openvino")
+pipe.vae_decoder = CustomOVModelVaeDecoder(model = OVBaseModel.load_model(f"{taesd_dir}/vae_decoder/openvino_model.xml"), parent_model = pipe, model_dir = taesd_dir)
+
 pipe.reshape(batch_size=batch_size, height=height, width=width, num_images_per_prompt=num_images)
 pipe.compile()