4

README.md

@@ -4,10 +4,10 @@ emoji: 🐢
 colorFrom: indigo
 colorTo: blue
 sdk: gradio
-sdk_version: 4.19.1
+sdk_version: 4.19.2
 app_file: app.py
 pinned: false
 license: mit
+disable_embedding: true
+header: mini
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,40 +1,79 @@
-import gradio as gr
+from diffusers import (
+    StableDiffusionXLPipeline,
+    EulerDiscreteScheduler,
+    UNet2DConditionModel,
+    AutoencoderTiny,
+)
 import torch
-from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
-from huggingface_hub import hf_hub_download
-from safetensors.torch import load_file
-import spaces
 import os
-from PIL import Image
+from huggingface_hub import hf_hub_download
+
 
+from PIL import Image
+import gradio as gr
+import time
+from safetensors.torch import load_file
+import time
+import tempfile
+from pathlib import Path
+
+# Constants
+BASE = "stabilityai/stable-diffusion-xl-base-1.0"
+REPO = "ByteDance/SDXL-Lightning"
+# 1-step
+CHECKPOINT = "sdxl_lightning_2step_unet.safetensors"
+taesd_model = "madebyollin/taesdxl"
+
+# {
+#     "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],
+# }
+
+
+SFAST_COMPILE = os.environ.get("SFAST_COMPILE", "0") == "1"
 SAFETY_CHECKER = os.environ.get("SAFETY_CHECKER", "0") == "1"
+USE_TAESD = os.environ.get("USE_TAESD", "0") == "1"
 
-base = "stabilityai/stable-diffusion-xl-base-1.0"
-repo = "ByteDance/SDXL-Lightning"
-checkpoints = {
-    "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],
-}
+# check if MPS is available OSX only M1/M2/M3 chips
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch_device = device
+torch_dtype = torch.float16
 
-print('--------------------------')
-print(torch.cuda.is_available())
-print('--------------------------')
-# Ensure model and scheduler are initialized in GPU-enabled function
+print(f"SAFETY_CHECKER: {SAFETY_CHECKER}")
+print(f"SFAST_COMPILE: {SFAST_COMPILE}")
+print(f"USE_TAESD: {USE_TAESD}")
+print(f"device: {device}")
 
-if torch.cuda.is_available():
-    pipe = StableDiffusionXLPipeline.from_pretrained(base, torch_dtype=torch.float16, variant="fp16").to("cuda")
-    
 
+unet = UNet2DConditionModel.from_config(BASE, subfolder="unet").to(
+    "cuda", torch.float16
+)
+unet.load_state_dict(load_file(hf_hub_download(REPO, CHECKPOINT), device="cuda"))
+pipe = StableDiffusionXLPipeline.from_pretrained(
+    BASE, unet=unet, torch_dtype=torch.float16, variant="fp16", safety_checker=False
+).to("cuda")
 
+if USE_TAESD:
+    pipe.vae = AutoencoderTiny.from_pretrained(
+        taesd_model, torch_dtype=torch_dtype, use_safetensors=True
+    ).to(device)
+
+
+# Ensure sampler uses "trailing" timesteps.
+pipe.scheduler = EulerDiscreteScheduler.from_config(
+    pipe.scheduler.config, timestep_spacing="trailing"
+)
+pipe.set_progress_bar_config(disable=True)
 if SAFETY_CHECKER:
     from safety_checker import StableDiffusionSafetyChecker
     from transformers import CLIPFeatureExtractor
 
     safety_checker = StableDiffusionSafetyChecker.from_pretrained(
         "CompVis/stable-diffusion-safety-checker"
-    ).to("cuda")
+    ).to(device)
     feature_extractor = CLIPFeatureExtractor.from_pretrained(
         "openai/clip-vit-base-patch32"
     )
@@ -42,67 +81,125 @@ if SAFETY_CHECKER:
     def check_nsfw_images(
         images: list[Image.Image],
     ) -> tuple[list[Image.Image], list[bool]]:
-        safety_checker_input = feature_extractor(images, return_tensors="pt").to("cuda")
+        safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)
         has_nsfw_concepts = safety_checker(
             images=[images],
-            clip_input=safety_checker_input.pixel_values.to("cuda")
+            clip_input=safety_checker_input.pixel_values.to(torch_device),
         )
 
         return images, has_nsfw_concepts
 
-# Function
-@spaces.GPU(enable_queue=True)
-def generate_image(prompt, ckpt):
-
-    checkpoint = checkpoints[ckpt][0]
-    num_inference_steps = checkpoints[ckpt][1]
-
-    if num_inference_steps==1:
-        # Ensure sampler uses "trailing" timesteps and "sample" prediction type for 1-step inference.
-        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing", prediction_type="sample")
-    else:
-        # Ensure sampler uses "trailing" timesteps.
-        pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
-
-    pipe.unet.load_state_dict(load_file(hf_hub_download(repo, checkpoint), device="cuda"))
-    results = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=0)
 
+if SFAST_COMPILE:
+    from sfast.compilers.diffusion_pipeline_compiler import compile, CompilationConfig
+
+    # sfast compilation
+    config = CompilationConfig.Default()
+    try:
+        import xformers
+
+        config.enable_xformers = True
+    except ImportError:
+        print("xformers not installed, skip")
+    try:
+        import triton
+
+        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):
+    generator = torch.manual_seed(seed)
+    last_time = time.time()
+    results = pipe(
+        prompt=prompt,
+        generator=generator,
+        num_inference_steps=2,
+        guidance_scale=0.0,
+        # width=768,
+        # height=768,
+        output_type="pil",
+    )
+    print(f"Pipe took {time.time() - last_time} seconds")
     if SAFETY_CHECKER:
         images, has_nsfw_concepts = check_nsfw_images(results.images)
         if any(has_nsfw_concepts):
             gr.Warning("NSFW content detected.")
             return Image.new("RGB", (512, 512))
-        return images[0]
-    return results.images[0]
-
-
-
+    image = results.images[0]
+    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmpfile:
+        image.save(tmpfile, "JPEG", quality=80, optimize=True, progressive=True)
+        return Path(tmpfile.name)
 
 
-
-# Gradio Interface
-description = """
-This demo utilizes the SDXL-Lightning model by ByteDance, which is a lightning-fast text-to-image generative model capable of producing high-quality images in 4 steps.
-As a community effort, this demo was put together by AngryPenguin. Link to model: https://huggingface.co/ByteDance/SDXL-Lightning
+css = """
+#container{
+    margin: 0 auto;
+    max-width: 40rem;
+}
+#intro{
+    max-width: 100%;
+    margin: 0 auto;
+}
 """
-
-with gr.Blocks(css="style.css") as demo:
-    gr.HTML("<h1><center>Text-to-Image with SDXL-Lightning ⚡</center></h1>")
-    gr.Markdown(description)
-    with gr.Group():
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="container"):
+        gr.Markdown(
+            """
+# SDXL-Lightning- Text To Image 2-Steps
+**Model**: https://huggingface.co/ByteDance/SDXL-Lightning
+            """,
+            elem_id="intro",
+        )
         with gr.Row():
-            prompt = gr.Textbox(label='Enter you image prompt:', scale=8)
-            ckpt = gr.Dropdown(label='Select inference steps',choices=['1-Step', '2-Step', '4-Step', '8-Step'], value='4-Step', interactive=True)
-            submit = gr.Button(scale=1, variant='primary')
-    img = gr.Image(label='SDXL-Lightning Generated Image')
-
-    prompt.submit(fn=generate_image,
-                 inputs=[prompt, ckpt],
-                 outputs=img,
-                 )
-    submit.click(fn=generate_image,
-                 inputs=[prompt, ckpt],
-                 outputs=img,
-                 )
-
-demo.queue().launch()
+            with gr.Row():
+                prompt = gr.Textbox(
+                    placeholder="Insert your prompt here:", scale=5, container=False
+                )
+                generate_bt = gr.Button("Generate", scale=1)
+
+        image = gr.Image(type="filepath")
+        with gr.Accordion("Advanced options", open=False):
+            seed = gr.Slider(
+                randomize=True, minimum=0, maximum=12013012031030, label="Seed", step=1
+            )
+        with gr.Accordion("Run with diffusers"):
+            gr.Markdown(
+                """## Running SDXL-Lightning with `diffusers`
+```py
+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_2step_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")
+# Ensure using the same inference steps as the loaded model and CFG set to 0.
+pipe("A girl smiling", num_inference_steps=2, guidance_scale=0).images[0].save("output.png")
+```
+            """
+            )
+
+        inputs = [prompt, seed]
+        outputs = [image]
+        generate_bt.click(
+            fn=predict, inputs=inputs, outputs=outputs, show_progress=False
+        )
+        prompt.input(fn=predict, inputs=inputs, outputs=outputs, show_progress=False)
+        seed.change(fn=predict, inputs=inputs, outputs=outputs, show_progress=False)
+
+demo.queue()
+demo.launch()

requirements.txt

@@ -1,5 +1,16 @@
+diffusers==0.26.3
 transformers
-diffusers
-torch
-accelerate
-gradio
+gradio==4.19.2
+torch==2.1.0
+fastapi==0.104.0
+uvicorn==0.23.2
+Pillow==10.1.0
+accelerate==0.24.0
+compel==2.0.2
+controlnet-aux==0.0.7
+peft==0.6.0
+xformers
+hf_transfer
+huggingface_hub
+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'