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
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
|
import argparse
import datetime
from pathlib import Path
from torch import autocast
from diffusers import StableDiffusionPipeline
import torch
import json
from diffusers import AutoencoderKL, StableDiffusionPipeline, UNet2DConditionModel, PNDMScheduler
from transformers import CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
from slugify import slugify
from pipelines.stable_diffusion.no_check import NoCheck
model_id = "path-to-your-trained-model"
prompt = "A photo of sks dog in a bucket"
def parse_args():
parser = argparse.ArgumentParser(
description="Simple example of a training script."
)
parser.add_argument(
"--model",
type=str,
default=None,
)
parser.add_argument(
"--prompt",
type=str,
default=None,
)
parser.add_argument(
"--batch_size",
type=int,
default=1,
)
parser.add_argument(
"--batch_num",
type=int,
default=50,
)
parser.add_argument(
"--steps",
type=int,
default=80,
)
parser.add_argument(
"--scale",
type=int,
default=7.5,
)
parser.add_argument(
"--seed",
type=int,
default=None,
)
parser.add_argument(
"--output_dir",
type=str,
default="output/inference",
)
parser.add_argument(
"--config",
type=str,
default=None,
)
args = parser.parse_args()
if args.config is not None:
with open(args.config, 'rt') as f:
args = parser.parse_args(
namespace=argparse.Namespace(**json.load(f)["args"]))
return args
def main():
args = parse_args()
seed = args.seed or torch.random.seed()
generator = torch.Generator(device="cuda").manual_seed(seed)
now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
output_dir = Path(args.output_dir).joinpath(f"{now}_{seed}_{slugify(args.prompt)[:80]}")
output_dir.mkdir(parents=True, exist_ok=True)
tokenizer = CLIPTokenizer.from_pretrained(args.model + '/tokenizer', torch_dtype=torch.bfloat16)
text_encoder = CLIPTextModel.from_pretrained(args.model + '/text_encoder', torch_dtype=torch.bfloat16)
vae = AutoencoderKL.from_pretrained(args.model + '/vae', torch_dtype=torch.bfloat16)
unet = UNet2DConditionModel.from_pretrained(args.model + '/unet', torch_dtype=torch.bfloat16)
feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32", torch_dtype=torch.bfloat16)
pipeline = StableDiffusionPipeline(
text_encoder=text_encoder,
vae=vae,
unet=unet,
tokenizer=tokenizer,
scheduler=PNDMScheduler(
beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", skip_prk_steps=True
),
safety_checker=NoCheck(),
feature_extractor=feature_extractor
)
pipeline.enable_attention_slicing()
pipeline.to("cuda")
with autocast("cuda"):
for i in range(args.batch_num):
images = pipeline(
[args.prompt] * args.batch_size,
num_inference_steps=args.steps,
guidance_scale=args.scale,
generator=generator,
).images
for j, image in enumerate(images):
image.save(output_dir.joinpath(f"{i * args.batch_size + j}.jpg"))
if __name__ == "__main__":
main()
|
