-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy path__init__.py
133 lines (105 loc) · 4.3 KB
/
__init__.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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
import random
import os
def process_txt_file(txtfile: str):
script_dir = os.path.dirname(os.path.abspath(__file__)) # Script directory
file_path = os.path.join(script_dir, "./txtfiles/", txtfile)
if os.path.exists(file_path):
with open(file_path, 'r', encoding='utf-8') as file:
lines = file.readlines()
processed_lines = [line.split("//")[0].strip() for line in lines if line.strip() and not line.lstrip().startswith("//")]
if processed_lines:
return processed_lines
else:
file_path = os.path.join(script_dir, "./txtfiles/", "example_" + txtfile)
else:
file_path = os.path.join(script_dir, "./txtfiles/", "example_" + txtfile)
with open(file_path, 'r', encoding='utf-8') as file:
lines = file.readlines()
processed_lines = [line.split("//")[0].strip() for line in lines if line.strip() and not line.lstrip().startswith("//")]
return processed_lines
humans = process_txt_file("humans.txt")
others = process_txt_file("others.txt")
poses = process_txt_file("poses.txt")
styles = process_txt_file("styles.txt")
test_prompts = process_txt_file("test.txt")
subject = ["human", "other", "dual_subject", "None"]
def generate_prompt(subject: str, pose: bool, style: bool, lora_trigger_or_prefix: str, refresh: bool, test: bool, seed: int):
if refresh:
global humans, others, poses, styles
humans = process_txt_file("humans.txt")
others = process_txt_file("others.txt")
poses = process_txt_file("poses.txt")
styles = process_txt_file("styles.txt")
if seed > 0:
random.seed(seed)
prompt_human = random.choice(humans)
prompt_other = random.choice(others)
prompt_pose = random.choice(poses)
prompt_style = random.choice(styles)
if subject == "human":
prompt_subject = prompt_human
if subject == "other":
prompt_subject = prompt_other
if subject == "dual_subject":
prompt_subject = prompt_human + ", " + prompt_other
if subject == "None":
prompt_subject = ""
if pose == True:
if prompt_subject:
prompt_subject = prompt_subject + ", " + prompt_pose
else:
prompt_subject = prompt_pose
if lora_trigger_or_prefix:
if lora_trigger_or_prefix.strip():
lora_trigger_or_prefix = lora_trigger_or_prefix.strip() + ", "
if style == True:
prompt = lora_trigger_or_prefix + prompt_subject + ", " + prompt_style
else:
prompt = lora_trigger_or_prefix + prompt_subject
if test:
if not test_prompts:
prompt = prompt
else:
prompt = lora_trigger_or_prefix + test_prompts.pop(0)
return prompt
class OneButtonPromptFlux:
CATEGORY = "MW-OneButtonPrompt"
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("prompt",)
FUNCTION = "fluxprompt"
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"refresh": ("BOOLEAN", {"default": False}),
},
"optional": {
"subject": (subject, {
"default": "human", "tooltip": "'dual_subject' including both. 'None' will be no subject."
}),
"pose": ("BOOLEAN", {"default": False, "tooltip": "The pose of any subject."}),
"style": ("BOOLEAN", {"default": False}),
"lora_trigger_or_prefix": ("STRING", {
"multiline": False,
"default": "", "tooltip": "Lora trigger words or custom prefix."
}),
"test": ("BOOLEAN", {"default": False}),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
},
}
def fluxprompt(
self,
subject: str = "human",
pose: bool = False,
style: bool = False,
lora_trigger_or_prefix: str = "",
refresh: bool = False,
test: bool = False,
seed: int = 0):
return (generate_prompt(subject, pose, style, lora_trigger_or_prefix, refresh, test, seed),)
NODE_CLASS_MAPPINGS = {
"OneButtonPromptFlux": OneButtonPromptFlux,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"OneButtonPromptFlux": "One Button Prompt Flux",
}