78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
import nodes
|
|
from comfy_api.latest import ComfyExtension, io
|
|
|
|
|
|
class AnyNode(io.ComfyNode):
|
|
ANYCLASS = True
|
|
OUTPUTS = []
|
|
INCLUDE = []
|
|
EXCLUDE = []
|
|
CATEGORY= "sampling/custom_sampler"
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
combo_options = []
|
|
for k in nodes.NODE_CLASS_MAPPINGS:
|
|
klass = nodes.NODE_CLASS_MAPPINGS[k]
|
|
if hasattr(klass, 'ANYCLASS'):
|
|
continue
|
|
if not hasattr(klass, 'define_schema'):
|
|
continue
|
|
if not all(x in k for x in cls.INCLUDE):
|
|
continue
|
|
if any(x in k for x in cls.EXCLUDE):
|
|
continue
|
|
s = klass.define_schema()
|
|
if len(s.outputs) != len(cls.OUTPUTS) or not all(o1.__class__ == o2.__class__ for o1, o2 in zip(s.outputs, cls.OUTPUTS, strict=True)):
|
|
continue
|
|
combo_options.append(io.DynamicCombo.Option(k, s.inputs))
|
|
return io.Schema(
|
|
node_id=cls.__name__,
|
|
display_name=cls.DISPLAY_NAME,
|
|
category=cls.CATEGORY,
|
|
outputs=cls.OUTPUTS,
|
|
inputs=[io.DynamicCombo.Input("selection", options=combo_options)])
|
|
|
|
|
|
@classmethod
|
|
def execute(cls, selection):
|
|
s = selection.pop("selection")
|
|
klass = nodes.NODE_CLASS_MAPPINGS[s]
|
|
|
|
return klass.execute(**selection)
|
|
|
|
class DynamicAnySampler(AnyNode):
|
|
OUTPUTS = [io.Sampler.Output("sampler")]
|
|
INCLUDE = ["Sampler"]
|
|
EXCLUDE = ["CustomSampler"]
|
|
DISPLAY_NAME = "Any Sampler (All-In-One)"
|
|
|
|
class DynamicAnyGuider(AnyNode):
|
|
OUTPUTS = [io.Guider.Output("guider")]
|
|
CATEGORY = "sampling/custom_sampling/"
|
|
DISPLAY_NAME = "Any Guider (All-In-One)"
|
|
|
|
class DynamicAnyScheduler(AnyNode):
|
|
OUTPUTS = [io.Sigmas.Output("sigmas")]
|
|
INCLUDE = ["Scheduler"]
|
|
CATEGORY = "sampling/custom_sampling/"
|
|
DISPLAY_NAME = "Any Scheduler (All-In-One)"
|
|
|
|
class DynamicAnyNoise(AnyNode):
|
|
OUTPUTS = [io.Noise.Output("sigmas")]
|
|
INCLUDE = ["Noise"]
|
|
CATEGORY = "sampling/custom_sampling/"
|
|
DISPLAY_NAME = "Any Noise (All-In-One)"
|
|
|
|
class AnyNodeExtension(ComfyExtension):
|
|
async def get_node_list(self):
|
|
return [
|
|
DynamicAnySampler,
|
|
DynamicAnyGuider,
|
|
DynamicAnyScheduler,
|
|
DynamicAnyNoise,
|
|
]
|
|
|
|
async def comfy_entrypoint() -> AnyNodeExtension:
|
|
return AnyNodeExtension()
|