File size: 1,187 Bytes
3ab117a |
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 |
from diffusers import StableDiffusionPipeline, DDIMScheduler
import torch
device = "cuda"
# use DDIM scheduler, you can modify it to use other scheduler
scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=True)
# modify the model path
pipe = StableDiffusionPipeline.from_pretrained(
f"./output-models/1500/",
scheduler=scheduler,
safety_checker=None,
torch_dtype=torch.float16,
).to(device)
# enable xformers memory attention
pipe.enable_xformers_memory_efficient_attention()
prompt = "photo of zwx dog with Texas bluebonnet"
negative_prompt = ""
num_samples = 4
guidance_scale = 7.5
num_inference_steps = 50
height = 512
width = 512
with torch.autocast("cuda"), torch.inference_mode():
images = pipe(
prompt,
height=height,
width=width,
negative_prompt=negative_prompt,
num_images_per_prompt=num_samples,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale
).images
count = 1
for image in images:
# save image to local directory
image.save(f"img-{count}.png")
count += 1
|