sfast

app.py

@@ -12,6 +12,11 @@ import numpy as np
 import gradio as gr
 import psutil
 import time
+from sfast.compilers.stable_diffusion_pipeline_compiler import (
+    compile,
+    CompilationConfig,
+)
+
 
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
 TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)
@@ -40,22 +45,18 @@ else:
     pipe = DiffusionPipeline.from_pretrained("Lykon/dreamshaper-7", safety_checker=None)
 
 pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
-pipe.to(device=torch_device, dtype=torch_dtype).to(device)
 pipe.unet.to(memory_format=torch.channels_last)
 pipe.set_progress_bar_config(disable=True)
-
-# check if computer has less than 64GB of RAM using sys or os
-if psutil.virtual_memory().total < 64 * 1024**3:
-    pipe.enable_attention_slicing()
-
-if TORCH_COMPILE:
-    pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-    pipe.vae = torch.compile(pipe.vae, mode="reduce-overhead", fullgraph=True)
-    pipe(prompt="warmup", num_inference_steps=1, guidance_scale=8.0)
-
-# Load LCM LoRA
 pipe.load_lora_weights("latent-consistency/lcm-lora-sdv1-5")
 pipe.fuse_lora()
+pipe.to(device=torch_device, dtype=torch_dtype).to(device)
+
+
+config = CompilationConfig.Default()
+config.enable_xformers = True
+config.enable_triton = True
+config.enable_cuda_graph = True
+pipe = compile(pipe, config=config)
 
 
 def predict(prompt, guidance, steps, seed=1231231):

requirements.txt

@@ -1,4 +1,5 @@
-diffusers==0.23.0
+# diffusers==0.22.2
+git+https://github.com/huggingface/diffusers.git@6110d7c95f630479cf01340cc8a8141c1e359f09
 transformers==4.34.1
 gradio==4.1.2
 --extra-index-url https://download.pytorch.org/whl/cu121
@@ -10,4 +11,6 @@ accelerate==0.24.0
 compel==2.0.2
 controlnet-aux==0.0.7
 peft==0.6.0
-xformers
+stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v0.0.15.post1/stable_fast-0.0.15.post1+torch211cu121-cp310-cp310-manylinux2014_x86_64.whl
+xformers
+triton