Update space

app.py

@@ -1,7 +1,9 @@
 import gradio as gr
 import numpy as np
 import random
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionXLPipeline, LCMScheduler, UNet2DConditionModel, EulerDiscreteScheduler
+from safetensors.torch import load_file
+from huggingface_hub import hf_hub_download
 import torch
 from typing import Tuple
 
@@ -66,39 +68,39 @@ 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
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+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")
 
-if torch.cuda.is_available():
-    # torch.cuda.max_memory_allocated(device=device)
-    pipe = StableDiffusionXLPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-    # pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
-# else: 
-#     pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", use_safetensors=True)
-#     pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, style_name=None):
+def infer(prompt, negative_prompt, width, height, guidance_scale, style_name=None):
+    seed = random.randint(0,4294967295)
 
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-        
     generator = torch.Generator().manual_seed(seed)
 
     prompt, negative_prompt = apply_style(style_name, prompt, negative_prompt)
-    
+
     image = pipe(
-        prompt = prompt, 
+        prompt = prompt,
         negative_prompt = negative_prompt,
-        guidance_scale = guidance_scale, 
-        num_inference_steps = num_inference_steps, 
-        width = width, 
+        guidance_scale = guidance_scale,
+        # num_inference_steps = num_inference_steps,
+        width = width,
         height = height,
         generator = generator
-    ).images[0] 
-    
+    ).images[0]
+
     return image
 
 examples = [
@@ -120,15 +122,15 @@ else:
     power_device = "CPU"
 
 with gr.Blocks(css=css) as demo:
-    
+
     with gr.Column(elem_id="col-container"):
         gr.Markdown(f"""
         # Text-to-Image Gradio Template
         Currently running on {power_device}.
         """)
-        
+
         with gr.Row():
-            
+
             prompt = gr.Text(
                 label="Prompt",
                 show_label=False,
@@ -136,74 +138,55 @@ with gr.Blocks(css=css) as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            
+
             run_button = gr.Button("Run", scale=0)
-        
+
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
             negative_prompt = gr.Textbox(
                 label="Negative prompt",
-                show_label=False,
+                # show_label=False,
                 max_lines=1,
                 placeholder="Enter a negative prompt",
-                elem_id="negative-prompt-text-input",
-                value="low-quality, text, blurry, fuzziness"
+                elem_id="negative-prompt-text-input"
             )
-            
+
             style_selection = gr.Radio(
                                show_label=True, container=True, interactive=True,
                                choices=STYLE_NAMES,
                                value=DEFAULT_STYLE_NAME,
                                label='Image Style'
              )
-            
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
-            )
-            
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-            
+
             with gr.Row():
-                
+
                 width = gr.Slider(
                     label="Width",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
-                
+
                 height = gr.Slider(
                     label="Height",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=512,
+                    value=1024,
                 )
-            
+
             with gr.Row():
-                
+
                 guidance_scale = gr.Slider(
                     label="Guidance scale",
                     minimum=0.0,
-                    maximum=10.0,
+                    maximum=50.0,
                     step=0.1,
-                    value=0.0,
+                    value=5,
                 )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=2,
-                )
-        
+
         gr.Examples(
             examples = examples,
             inputs = [prompt]
@@ -211,7 +194,7 @@ with gr.Blocks(css=css) as demo:
 
     run_button.click(
         fn = infer,
-        inputs = [prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, style_selection],
+        inputs = [prompt, negative_prompt, width, height, guidance_scale, style_selection],
         outputs = [result]
     )