forked from CrazyBoyM/dreambooth-for-diffusion
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_model.py
More file actions
27 lines (24 loc) · 780 Bytes
/
test_model.py
File metadata and controls
27 lines (24 loc) · 780 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
26
27
from diffusers import StableDiffusionPipeline
import torch
from diffusers import DDIMScheduler
model_path = "./new_model"
prompt = "a cute girl, blue eyes, brown hair"
pipe = StableDiffusionPipeline.from_pretrained(
model_path,
torch_dtype=torch.float16,
scheduler=DDIMScheduler(
beta_start=0.00085,
beta_end=0.012,
beta_schedule="scaled_linear",
clip_sample=False,
set_alpha_to_one=True,
),
safety_checker=None
)
# def dummy(images, **kwargs):
# return images, False
# pipe.safety_checker = dummy
pipe = pipe.to("cuda")
images = pipe(prompt, num_inference_steps=30, num_images_per_prompt=3).images
for i, image in enumerate(images):
image.save(f"test-{i}.png")