Lightning

README.md

@@ -1,10 +1,10 @@
 ---
-title: Unofficial SDXL Turbo Real Time Text To Image
+title: Real-Time Text-to-Image SDXL Lightning
 emoji: 🏆
 colorFrom: yellow
 colorTo: purple
 sdk: gradio
-sdk_version: 4.7.1
+sdk_version: 4.19.1
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,6 +1,7 @@
-from diffusers import DiffusionPipeline
+from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
 import torch
 import os
+from huggingface_hub import hf_hub_download
 
 try:
     import intel_extension_for_pytorch as ipex
@@ -8,14 +9,26 @@ except:
     pass
 
 from PIL import Image
-import numpy as np
 import gradio as gr
-import psutil
 import time
+from safetensors.torch import load_file
 
-SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", None)
-TORCH_COMPILE = os.environ.get("TORCH_COMPILE", None)
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
+
+# Constants
+BASE = "stabilityai/stable-diffusion-xl-base-1.0"
+REPO = "ByteDance/SDXL-Lightning"
+# 1-step
+CHECKPOINT = "sdxl_lightning_1step_unet_x0.safetensors"
+
+# {
+#     "1-Step": ["sdxl_lightning_1step_unet_x0.safetensors", 1],
+#     "2-Step": ["sdxl_lightning_2step_unet.safetensors", 2],
+#     "4-Step": ["sdxl_lightning_4step_unet.safetensors", 4],
+#     "8-Step": ["sdxl_lightning_8step_unet.safetensors", 8],
+# }
+
+
+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()
@@ -25,7 +38,6 @@ device = torch.device(
 torch_device = device
 torch_dtype = torch.float16
 
-print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
 print(f"TORCH_COMPILE: {TORCH_COMPILE}")
 print(f"device: {device}")
 
@@ -34,32 +46,30 @@ if mps_available:
     torch_device = "cpu"
     torch_dtype = torch.float32
 
-if SAFETY_CHECKER == "True":
-    pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/sdxl-turbo",
-        torch_dtype=torch_dtype,
-        variant="fp16" if torch_dtype == torch.float16 else "fp32")
-else:
-    pipe = DiffusionPipeline.from_pretrained(
-        "stabilityai/sdxl-turbo",
-        safety_checker=None,
-        torch_dtype=torch_dtype,
-        variant="fp16" if torch_dtype == torch.float16 else "fp32",
-    )
 
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    BASE, torch_dtype=torch.float16, variant="fp16"
+)
+
+pipe.scheduler = EulerDiscreteScheduler.from_config(
+    pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample"
+)
+
+pipe.unet.load_state_dict(
+    torch.load(load_file(hf_hub_download(REPO, CHECKPOINT)), map_location="cuda")
+)
 
 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)
 
 
-def predict(prompt, steps, seed=1231231):
+def predict(prompt, seed=1231231):
     generator = torch.manual_seed(seed)
     last_time = time.time()
     results = pipe(
         prompt=prompt,
         generator=generator,
-        num_inference_steps=steps,
+        num_inference_steps=1,
         guidance_scale=0.0,
         width=512,
         height=512,
@@ -108,37 +118,52 @@ with gr.Blocks(css=css) as demo:
 
         image = gr.Image(type="filepath")
         with gr.Accordion("Advanced options", open=False):
-            steps = gr.Slider(label="Steps", value=2, minimum=1, maximum=10, step=1)
             seed = gr.Slider(
                 randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
             )
         with gr.Accordion("Run with diffusers"):
             gr.Markdown(
                 """## Running SDXL Turbo with `diffusers`
-            ```bash
-            pip install diffusers==0.23.1
-            ```
             ```py
-            from diffusers import DiffusionPipeline
-
-            pipe = DiffusionPipeline.from_pretrained(
-                "stabilityai/sdxl-turbo", variant="fp16", torch_dtype=torch.float16
-            ).to("cuda")
-            results = pipe(
-                prompt="A cinematic shot of a baby racoon wearing an intricate italian priest robe",
-                num_inference_steps=1,
-                guidance_scale=0.0,
-            )
-            imga = results.images[0]
-            imga.save("image.png")
+import torch
+from diffusers import (
+    StableDiffusionXLPipeline,
+    UNet2DConditionModel,
+    EulerDiscreteScheduler,
+)
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+ckpt = "sdxl_lightning_1step_unet_x0.safetensors"  # Use the correct ckpt for your step setting!
+
+# Load model.
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(
+    "cuda", torch.float16
+)
+unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    base, unet=unet, torch_dtype=torch.float16, variant="fp16"
+).to("cuda")
+
+# Ensure sampler uses "trailing" timesteps and "sample" prediction type.
+pipe.scheduler = EulerDiscreteScheduler.from_config(
+    pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample"
+)
+
+# Ensure using the same inference steps as the loaded model and CFG set to 0.
+pipe("A girl smiling", num_inference_steps=1, guidance_scale=0).images[0].save(
+    "output.png"
+)
+
             ```
             """
             )
 
-        inputs = [prompt, steps, seed]
+        inputs = [prompt, seed]
         generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
         prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
-        steps.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
         seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
 
 demo.queue()

requirements.txt

@@ -1,8 +1,7 @@
-diffusers==0.23.1
+diffusers==0.26.3
 transformers
-gradio==4.7.1
---extra-index-url https://download.pytorch.org/whl/cu121
-torch==2.1.0
+gradio==4.19.1
+torch==2.2.0
 fastapi==0.104.0
 uvicorn==0.23.2
 Pillow==10.1.0
@@ -11,4 +10,6 @@ compel==2.0.2
 controlnet-aux==0.0.7
 peft==0.6.0
 xformers
-hf_transfer
+hf_transfer
+huggingface_hub
+safetensors