Spaces:
Runtime error
Runtime error
File size: 5,190 Bytes
cac33a7 7c7890b cac33a7 7c7890b cac33a7 7c7890b cac33a7 7c7890b cac33a7 7c7890b cac33a7 7c7890b cac33a7 7c7890b cac33a7 7c7890b cac33a7 7c7890b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
from diffusers import StableDiffusionXLPipeline, EulerDiscreteScheduler
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
# 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()
device = torch.device(
"cuda" if torch.cuda.is_available() else "xpu" if xpu_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
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.set_progress_bar_config(disable=True)
def predict(prompt, seed=1231231):
generator = torch.manual_seed(seed)
last_time = time.time()
results = pipe(
prompt=prompt,
generator=generator,
num_inference_steps=1,
guidance_scale=0.0,
width=512,
height=512,
# original_inference_steps=params.lcm_steps,
output_type="pil",
)
print(f"Pipe took {time.time() - last_time} seconds")
nsfw_content_detected = (
results.nsfw_content_detected[0]
if "nsfw_content_detected" in results
else False
)
if nsfw_content_detected:
gr.Warning("NSFW content detected.")
return Image.new("RGB", (512, 512))
return results.images[0]
css = """
#container{
margin: 0 auto;
max-width: 40rem;
}
#intro{
max-width: 100%;
text-align: center;
margin: 0 auto;
}
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="container"):
gr.Markdown(
"""# SDXL Turbo - Text To Image
## 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
""",
elem_id="intro",
)
with gr.Row():
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 Turbo 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_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, seed]
generate_bt.click(fn=predict, inputs=inputs, outputs=image, show_progress=False)
prompt.input(fn=predict, inputs=inputs, outputs=image, show_progress=False)
seed.change(fn=predict, inputs=inputs, outputs=image, show_progress=False)
demo.queue()
demo.launch()
|