-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathsd_maker.py
More file actions
25 lines (19 loc) · 778 Bytes
/
sd_maker.py
File metadata and controls
25 lines (19 loc) · 778 Bytes
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
from diffusers import DiffusionPipeline, LCMScheduler
import torch
from os import path
cache_path = path.join(path.dirname(path.abspath(__file__)), "models")
def make_img(p, model_id="Lykon/dreamshaper-7"):
lcm_lora_id = "latent-consistency/lcm-lora-sdv1-5"
if model_id == "stabilityai/stable-diffusion-xl-base-1.0":
lcm_lora_id = "latent-consistency/lcm-lora-sdxl"
pipe = DiffusionPipeline.from_pretrained(model_id, variant="fp16", cache_dir=cache_path)
pipe.load_lora_weights(lcm_lora_id)
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to(device="cuda", dtype=torch.float16)
images = pipe(
prompt=p,
num_inference_steps=6,
guidance_scale=1,
).images[0]
del pipe
return images