-
Notifications
You must be signed in to change notification settings - Fork 63
/
test_model.py
28 lines (25 loc) · 832 Bytes
/
test_model.py
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
from diffusers import StableDiffusionPipeline
import torch
from diffusers import DDIMScheduler
model_path = "./new_model"
prompt = "a cute girl, blue eyes, brown hair"
torch.manual_seed(123123123)
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, width=512, height=512, num_inference_steps=30, num_images_per_prompt=3).images
for i, image in enumerate(images):
image.save(f"test-{i}.png")