Spaces:
Running
on
Zero
Running
on
Zero
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() | |