{{ message }}
Repository navigation
Expand file tree
/
Copy path__init__.py
More file actions
202 lines (171 loc) · 8.68 KB
/
Copy path__init__.py
File metadata and controls
202 lines (171 loc) · 8.68 KB
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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
success = True
import sys
from pathlib import Path
if sys.version_info < (3, 9):
success = False
print('\033[34mComfyScript: \033[91mPython 3.9+ is required.\033[0m')
try:
print('\033[34mComfyScript: \033[93mLoading nodes...\033[0m')
# If there are conflicts, the later one will override the former one.
from .nodes import ComfyUI_Ib_CustomNodes
NODE_CLASS_MAPPINGS.update(ComfyUI_Ib_CustomNodes.NODE_CLASS_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(ComfyUI_Ib_CustomNodes.NODE_DISPLAY_NAME_MAPPINGS)
from .nodes import comfyui_tooling_nodes
NODE_CLASS_MAPPINGS.update(comfyui_tooling_nodes.NODE_CLASS_MAPPINGS)
NODE_DISPLAY_NAME_MAPPINGS.update(comfyui_tooling_nodes.NODE_DISPLAY_NAME_MAPPINGS)
except (ImportError, AttributeError) as e:
success = False
print(
f'''\033[34mComfyScript: \033[91mFailed to load nodes due to missing submodules: {e}.
If you need them, try to run:
git -C "{Path(__file__).resolve().parent}" submodule update --init --recursive
\033[0m''')
def setup_script():
from .script import transpile
import inspect
import traceback
import json
import PIL
if sys.version_info < (3, 6):
print('ComfyScript: Python 3.6+ is required to preserve insertion order of input types.')
# pnginfo_add_text_original = PIL.PngImagePlugin.PngInfo.add_text
# def pnginfo_add_text_hook(self, key, value, zip=False):
# try:
# if key == 'workflow':
# workflow = value
# end_nodes = None
# frame = inspect.currentframe()
# while frame := frame.f_back:
# if 'unique_id' in frame.f_locals:
# end_nodes = [int(frame.f_locals['unique_id'])]
# break
# else:
# print('ComfyScript: Failed to resolve the id of current node.')
# comfy_script = transpile.WorkflowToScriptTranspiler(workflow).to_script(end_nodes)
# print('ComfyScript:', comfy_script, sep='\n')
# chunks = self.chunks
# self.chunks = []
# pnginfo_add_text_original(self, 'ComfyScript', comfy_script, zip)
# self.chunks.extend(chunks)
# elif key == 'ComfyScriptSource':
# value = json.loads(value)
# # print(value)
# chunks = self.chunks
# self.chunks = []
# r = pnginfo_add_text_original(self, key, value, zip)
# self.chunks.extend(chunks)
# return r
# except Exception:
# # Print stack trace, but do not block the original saving
# traceback.print_exc()
# return pnginfo_add_text_original(self, key, value, zip)
# PIL.PngImagePlugin.PngInfo.add_text = pnginfo_add_text_hook
class HookedPngInfo(PIL.PngImagePlugin.PngInfo):
def __new__(cls):
return object.__new__(cls)
def __init__(self):
self._chunks = []
self._texts = {}
def add(self, cid, data, after_idat=False):
chunk = [cid, data]
if after_idat:
chunk.append(True)
self._chunks.append(tuple(chunk))
def add_text(self, key, value, zip=False):
self._texts[key] = value, zip
@property
def chunks(self):
if self._texts:
chunks = self._chunks
self._chunks = []
if 'workflow' in self._texts or 'prompt' in self._texts:
try:
workflow, zip = self._texts['workflow' if 'workflow' in self._texts else 'prompt']
end_nodes = None
frame = inspect.currentframe()
while frame := frame.f_back:
if 'unique_id' in frame.f_locals:
end_nodes = [frame.f_locals['unique_id']]
break
else:
print('ComfyScript: Failed to resolve the id of current node.')
comfy_script = transpile.WorkflowToScriptTranspiler(workflow).to_script(end_nodes)
# TODO: Syntax highlight?
print('ComfyScript:', comfy_script, sep='\n')
super().add_text('ComfyScript', comfy_script, zip)
except Exception:
# Print stack trace, but do not block the original saving
traceback.print_exc()
else:
print("ComfyScript: Failed to save ComfyScript because neither of workflow and prompt is in extra_pnginfo")
if 'ComfyScriptSource' in self._texts:
try:
source, zip = self._texts['ComfyScriptSource']
source = json.loads(source)
# print(source)
super().add_text('ComfyScriptSource', source, zip)
del self._texts['ComfyScriptSource']
except Exception:
# Print stack trace, but do not block the original saving
traceback.print_exc()
for key, (value, zip) in self._texts.items():
super().add_text(key, value, zip)
self._chunks.extend(chunks)
self._texts.clear()
return self._chunks
# PIL.PngImagePlugin.PngInfo = HookedPngInfo
PIL.PngImagePlugin.PngInfo.__new__ = lambda _cls: HookedPngInfo.__new__(HookedPngInfo)
# import nodes
# SaveImage = nodes.NODE_CLASS_MAPPINGS['SaveImage']
# save_images_orginal = getattr(SaveImage, SaveImage.FUNCTION)
# def save_images_hook(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
# # print(extra_pnginfo)
# if extra_pnginfo is None or ('workflow' not in extra_pnginfo and 'ComfyScriptSource' not in extra_pnginfo):
# print("ComfyScript: Failed to save ComfyScript because workflow is not in extra_pnginfo")
# # elif 'ComfyScriptSource' in extra_pnginfo:
# # # Values in extra_pnginfo will be serialized as JSON
# # pnginfo_init_original = PIL.PngImagePlugin.PngInfo.__init__
# # comfy_script_source = extra_pnginfo['ComfyScriptSource']
# # def pnginfo_init_hook(self):
# # pnginfo_init_original(self)
# # self.add_text('ComfyScriptSource', comfy_script_source)
# # PIL.PngImagePlugin.PngInfo.__init__ = pnginfo_init_hook
# # del extra_pnginfo['ComfyScriptSource']
# # r = save_images_orginal(self, images, filename_prefix, prompt, extra_pnginfo)
# # extra_pnginfo['ComfyScriptSource'] = comfy_script_source
# # PIL.PngImagePlugin.PngInfo.__init__ = pnginfo_init_original
# # return r
# # elif 'workflow' in extra_pnginfo:
# # workflow = extra_pnginfo['workflow']
# # try:
# # comfy_script = script.WorkflowToScriptTranspiler(workflow).to_script()
# # # print(comfy_script)
# # # Values in extra_pnginfo will be serialized as JSON
# # # extra_pnginfo['ComfyScript'] = comfy_script
# # pnginfo_init_original = PIL.PngImagePlugin.PngInfo.__init__
# # def pnginfo_init_hook(self):
# # pnginfo_init_original(self)
# # self.add_text('ComfyScript', comfy_script)
# # PIL.PngImagePlugin.PngInfo.__init__ = pnginfo_init_hook
# # r = save_images_orginal(self, images, filename_prefix, prompt, extra_pnginfo)
# # PIL.PngImagePlugin.PngInfo.__init__ = pnginfo_init_original
# # return r
# # except Exception:
# # # Print stack trace, but do not block the original saving
# # traceback.print_exc()
# return save_images_orginal(self, images, filename_prefix, prompt, extra_pnginfo)
# setattr(SaveImage, SaveImage.FUNCTION, save_images_hook)
try:
setup_script()
except ImportError as e:
success = False
print(
f'''\033[34mComfyScript: \033[91mFailed to setup script translation due to missing dependencies: {e}.
If you need this feature, try to run:
python -m pip install -r "{Path(__file__).resolve().parent / 'script' / 'transpile' / 'requirements.txt'}"
\033[0m''')
if success:
print('\033[34mComfyScript: \033[92mLoaded\033[0m')
You can’t perform that action at this time.