Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import torch
-from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline
+from optimum.intel.openvino.modeling_diffusion import OVModelVaeDecoder, OVBaseModel, OVStableDiffusionPipeline
+from huggingface_hub import snapshot_download
 
 model_id = "hsuwill000/Fluently-v4-LCM-openvino"
 
@@ -9,6 +10,14 @@ WIDTH = 512
 
 batch_size = -1  # Or set it to a specific positive integer if needed
 
+class CustomOVModelVaeDecoder(OVModelVaeDecoder):
+    def __init__(
+        self, model: ov.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)
+
+
+
 pipe = OVStableDiffusionPipeline.from_pretrained(
     model_id,
     compile=False,
@@ -17,6 +26,13 @@ pipe = OVStableDiffusionPipeline.from_pretrained(
     safety_checker=None,
     use_safetensors=False,
 )
+
+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, 
+                                           
+
 print(pipe.scheduler.compatibles)
 
 pipe.reshape(batch_size=batch_size, height=HIGH, width=WIDTH, num_images_per_prompt=1)