1 files changed, 4 insertions, 3 deletions

diff --git a/infer.py b/infer.py
index 329c60b..13219f8 100644
--- a/infer.py
+++ b/infer.py
@@ -21,7 +21,8 @@ from diffusers import (
     LMSDiscreteScheduler,
     EulerAncestralDiscreteScheduler,
     KDPM2DiscreteScheduler,
-    KDPM2AncestralDiscreteScheduler
+    KDPM2AncestralDiscreteScheduler,
+    UniPCMultistepScheduler
 )
 from transformers import CLIPTextModel
 
@@ -29,7 +30,6 @@ from data.keywords import prompt_to_keywords, keywords_to_prompt
 from models.clip.embeddings import patch_managed_embeddings
 from models.clip.tokenizer import MultiCLIPTokenizer
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
-from schedulers.scheduling_unipc_multistep import UniPCMultistepScheduler
 from util import load_config, load_embeddings_from_dir
 
 
@@ -245,7 +245,8 @@ def create_pipeline(model, dtype):
         tokenizer=tokenizer,
         scheduler=scheduler,
     )
-    pipeline.enable_xformers_memory_efficient_attention()
+    # pipeline.enable_xformers_memory_efficient_attention()
+    pipeline.unet = torch.compile(pipeline.unet)
     pipeline.enable_vae_slicing()
     pipeline.to("cuda")