|
1 | | -import sys |
2 | | -from os.path import dirname |
3 | | -from copy import deepcopy |
4 | | -from random import randint |
5 | | -LPP_ROOT_DIR = dirname(__file__) |
6 | | -sys.path.append(LPP_ROOT_DIR) |
7 | | - |
8 | | -from .lpp.sources import Derpibooru, E621 |
9 | | -from .lpp.backend import SourcesManager, PromptsManager, CacheManager |
10 | | - |
11 | | -sm = SourcesManager(LPP_ROOT_DIR) |
12 | | -cm = CacheManager(LPP_ROOT_DIR) |
13 | | - |
14 | | - |
15 | | -class ComfyTagSourceBase: |
16 | | - tag_source_input_types_base = { |
17 | | - "required": { |
18 | | - "query": ("STRING", { |
19 | | - "multiline": True |
20 | | - }), |
21 | | - "tag_filter": ("STRING", { |
22 | | - "multiline": False |
23 | | - }), |
24 | | - "count": ("INT", { |
25 | | - "default": 100, |
26 | | - "min": 5, |
27 | | - "max": 300, |
28 | | - "step": 5, |
29 | | - "display": "slider" |
30 | | - }), |
31 | | - "send_request": ("BOOLEAN", { |
32 | | - "default": False |
33 | | - }) |
34 | | - }, |
35 | | - "optional": { |
36 | | - "update_dummy": ("INT", { |
37 | | - "default": 0, |
38 | | - "min": 0, |
39 | | - "max": 0xffffffffffffffff, |
40 | | - }), |
41 | | - "prompt_template": ("STRING", { |
42 | | - "multiline": False |
43 | | - }) |
44 | | - } |
45 | | - } |
46 | | - |
47 | | - def __init__(self, source): |
48 | | - self._sm: SourcesManager = SourcesManager(LPP_ROOT_DIR, [source]) |
49 | | - self._pm: PromptsManager = PromptsManager(self._sm) |
50 | | - |
51 | | - RETURN_TYPES = ("STRING", "LPP_TAG_DATA") |
52 | | - RETURN_NAMES = ("Prompt", "LPP Tag Data") |
53 | | - CATEGORY = "LPP/sources" |
54 | | - FUNCTION = "get_prompt" |
55 | | - |
56 | | - def get_prompt(self): |
57 | | - pass |
58 | | - |
59 | | - |
60 | | -class ComfyDerpibooru(ComfyTagSourceBase): |
61 | | - def __init__(self): |
62 | | - ComfyTagSourceBase.__init__(self, Derpibooru) |
63 | | - |
64 | | - @classmethod |
65 | | - def INPUT_TYPES(self): |
66 | | - types = deepcopy(self.tag_source_input_types_base) |
67 | | - types["required"]["filter"] = (sm.sources["Derpibooru"].get_filters(),) |
68 | | - types["required"]["sort_by"] = (sm.sources["Derpibooru"].get_sort_options(),) |
69 | | - types["required"]["format"] = (sm.sources["Derpibooru"].get_model_names(),) |
70 | | - types["optional"]["tag_data"] = ("LPP_TAG_DATA_DERPIBOORU",) |
71 | | - return types |
72 | | - |
73 | | - def get_prompt( |
74 | | - self, query, count, filter, sort_by, format, tag_filter, |
75 | | - send_request, tag_data=None, update_dummy=0, prompt_template="" |
76 | | - ): |
77 | | - if tag_data: |
78 | | - self._pm.tag_data = tag_data |
79 | | - elif self._pm.get_loaded_prompts_count() == 0 or send_request: |
80 | | - self._pm.tag_data = self._sm.request_prompts( |
81 | | - "Derpibooru", query, count, filter, sort_by |
82 | | - ) |
83 | | - return ( |
84 | | - self._pm.choose_prompts(format, prompt_template, 1, tag_filter)[0], |
85 | | - (self._pm.tag_data, tag_filter) |
86 | | - ) |
87 | | - |
88 | | - |
89 | | -class ComfyE621(ComfyTagSourceBase): |
90 | | - def __init__(self): |
91 | | - ComfyTagSourceBase.__init__(self, E621) |
92 | | - |
93 | | - @classmethod |
94 | | - def INPUT_TYPES(self): |
95 | | - types = deepcopy(self.tag_source_input_types_base) |
96 | | - types["required"]["rating"] = (sm.sources["E621"].get_ratings(),) |
97 | | - types["required"]["sort_by"] = (sm.sources["E621"].get_sort_options(),) |
98 | | - types["required"]["format"] = (sm.sources["E621"].get_model_names(),) |
99 | | - types["optional"]["tag_data"] = ("LPP_TAG_DATA_E621",) |
100 | | - return types |
101 | | - |
102 | | - def get_prompt( |
103 | | - self, query, count, rating, sort_by, format, tag_filter, |
104 | | - send_request, tag_data=None, update_dummy=0, prompt_template="" |
105 | | - ): |
106 | | - if tag_data: |
107 | | - self._pm.tag_data = tag_data |
108 | | - elif self._pm.get_loaded_prompts_count() == 0 or send_request: |
109 | | - self._pm.tag_data = self._sm.request_prompts( |
110 | | - "E621", query, count, rating, sort_by |
111 | | - ) |
112 | | - return ( |
113 | | - self._pm.choose_prompts(format, prompt_template, 1, tag_filter)[0], |
114 | | - (self._pm.tag_data, tag_filter) |
115 | | - ) |
116 | | - |
117 | | - |
118 | | -class LPPSaver: |
119 | | - @classmethod |
120 | | - def INPUT_TYPES(self): |
121 | | - return { |
122 | | - "required": { |
123 | | - "tag_data": ("LPP_TAG_DATA",), |
124 | | - "name": ("STRING", { |
125 | | - "multiline": False, |
126 | | - }), |
127 | | - "overwrite": ("BOOLEAN", { |
128 | | - "default": False |
129 | | - }) |
130 | | - } |
131 | | - } |
132 | | - RETURN_TYPES = () |
133 | | - CATEGORY = "LPP" |
134 | | - FUNCTION = "save_tag_data" |
135 | | - OUTPUT_NODE = True |
136 | | - |
137 | | - def save_tag_data(self, tag_data, name, overwrite): |
138 | | - existing_names = cm.get_saved_names() |
139 | | - if (name in existing_names and overwrite) \ |
140 | | - or name not in existing_names: |
141 | | - cm.cache_tag_data(name, tag_data[0], tag_data[1]) |
142 | | - return {} |
143 | | - |
144 | | - |
145 | | -class ForceRunBase: |
146 | | - @classmethod |
147 | | - def IS_CHANGED(self): |
148 | | - return randint(0, 0xffffffffffffffff) |
149 | | - |
150 | | - |
151 | | -class LPPLoaderDerpibooru(ForceRunBase): |
152 | | - @classmethod |
153 | | - def INPUT_TYPES(self): |
154 | | - return { |
155 | | - "required": { |
156 | | - "collection_name": (cm.get_saved_names("Derpibooru"),) |
157 | | - } |
158 | | - } |
159 | | - RETURN_TYPES = ("LPP_TAG_DATA_DERPIBOORU",) |
160 | | - RETURN_NAMES = ("tag data",) |
161 | | - CATEGORY = "LPP/loaders" |
162 | | - FUNCTION = "load_tag_data" |
163 | | - |
164 | | - def load_tag_data(self, collection_name): |
165 | | - return (cm.get_tag_data(collection_name),) |
166 | | - |
167 | | - |
168 | | -class LPPLoaderE621(ForceRunBase): |
169 | | - @classmethod |
170 | | - def INPUT_TYPES(self): |
171 | | - return { |
172 | | - "required": { |
173 | | - "collection_name": (cm.get_saved_names("E621"),) |
174 | | - } |
175 | | - } |
176 | | - RETURN_TYPES = ("LPP_TAG_DATA_E621",) |
177 | | - RETURN_NAMES = ("tag data",) |
178 | | - CATEGORY = "LPP/loaders" |
179 | | - FUNCTION = "load_tag_data" |
180 | | - |
181 | | - def load_tag_data(self, collection_name): |
182 | | - return (cm.get_tag_data(collection_name),) |
183 | | - |
184 | | - |
185 | | -class LPPDeleter(ForceRunBase): |
186 | | - @classmethod |
187 | | - def INPUT_TYPES(self): |
188 | | - return { |
189 | | - "required": { |
190 | | - "collection_name": (cm.get_saved_names(),) |
191 | | - } |
192 | | - } |
193 | | - RETURN_TYPES = () |
194 | | - CATEGORY = "LPP" |
195 | | - FUNCTION = "delete_tag_data" |
196 | | - OUTPUT_NODE = True |
197 | | - |
198 | | - def delete_tag_data(self, collection_name): |
199 | | - cm.delete_tag_data(collection_name) |
200 | | - return {} |
| 1 | +from .lpp.ui.comfy import ComfyDerpibooru, LPPLoaderDerpibooru |
| 2 | +from .lpp.ui.comfy import ComfyE621, LPPLoaderE621 |
| 3 | +from .lpp.ui.comfy import ComfyDanbooru, LPPLoaderDanbooru |
| 4 | +from .lpp.ui.comfy import LPPSaver, LPPDeleter |
201 | 5 |
|
202 | 6 |
|
203 | 7 | NODE_CLASS_MAPPINGS = { |
204 | 8 | "LPP_Derpibooru": ComfyDerpibooru, |
205 | 9 | "LPP_E621": ComfyE621, |
| 10 | + "LPP_Danbooru": ComfyDanbooru, |
206 | 11 | "LPP_Saver": LPPSaver, |
207 | 12 | "LPP_Loader_Derpibooru": LPPLoaderDerpibooru, |
208 | 13 | "LPP_Loader_E621": LPPLoaderE621, |
| 14 | + "LPP_Loader_Danbooru": LPPLoaderDanbooru, |
209 | 15 | "LPP_Deleter": LPPDeleter |
210 | 16 | } |
211 | 17 |
|
212 | 18 | NODE_DISPLAY_NAME_MAPPINGS = { |
213 | 19 | "LPP_Derpibooru": "Derpibooru", |
214 | 20 | "LPP_E621": "E621", |
| 21 | + "LPP_Danbooru": "Danbooru", |
215 | 22 | "LPP_Saver": "Tag Data Saver", |
216 | 23 | "LPP_Loader_Derpibooru": "Tag Data Loader (Derpibooru)", |
217 | 24 | "LPP_Loader_E621": "Tag Data Loader (E621)", |
| 25 | + "LPP_Loader_Danbooru": "Tag Data Loader (Danbooru)", |
218 | 26 | "LPP_Deleter": "Tag Data Deleter" |
219 | 27 | } |
| 28 | + |
| 29 | +__all__ = [NODE_CLASS_MAPPINGS] |
0 commit comments