Skip to content

Instantly share code, notes, and snippets.

@lucataco
Created December 7, 2023 03:21
Show Gist options
  • Save lucataco/5cace44ef82c616fa7795bb6c202fb14 to your computer and use it in GitHub Desktop.
Save lucataco/5cace44ef82c616fa7795bb6c202fb14 to your computer and use it in GitHub Desktop.
Run SDXL locally
from diffusers import DiffusionPipeline
import torch
import time
# load both base & refiner
t1 = time.time()
base = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
)
base.to("cuda")
refiner = DiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-refiner-1.0",
text_encoder_2=base.text_encoder_2,
vae=base.vae,
torch_dtype=torch.float16,
use_safetensors=True,
variant="fp16",
)
refiner.to("cuda")
t2 = time.time()
print("Init took - ",t2-t1,"seconds")
# Define how many steps and what % of steps to be run on each experts (80/20) here
n_steps = 40
high_noise_frac = 0.8
prompt = "A majestic lion jumping from a big stone at night"
t3 = time.time()
for _ in range(10):
# run both experts
image = base(
prompt=prompt,
num_inference_steps=n_steps,
denoising_end=high_noise_frac,
output_type="latent",
).images
image = refiner(
prompt=prompt,
num_inference_steps=n_steps,
denoising_start=high_noise_frac,
image=image,
).images[0]
t4 = time.time()
print("Inference took - ",t4-t3,"seconds")
image.save("output.png")
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment