Update app.py

app.py

@@ -9,9 +9,10 @@ import openvino.runtime as ov
 from typing import Optional, Dict
 from diffusers import EulerAncestralDiscreteScheduler, LCMScheduler, DDIMScheduler
 
-#LCMScheduler 產生垃圾
+
+
 #EulerDiscreteScheduler 尚可
-#EulerAncestralDiscreteScheduler 很不錯chatgpt推薦
+#EulerAncestralDiscreteScheduler 
 
 
 model_id = "hsuwill000/LCM-absolutereality-openvino-8bit"
@@ -19,12 +20,20 @@ model_id = "hsuwill000/LCM-absolutereality-openvino-8bit"
 #model_id = "spamsoms/LCM-anything-v5-openvino2"
 #adapter_id = "latent-consistency/lcm-lora-sdv1-5"
 
-#512*512 好 太大會變形
+
 HIGH=1024
-WIDTH=512
+WIDTH=1024
 
 batch_size = -1
 
+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, 
@@ -36,11 +45,19 @@ pipe = OVStableDiffusionPipeline.from_pretrained(
         use_safetensors=False,
         )
 print(pipe.scheduler.compatibles)
+
 #pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
 #pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
 
 #pipe.load_lora_weights(adapter_id)
 #pipe.fuse_lora()
+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=-1, height=HIGH, width=WIDTH, num_images_per_prompt=1)
@@ -93,7 +110,7 @@ with gr.Blocks(css=css) as demo:
     
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
-        # LCM-absolutereality-openvino-8bit {WIDTH}x{HIGH}
+        # {model_id.split('/')[1]} {WIDTH}x{HIGH}
         Currently running on {power_device}.
         """)