Update app.py

app.py

@@ -22,9 +22,10 @@ import paramiko
 import gc
 import time
 import datetime
-from diffusers.schedulers import AysSchedules
+#from diffusers.schedulers import AysSchedules
 
 from gradio import themes
+from hidiffusion import apply_hidiffusion, remove_hidiffusion
 
 torch.backends.cuda.matmul.allow_tf32 = False
 torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
@@ -91,7 +92,7 @@ DEFAULT_STYLE_NAME = "Style Zero"
 STYLE_NAMES = list(styles.keys())
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"]
+#sampling_schedule = AysSchedules["StableDiffusionXLTimesteps"]
 
 def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
     if style_name in styles:
@@ -149,6 +150,8 @@ def load_and_prepare_model(model_id):
     pipe.to(device)
     pipe.to(torch.bfloat16)
     
+    apply_hidiffusion(pipe)
+
     pipe.unet.set_default_attn_processor()
     pipe.vae.set_default_attn_processor()
     
@@ -229,7 +232,7 @@ def generate_30(
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-        "timesteps": sampling_schedule,
+     #   "timesteps": sampling_schedule,
         "output_type": "pil",
     }
     if use_resolution_binning:
@@ -296,7 +299,7 @@ def generate_60(
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-        "timesteps": sampling_schedule,
+     #   "timesteps": sampling_schedule,
         "output_type": "pil",
     }
     if use_resolution_binning:
@@ -363,7 +366,7 @@ def generate_90(
         "guidance_scale": guidance_scale,
         "num_inference_steps": num_inference_steps,
         "generator": generator,
-        "timesteps": sampling_schedule,
+     #   "timesteps": sampling_schedule,
         "output_type": "pil",
     }
     if use_resolution_binning: