-
Notifications
You must be signed in to change notification settings - Fork 63
/
a1111.py
118 lines (103 loc) · 4.11 KB
/
a1111.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
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
__author__ = "receyuki"
__filename__ = "a1111.py"
__copyright__ = "Copyright 2023"
__email__ = "[email protected]"
import re
from ..format.base_format import BaseFormat
from ..utility import add_quotes, concat_strings
class A1111(BaseFormat):
PROMPT_MAPPING = {
# "Model": ("sd_model", True),
# "prompt",
# "negative_prompt",
"Seed": ("seed", False),
"Variation seed strength": ("subseed_strength", False),
# "seed_resize_from_h",
# "seed_resize_from_w",
"Sampler": ("sampler_name", True),
"Steps": ("steps", False),
"CFG scale": ("cfg_scale", False),
# "width",
# "height",
"Face restoration": ("restore_faces", False),
}
SETTING_KEY = ["Model", "Sampler", "Seed", "CFG scale", "Steps", "Size"]
def __init__(self, info: dict = None, raw: str = ""):
super().__init__(info, raw)
self._extra = ""
def _process(self):
if not self._raw:
self._raw = self._info.get("parameters", "")
self._extra = self._info.get("postprocessing", "")
self._sd_format()
def _sd_format(self):
if not self._raw and not self._extra:
return
steps_index = self._raw.find("\nSteps:")
# w/ setting
if steps_index != -1:
self._positive = self._raw[:steps_index].strip()
self._setting = self._raw[steps_index:].strip()
# w/ neg
if "Negative prompt:" in self._raw:
prompt_index = self._raw.find("\nNegative prompt:")
# w/ neg and w/ setting
if steps_index != -1:
self._negative = self._raw[
prompt_index + len("Negative prompt:") + 1 : steps_index
].strip()
# w/ neg and w/o setting
else:
self._negative = self._raw[
prompt_index + len("Negative prompt:") + 1 :
].strip()
self._positive = self._raw[:prompt_index].strip()
# w/o neg and w/o setting
elif steps_index == -1:
self._positive = self._raw
# match parameters like "Steps: x",
pattern = r"\s*([^:,]+):\s*([^,]+)"
matches = re.findall(pattern, self._setting)
setting_dict = {}
for key, value in matches:
if key not in setting_dict:
setting_dict[key] = value
[self._width, self._height] = setting_dict.get("Size", "0x0").split("x")
for p, s in zip(super().PARAMETER_KEY, A1111.SETTING_KEY):
self._parameter[p] = setting_dict.get(s)
if self._extra:
self._raw = concat_strings(self._raw, self._extra)
self._setting = concat_strings(self._setting, self._extra)
def prompt_to_line(self):
if not self._setting:
return ""
single_line_prompt = "--prompt " + add_quotes(self._positive).replace("\n", "")
if self._negative:
single_line_prompt += " --negative_prompt " + add_quotes(
self._negative
).replace("\n", "")
setting = dict(
filter(
lambda x: len(x) == 2,
(param.split(": ") for param in self._setting.split(", ")),
)
)
for key, value in setting.items():
if key == "Size":
width, height = value.split("x")
single_line_prompt += " --width " + width
single_line_prompt += " --height " + height
if key == "Seed resize from":
seed_resize_from_w, seed_resize_from_h = value.split("x")
single_line_prompt += " --seed_resize_from_w " + seed_resize_from_w
single_line_prompt += " --seed_resize_from_h " + seed_resize_from_h
try:
(tag, is_str) = A1111.PROMPT_MAPPING.get(key)
except:
pass
else:
if is_str:
single_line_prompt += " --" + tag + " " + add_quotes(str(value))
else:
single_line_prompt += " --" + tag + " " + value
return single_line_prompt