Update space

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import StableDiffusionXLPipeline, LCMScheduler, UNet2DConditionModel, EulerDiscreteScheduler
+from diffusers import DiffusionPipeline
 from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 import torch
@@ -53,25 +53,16 @@ def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str
     p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
     return p.replace("{prompt}", positive), n + negative
 
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "ByteDance/SDXL-Lightning"
-ckpt = "sdxl_lightning_4step_unet.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.
-pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16")
+pipe.to("cuda")
 
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt, negative_prompt, width, height, style_name=None):
+def infer(prompt, negative_prompt, width, height, guidance_scale, style_name=None):
     seed = random.randint(0,4294967295)
-    guidance_scale = 0
+    # guidance_scale = 7.5
 
     generator = torch.Generator().manual_seed(seed)
 
@@ -176,7 +167,15 @@ with gr.Blocks(css=css) as demo:
                     step=32,
                     value=1024,
                 )
+            with gr.Row():
 
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=50.0,
+                    step=0.1,
+                    value=10,
+                )
         gr.Examples(
             examples = examples,
             inputs = [prompt]
@@ -194,7 +193,7 @@ Used Stable Diffusion XL (SDXL) Model by <a href="https://huggingface.co/stabili
 
     run_button.click(
         fn = infer,
-        inputs = [prompt, negative_prompt, width, height, style_selection],
+        inputs = [prompt, negative_prompt, width, height, guidance_scale, style_selection],
         outputs = [result]
     )