ComfyUI_ColorMod/__init__.py at main · city96/ComfyUI_ColorMod · GitHub

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
# only import if running as a custom node
try:
	import comfy.utils
except ImportError:
	pass
else:
	WEB_DIRECTORY = "./web"
	NODE_CLASS_MAPPINGS = {}

	# main nodes (no deps)
	from .nodes.mod import NODE_CLASS_MAPPINGS as mod_nodes
	NODE_CLASS_MAPPINGS.update(mod_nodes)

	# pypng dep
	try:
		import png
	except ImportError:
		print("ColorMod: Can't find pypng! Please install to enable 16bit image support.")
	else:
		# 10bit PNG nodes
		from .nodes.save_png import NODE_CLASS_MAPPINGS as save_png_nodes
		NODE_CLASS_MAPPINGS.update(save_png_nodes)

	# cv2 dep
	try:
		import cv2
	except ImportError:
		print("ColorMod: Can't find opencv! Please install to enable HDR/tonemapping support.")
	else:
		# HDR creation/etc nodes
		from .nodes.hdr import NODE_CLASS_MAPPINGS as hdr_nodes
		NODE_CLASS_MAPPINGS.update(hdr_nodes)

		# HDR save/load nodes
		from .nodes.save_hdr import NODE_CLASS_MAPPINGS as save_hdr_nodes
		NODE_CLASS_MAPPINGS.update(save_hdr_nodes)

		# HDR tonemap nodes
		from .nodes.tonemap import NODE_CLASS_MAPPINGS as tonemap_nodes
		NODE_CLASS_MAPPINGS.update(tonemap_nodes)

		# color space nodes
		from .nodes.colorspace import NODE_CLASS_MAPPINGS as colorspace_nodes
		NODE_CLASS_MAPPINGS.update(colorspace_nodes)

	# export
	NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()}
	__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]