sfast

README.md

@@ -1,6 +1,6 @@
 ---
 title: Real-Time Text-to-Image SDXL Lightning
-emoji: 🏆
+emoji: ⚡️⚡️⚡️⚡️
 colorFrom: yellow
 colorTo: purple
 sdk: gradio

app.py

@@ -1,18 +1,18 @@
-from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
+from diffusers import (
+    StableDiffusionXLPipeline,
+    EulerDiscreteScheduler,
+    UNet2DConditionModel,
+)
 import torch
 import os
 from huggingface_hub import hf_hub_download
 
-try:
-    import intel_extension_for_pytorch as ipex
-except:
-    pass
 
 from PIL import Image
 import gradio as gr
 import time
 from safetensors.torch import load_file
-
+from sfast.compilers.diffusion_pipeline_compiler import compile, CompilationConfig
 
 # Constants
 BASE = "stabilityai/stable-diffusion-xl-base-1.0"
@@ -30,37 +30,45 @@ CHECKPOINT = "sdxl_lightning_1step_unet_x0.safetensors"
 
 TORCH_COMPILE = os.environ.get("TORCH_COMPILE", "0") == "1"
 # check if MPS is available OSX only M1/M2/M3 chips
-mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
-xpu_available = hasattr(torch, "xpu") and torch.xpu.is_available()
-device = torch.device(
-    "cuda" if torch.cuda.is_available() else "xpu" if xpu_available else "cpu"
-)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 torch_device = device
 torch_dtype = torch.float16
 
 print(f"TORCH_COMPILE: {TORCH_COMPILE}")
 print(f"device: {device}")
 
-if mps_available:
-    device = torch.device("mps")
-    torch_device = "cpu"
-    torch_dtype = torch.float32
-
-
+# Load model.
+unet = UNet2DConditionModel.from_config(BASE, subfolder="unet").to(device, torch_dtype)
+unet.load_state_dict(load_file(hf_hub_download(REPO, CHECKPOINT), device="cuda"))
 pipe = StableDiffusionXLPipeline.from_pretrained(
-    BASE, torch_dtype=torch.float16, variant="fp16"
-)
+    BASE, unet=unet, torch_dtype=torch_dtype, variant="fp16"
+).to(device)
 
+# Ensure sampler uses "trailing" timesteps and "sample" prediction type.
 pipe.scheduler = EulerDiscreteScheduler.from_config(
     pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample"
 )
+pipe.set_progress_bar_config(disable=True)
+config = CompilationConfig.Default()
+try:
+    import xformers
 
-pipe.unet.load_state_dict(
-    torch.load(load_file(hf_hub_download(REPO, CHECKPOINT)), map_location="cuda")
-)
+    config.enable_xformers = True
+except ImportError:
+    print("xformers not installed, skip")
+try:
+    import triton
 
-pipe.to(device=torch_device, dtype=torch_dtype).to(device)
-pipe.set_progress_bar_config(disable=True)
+    config.enable_triton = True
+except ImportError:
+    print("Triton not installed, skip")
+# CUDA Graph is suggested for small batch sizes and small resolutions to reduce CPU overhead.
+# But it can increase the amount of GPU memory used.
+# For StableVideoDiffusionPipeline it is not needed.
+config.enable_cuda_graph = True
+
+pipe = compile(pipe, config)
 
 
 def predict(prompt, seed=1231231):
@@ -71,8 +79,6 @@ def predict(prompt, seed=1231231):
         generator=generator,
         num_inference_steps=1,
         guidance_scale=0.0,
-        width=512,
-        height=512,
         # original_inference_steps=params.lcm_steps,
         output_type="pil",
     )
@@ -102,10 +108,10 @@ css = """
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="container"):
         gr.Markdown(
-            """# SDXL Turbo - Text To Image
+            """# SDXL-Lightning- Text To Image 1-Step
             ## Unofficial Demo
-            SDXL Turbo model can generate high quality images in a single pass read more on [stability.ai post](https://stability.ai/news/stability-ai-sdxl-turbo).  
-            **Model**: https://huggingface.co/stabilityai/sdxl-turbo
+            SDXL-Lightining https://huggingface.co/ByteDance/SDXL-Lightning   
+            **Model**: https://huggingface.co/ByteDance/SDXL-Lightning
             """,
             elem_id="intro",
         )
@@ -123,8 +129,8 @@ with gr.Blocks(css=css) as demo:
             )
         with gr.Accordion("Run with diffusers"):
             gr.Markdown(
-                """## Running SDXL Turbo with `diffusers`
-            ```py
+                """## Running SDXL-Lightning with `diffusers`
+```py
 import torch
 from diffusers import (
     StableDiffusionXLPipeline,
@@ -156,8 +162,7 @@ pipe.scheduler = EulerDiscreteScheduler.from_config(
 pipe("A girl smiling", num_inference_steps=1, guidance_scale=0).images[0].save(
     "output.png"
 )
-
-            ```
+```
             """
             )
 

requirements.txt

@@ -1,7 +1,7 @@
 diffusers==0.26.3
 transformers
 gradio==4.19.1
-torch==2.2.0
+torch==2.1.0
 fastapi==0.104.0
 uvicorn==0.23.2
 Pillow==10.1.0
@@ -12,4 +12,5 @@ peft==0.6.0
 xformers
 hf_transfer
 huggingface_hub
-safetensors
+safetensors
+stable_fast @ https://github.com/chengzeyi/stable-fast/releases/download/v1.0.2/stable_fast-1.0.2+torch211cu121-cp310-cp310-manylinux2014_x86_64.whl; sys_platform != 'darwin' or platform_machine != 'arm64'