|
import gradio as gr |
|
from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler, LCMScheduler, AutoencoderKL |
|
import torch |
|
from huggingface_hub import hf_hub_download |
|
from safetensors.torch import load_file |
|
import spaces |
|
|
|
vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16) |
|
|
|
|
|
pipe_turbo = StableDiffusionXLPipeline.from_pretrained("stabilityai/sdxl-turbo", |
|
vae=vae, |
|
torch_dtype=torch.float16, |
|
variant="fp16" |
|
) |
|
|
|
|
|
|
|
base = "stabilityai/stable-diffusion-xl-base-1.0" |
|
repo = "ByteDance/SDXL-Lightning" |
|
ckpt = "sdxl_lightning_1step_unet_x0.safetensors" |
|
|
|
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(torch.float16) |
|
unet.load_state_dict(load_file(hf_hub_download(repo, ckpt))) |
|
pipe_lightning = StableDiffusionXLPipeline.from_pretrained(base, |
|
unet=unet, |
|
vae=vae, |
|
text_encoder_1=pipe_turbo.text_encoder_1, |
|
text_encoder_2=pipe_turbo.text_encoder_2, |
|
tokenizer=pipe_turbo.tokenizer, |
|
tokenizer_2=pipe_turbo.tokenizer_2, |
|
torch_dtype=torch.float16, |
|
variant="fp16" |
|
) |
|
del unet |
|
pipe_lightning.scheduler = EulerDiscreteScheduler.from_config(pipe_lightning.scheduler.config, timestep_spacing="trailing", prediction_type="sample") |
|
|
|
|
|
|
|
repo_name = "ByteDance/Hyper-SD" |
|
ckpt_name = "Hyper-SDXL-1step-Unet.safetensors" |
|
|
|
unet = UNet2DConditionModel.from_config(base, subfolder="unet").to(torch.float16) |
|
unet.load_state_dict(load_file(hf_hub_download(repo_name, ckpt_name))) |
|
pipe_hyper = StableDiffusionXLPipeline.from_pretrained(base, |
|
unet=unet, |
|
vae=vae, |
|
text_encoder_1=pipe_turbo.text_encoder_1, |
|
text_encoder_2=pipe_turbo.text_encoder_2, |
|
tokenizer=pipe_turbo.tokenizer, |
|
tokenizer_2=pipe_turbo.tokenizer_2, |
|
torch_dtype=torch.float16, |
|
variant="fp16" |
|
) |
|
pipe_hyper.scheduler = LCMScheduler.from_config(pipe_hyper.scheduler.config) |
|
|
|
del unet |
|
|
|
@spaces.GPU |
|
def run_comparison(prompt): |
|
image_turbo.to("cuda") |
|
image_turbo=pipe_turbo(prompt=prompt, num_inference_steps=1, guidance_scale=0).images[0] |
|
image_turbo.to("cpu") |
|
image_lightning.to("cuda") |
|
image_lightning=pipe_lightning(prompt=prompt, num_inference_steps=1, guidance_scale=0).images[0] |
|
image_lightning.to("cpu") |
|
image_hyper.to("cuda") |
|
image_hyper=pipe_hyper(prompt=prompt, num_inference_steps=1, guidance_scale=0, timesteps=[800]).images[0] |
|
image_turbo.to("cpu") |
|
return image_turbo, image_lightning, image_hyper |
|
|
|
|
|
css = ''' |
|
.gradio-container{max-width: 768px !important} |
|
''' |
|
with gr.Blocks(css=css) as demo: |
|
prompt = gr.Textbox(label="Prompt") |
|
run = gr.Button("Run") |
|
with gr.Row(): |
|
image_turbo = gr.Image(label="SDXL Turbo") |
|
image_lightning = gr.Image(label="SDXL Lightning") |
|
image_hyper = gr.Image("Hyper SDXL") |
|
|
|
run.click(fn=run_comparison, inputs=prompt, outputs=[image_turbo, image_lightning, image_hyper]) |
|
|