Refactoring: Started migration to ComfyUI API v3, bugfixes, added TT_ContextNode and TT_ContextFlux2Node, new workflow nodes
This commit is contained in:
TenserTensor
@@ -1,141 +0,0 @@
|
||||
# (c) TenserTensor || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
import dataclasses
|
||||
import re
|
||||
from typing import override
|
||||
|
||||
import torch
|
||||
|
||||
from comfy.clip_vision import Output
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from .lib.common import Context, TTContext
|
||||
|
||||
CATEGORY = "TenserTensor/Context"
|
||||
|
||||
|
||||
class ContextNode(IO.ComfyNode):
|
||||
"""
|
||||
Abstract base class for all context nodes.
|
||||
Handles context creation, delegation to handle(), and building return tuple.
|
||||
Subclasses should set INPUT_TYPES_DICT, RETURN_TYPES, RETURN_NAMES from a Schema,
|
||||
and override handle() to implement node-specific logic.
|
||||
"""
|
||||
REQUIRED = ()
|
||||
OPTIONAL = ()
|
||||
RETURNS = ()
|
||||
|
||||
F_TYPES = {}
|
||||
|
||||
@classmethod
|
||||
def fill_field_types(cls):
|
||||
cls.F_TYPES = {f.name: f.type.__args__[0] for f in dataclasses.fields(TTContext)}
|
||||
|
||||
@classmethod
|
||||
def inputs(cls) -> list[Input]:
|
||||
def get_input(input_name, optional):
|
||||
return Context.Input("context", optional=optional) if input_name == "context" else cls.F_TYPES[input_name].Input(input_name,
|
||||
optional=optional)
|
||||
|
||||
len(cls.F_TYPES) == 0 and cls.fill_field_types()
|
||||
input_schema = []
|
||||
|
||||
for f_name in cls.REQUIRED:
|
||||
input_schema.append(get_input(f_name, optional=False))
|
||||
|
||||
for f_name in cls.OPTIONAL:
|
||||
input_schema.append(get_input(f_name, optional=True))
|
||||
|
||||
return input_schema
|
||||
|
||||
@classmethod
|
||||
def outputs(cls) -> list[Output]:
|
||||
len(cls.F_TYPES) == 0 and cls.fill_field_types()
|
||||
output_schema = []
|
||||
|
||||
for f_name in cls.RETURNS if len(cls.RETURNS) > 0 else cls.OPTIONAL:
|
||||
output_schema.append(
|
||||
Context.Output(display_name="CONTEXT") if f_name == "context" else cls.F_TYPES[f_name].Output(display_name=f_name.upper()))
|
||||
|
||||
return output_schema
|
||||
|
||||
@classmethod
|
||||
def node_output(cls, context, **kwargs) -> IO.NodeOutput:
|
||||
args = {}
|
||||
for f_name in cls.RETURNS if len(cls.RETURNS) > 0 else cls.OPTIONAL:
|
||||
if f_name != "context":
|
||||
args[f_name.upper()] = kwargs.get(f_name)
|
||||
|
||||
return IO.NodeOutput(**{{"context": context}, args})
|
||||
|
||||
@classmethod
|
||||
def display_name(cls):
|
||||
return re.sub(r'(?<=[a-z])(?=[A-Z])', ' ', cls.__name__.replace("TT_", "TT ").replace("Node", ""))
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id=cls.__name__,
|
||||
display_name=cls.display_name(),
|
||||
category=CATEGORY,
|
||||
inputs=cls.inputs(),
|
||||
outputs=cls.outputs(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def handle(cls, context, **kwargs):
|
||||
return (context, kwargs,)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, context=None, **kwargs) -> IO.NodeOutput:
|
||||
context = TTContext.create(**kwargs) if not context else context
|
||||
context, kwargs = cls.handle(context, **kwargs)
|
||||
|
||||
return cls.node_output(context, **kwargs)
|
||||
|
||||
|
||||
class TT_BaseContextNode(ContextNode):
|
||||
"""Base context node. Accepts model, clip, vae, latent and optional workflow_config."""
|
||||
REQUIRED = ("model", "clip", "vae", "latent")
|
||||
OPTIONAL = ("workflow_config",)
|
||||
RETURNS = ("context",)
|
||||
|
||||
|
||||
class TT_BaseContextFlux2Node(ContextNode):
|
||||
"""
|
||||
Base context node for FLUX2 workflows. Accepts model, clip, vae, latent, sampler, noise, and sigmas.
|
||||
Optionally accepts workflow_config. Returns context.
|
||||
"""
|
||||
REQUIRED = ("model", "clip", "vae", "latent", "sampler", "rnd_noise", "sigmas",)
|
||||
OPTIONAL = ("workflow_config",)
|
||||
RETURNS = ("context",)
|
||||
|
||||
|
||||
class TT_BaseContextPassthroughNode(ContextNode):
|
||||
"""Passes context through without modification."""
|
||||
REQUIRED = ("context",)
|
||||
OPTIONAL = ()
|
||||
RETURNS = ("context",)
|
||||
|
||||
|
||||
class ContextExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
TT_BaseContextNode,
|
||||
TT_BaseContextFlux2Node,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> ContextExtension:
|
||||
return ContextExtension()
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Re-exports for backward compatibility
|
||||
# ==============================================================================
|
||||
|
||||
__all__ = [
|
||||
"TT_BaseContextNode",
|
||||
"TT_BaseContextFlux2Node",
|
||||
"TT_BaseContextPassthroughNode",
|
||||
]
|
||||
-115
@@ -1,12 +1,7 @@
|
||||
# (c) TenserTensor || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
|
||||
import dataclasses
|
||||
from dataclasses import dataclass
|
||||
|
||||
import comfy.samplers as S
|
||||
from comfy_api.latest import IO
|
||||
from .workflow_classes import TTWorkflowSettings
|
||||
|
||||
|
||||
class BasicGuider(S.CFGGuider):
|
||||
@@ -19,113 +14,3 @@ class BasicGuider(S.CFGGuider):
|
||||
|
||||
class CommonTypes():
|
||||
MEGAPIXELS = ["0.25 MP", "0.5 MP", "1 MP", "2 MP", "4 MP", "8 MP"]
|
||||
|
||||
|
||||
@IO.comfytype(io_type="TT_WORKFLOW_CONFIG")
|
||||
class WorkflowSettings:
|
||||
Type = TTWorkflowSettings
|
||||
|
||||
class Input(IO.Input):
|
||||
def __init__(self, id: str, **kwargs):
|
||||
super().__init__(id, **kwargs)
|
||||
|
||||
class Output(IO.Output):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TTContext():
|
||||
"""
|
||||
Shared context object passed between TenserTensor context nodes.
|
||||
|
||||
Holds all pipeline state as optional fields.
|
||||
All fields default to None and are populated incrementally as the workflow progresses.
|
||||
|
||||
Use create() to instantiate from kwargs, and set_attr/get_attr for safe field access.
|
||||
Workflow config fields are accessible via set_workflow_config_attr/get_workflow_config_attr.
|
||||
"""
|
||||
model: IO.Model | None = None
|
||||
clip: IO.Clip | None = None
|
||||
vae: IO.Vae | None = None
|
||||
guider: IO.Guider | None = None
|
||||
latent: IO.Latent | None = None
|
||||
positive: IO.Conditioning | None = None
|
||||
negative: IO.Conditioning | None = None
|
||||
conditioning: IO.Conditioning | None = None
|
||||
image: IO.Image | None = None
|
||||
mask: IO.Mask | None = None
|
||||
control_net: IO.ControlNet | None = None
|
||||
workflow_config: WorkflowSettings | None = None
|
||||
sampler_name: IO.String | None = None
|
||||
sampler: IO.Sampler | None = None
|
||||
scheduler: list[str] | None = None
|
||||
rnd_noise: IO.Noise | None = None
|
||||
sigmas: IO.Sigmas | None = None
|
||||
|
||||
@classmethod
|
||||
def create(cls, **kwargs):
|
||||
ctx = cls()
|
||||
for field in dataclasses.fields(ctx):
|
||||
setattr(ctx, field.name, kwargs.get(field.name, None))
|
||||
|
||||
return ctx
|
||||
|
||||
@classmethod
|
||||
def get_attr_type(cls, key):
|
||||
return
|
||||
|
||||
def set_attr(self, key, attr):
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, attr)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Context has no attribute '{key}'")
|
||||
|
||||
def get_attr(self, key, default=None):
|
||||
if hasattr(self, key):
|
||||
val = getattr(self, key)
|
||||
return val if val is not None else default
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Context has no attribute '{key}'")
|
||||
|
||||
def update_attrs(self, **kwargs):
|
||||
for key in kwargs.keys():
|
||||
if hasattr(self, key):
|
||||
val = kwargs.get(key)
|
||||
if val is not None:
|
||||
setattr(self, key, val)
|
||||
|
||||
def set_workflow_config_attr(self, key, attr):
|
||||
if self.workflow_config:
|
||||
self.workflow_config.set_attr(key, attr)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Workflow Config is not set in context")
|
||||
|
||||
def get_workflow_config_attr(self, key, default=None):
|
||||
if self.workflow_config:
|
||||
return self.workflow_config.get_attr(key, default)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Workflow Config is not set in context")
|
||||
|
||||
def update_workflow_config_attrs(self, **kwargs):
|
||||
if self.workflow_config:
|
||||
for key in kwargs.keys():
|
||||
if hasattr(self.workflow_config, key):
|
||||
val = kwargs.get(key)
|
||||
if val is not None:
|
||||
self.workflow_config.set_attr(key, val)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Workflow Config is not set in context")
|
||||
|
||||
|
||||
@IO.comfytype(io_type="TT_CONTEXT")
|
||||
class Context:
|
||||
Type = TTContext
|
||||
|
||||
class Input(IO.Input):
|
||||
def __init__(self, id: str, **kwargs):
|
||||
super().__init__(id, **kwargs)
|
||||
|
||||
class Output(IO.Output):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
# (c) TenserTensor || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
import dataclasses
|
||||
from dataclasses import dataclass
|
||||
|
||||
from comfy.samplers import KSampler
|
||||
|
||||
|
||||
@dataclass
|
||||
class TTWorkflowSettings():
|
||||
seed: int | None = None
|
||||
steps: int | None = None
|
||||
cfg: int | None = None
|
||||
sampler_name: list[str] | None = None
|
||||
scheduler_name: list[str] | None = None
|
||||
clip_l_positive: str | None = None
|
||||
clip_g_positive: str | None = None
|
||||
t5xxl_positive: str | None = None
|
||||
clip_l_negative: str | None = None
|
||||
clip_g_negative: str | None = None
|
||||
t5xxl_negative: str | None = None
|
||||
guidance: float | None = None
|
||||
ascore_positive: float | None = None
|
||||
ascore_negative: float | None = None
|
||||
width: int | None = None
|
||||
height: int | None = None
|
||||
target_width: int | None = None
|
||||
target_height: int | None = None
|
||||
|
||||
@classmethod
|
||||
def create(cls, **kwargs):
|
||||
wf = cls()
|
||||
for field in dataclasses.fields(wf):
|
||||
setattr(wf, field.name, kwargs.get(field.name, None))
|
||||
|
||||
setattr(wf, "sampler_name", KSampler.SAMPLERS)
|
||||
setattr(wf, "scheduler_name", KSampler.SCHEDULERS)
|
||||
|
||||
return wf
|
||||
|
||||
def set_attr(self, key, attr):
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, attr)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Context has no attribute '{key}'")
|
||||
|
||||
def get_attr(self, key, default=None):
|
||||
if hasattr(self, key):
|
||||
val = getattr(self, key)
|
||||
return val if val is not None else default
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Context has no attribute '{key}'")
|
||||
+14
-3
@@ -15,7 +15,7 @@ NODES_COUNT = 0
|
||||
|
||||
try:
|
||||
before = set(dir())
|
||||
from .context_nodes import *
|
||||
from .nodes_context import *
|
||||
from .nodes.context.base_context import TT_BaseContext
|
||||
from .nodes.context.context import TT_Context
|
||||
from .nodes.context.context_passthrough import TT_ContextPassthrough
|
||||
@@ -118,6 +118,7 @@ try:
|
||||
from .nodes.workflow.sdxl_workflow_settings_advanced import TT_SdxlWorkflowSettingsAdvanced
|
||||
from .nodes.workflow.flux2_workflow_settings import TT_Flux2WorkflowSettings
|
||||
from .nodes.workflow.flux2_workflow_settings_advanced import TT_Flux2WorkflowSettingsAdvanced
|
||||
from .nodes_workflow import *
|
||||
|
||||
NODES_COUNT += len(set(dir()) - before)
|
||||
except ImportError:
|
||||
@@ -138,6 +139,8 @@ NODE_CLASS_MAPPINGS = {
|
||||
"TT_BaseContextNode": TT_BaseContextNode,
|
||||
"TT_BaseContextFlux2Node": TT_BaseContextFlux2Node,
|
||||
"TT_BaseContextPassthroughNode": TT_BaseContextPassthroughNode,
|
||||
"TT_ContextNode": TT_ContextNode,
|
||||
"TT_ContextFlux2Node": TT_ContextFlux2Node,
|
||||
# Image
|
||||
"TT_ImagePreviewSave": TT_ImagePreviewSave,
|
||||
"TT_ImagePreviewUpscaleSave": TT_ImagePreviewUpscaleSave,
|
||||
@@ -183,6 +186,9 @@ NODE_CLASS_MAPPINGS = {
|
||||
"TT_SdxlWorkflowSettingsAdvanced": TT_SdxlWorkflowSettingsAdvanced,
|
||||
"TT_Flux2WorkflowSettings": TT_Flux2WorkflowSettings,
|
||||
"TT_Flux2WorkflowSettingsAdvanced": TT_Flux2WorkflowSettingsAdvanced,
|
||||
# Workflow V3 Nodes
|
||||
"TT_Flux2WorkflowSettingsNode": TT_Flux2WorkflowSettingsNode,
|
||||
"TT_Flux2WorkflowSettingsAdvancedNode": TT_Flux2WorkflowSettingsAdvancedNode,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
@@ -200,6 +206,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TT_BaseContextNode": "TT Base Context",
|
||||
"TT_BaseContextFlux2Node": "TT FLUX2 Base Context",
|
||||
"TT_BaseContextPassthroughNode": "TT Context Passthrough",
|
||||
"TT_ContextNode": "TT Context",
|
||||
"TT_ContextFlux2Node": "TT FLUX2 Context",
|
||||
# Image
|
||||
"TT_ImagePreviewSave": "TT Image Preview / Save",
|
||||
"TT_ImagePreviewUpscaleSave": "TT Image Preview / Upscale / Save",
|
||||
@@ -243,6 +251,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TT_FluxWorkflowSettingsAdvanced": "TT FLUX Workflow Settings (Advanced)",
|
||||
"TT_SdxlWorkflowSettings": "TT SDXL Workflow Settings",
|
||||
"TT_SdxlWorkflowSettingsAdvanced": "TT SDXL Workflow Settings (Advanced)",
|
||||
"TT_Flux2WorkflowSettings": "TT FLUX2 Workflow Settings",
|
||||
"TT_Flux2WorkflowSettingsAdvanced": "TT FLUX2 Workflow Settings (Advanced)",
|
||||
"TT_Flux2WorkflowSettings": "TT FLUX2 Workflow Settings (Deprecated)",
|
||||
"TT_Flux2WorkflowSettingsAdvanced": "TT FLUX2 Workflow Settings (Deprecated/Advanced)",
|
||||
# Workflow V3 Nodes
|
||||
"TT_Flux2WorkflowSettingsNode": "TT FLUX2 Workflow Settings",
|
||||
"TT_Flux2WorkflowSettingsAdvancedNode": "TT FLUX2 Workflow Settings (Advanced)",
|
||||
}
|
||||
|
||||
@@ -46,7 +46,7 @@ class TT_Flux2WorkflowSettings:
|
||||
"GUIDANCE"
|
||||
)
|
||||
FUNCTION = "get_settings"
|
||||
CATEGORY = "TenserTensor/Workflow/FLUX2"
|
||||
CATEGORY = "TenserTensor/Deprecated/Workflow/FLUX2"
|
||||
|
||||
def get_settings(self, seed, steps, cfg, sampler_name, width, height, guidance):
|
||||
scheduler = get_schedule(steps, width, height)
|
||||
|
||||
@@ -52,7 +52,7 @@ class TT_Flux2WorkflowSettingsAdvanced:
|
||||
"GUIDANCE"
|
||||
)
|
||||
FUNCTION = "get_settings"
|
||||
CATEGORY = "TenserTensor/Workflow/FLUX2"
|
||||
CATEGORY = "TenserTensor/Deprecated/Workflow/FLUX2"
|
||||
|
||||
def get_settings(self, seed, steps, cfg, sampler_name, width, height, prompt, lora_triggers, guidance):
|
||||
scheduler = get_schedule(steps, width, height)
|
||||
|
||||
@@ -0,0 +1,293 @@
|
||||
# (c) TenserTensor || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
import dataclasses
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import override
|
||||
|
||||
import torch
|
||||
|
||||
from comfy.clip_vision import Output
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from .nodes_workflow import TTWorkflowSettings, WorkflowSettings
|
||||
|
||||
CATEGORY = "TenserTensor/Context"
|
||||
|
||||
|
||||
@dataclass
|
||||
class TTContext():
|
||||
"""
|
||||
Shared context object passed between TenserTensor context nodes.
|
||||
|
||||
Holds all pipeline state as optional fields.
|
||||
All fields default to None and are populated incrementally as the workflow progresses.
|
||||
|
||||
Use create() to instantiate from kwargs, and set_attr/get_attr for safe field access.
|
||||
Workflow config fields are accessible via set_workflow_config_attr/get_workflow_config_attr.
|
||||
"""
|
||||
cfg: IO.Float | None = None
|
||||
clip: IO.Clip | None = None
|
||||
control_net: IO.ControlNet | None = None
|
||||
conditioning: IO.Conditioning | None = None
|
||||
guider: IO.Guider | None = None
|
||||
image: IO.Image | None = None
|
||||
latent: IO.Latent | None = None
|
||||
mask: IO.Mask | None = None
|
||||
model: IO.Model | None = None
|
||||
negative: IO.Conditioning | None = None
|
||||
positive: IO.Conditioning | None = None
|
||||
rnd_noise: IO.Noise | None = None
|
||||
sampler: IO.Sampler | None = None
|
||||
seed: IO.Int | None = None
|
||||
sigmas: IO.Sigmas | None = None
|
||||
steps: IO.Int | None = None
|
||||
vae: IO.Vae | None = None
|
||||
workflow_config: WorkflowSettings | None = None
|
||||
sampler_name: IO.String | None = None
|
||||
scheduler: IO.String | None = None
|
||||
|
||||
@classmethod
|
||||
def create(cls, **kwargs):
|
||||
ctx = cls()
|
||||
for field in dataclasses.fields(ctx):
|
||||
ctx.set_attr(field.name, kwargs.get(field.name, None))
|
||||
|
||||
return ctx
|
||||
|
||||
def set_attr(self, key, attr):
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, attr)
|
||||
else:
|
||||
self.set_workflow_config_attr(key, attr)
|
||||
|
||||
def get_attr(self, key, default=None):
|
||||
if hasattr(self, key):
|
||||
val = getattr(self, key)
|
||||
return val if val is not None else default
|
||||
else:
|
||||
self.get_workflow_config_attr(key, default)
|
||||
|
||||
def update_attrs(self, **kwargs):
|
||||
for key in kwargs.keys():
|
||||
if hasattr(self, key):
|
||||
val = kwargs.get(key)
|
||||
if val is not None:
|
||||
setattr(self, key, val)
|
||||
|
||||
def set_workflow_config_attr(self, key, attr):
|
||||
if self.workflow_config:
|
||||
self.workflow_config.set_attr(key, attr)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Workflow Config is not set in context")
|
||||
|
||||
def get_workflow_config_attr(self, key, default=None):
|
||||
if self.workflow_config:
|
||||
return self.workflow_config.get_attr(key, default)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Workflow Config is not set in context")
|
||||
|
||||
def update_workflow_config_attrs(self, **kwargs):
|
||||
if self.workflow_config:
|
||||
for key in kwargs.keys():
|
||||
if hasattr(self.workflow_config, key):
|
||||
val = kwargs.get(key)
|
||||
if val is not None:
|
||||
self.workflow_config.set_attr(key, val)
|
||||
else:
|
||||
raise AttributeError(f"ERROR: Workflow Config is not set in context")
|
||||
|
||||
|
||||
@IO.comfytype(io_type="TT_CONTEXT")
|
||||
class Context:
|
||||
Type = TTContext
|
||||
|
||||
class Input(IO.Input):
|
||||
def __init__(self, id: str, **kwargs):
|
||||
super().__init__(id, **kwargs)
|
||||
|
||||
class Output(IO.Output):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
|
||||
class ContextNode(IO.ComfyNode):
|
||||
"""
|
||||
Abstract base class for all context nodes.
|
||||
Handles context creation, delegation to handle(), and building return tuple.
|
||||
Subclasses should set INPUT_TYPES_DICT, RETURN_TYPES, RETURN_NAMES from a Schema,
|
||||
and override handle() to implement node-specific logic.
|
||||
"""
|
||||
REQUIRED = ()
|
||||
OPTIONAL = ()
|
||||
RETURNS = ()
|
||||
|
||||
F_TYPES = {}
|
||||
|
||||
@classmethod
|
||||
def fill_field_types(cls):
|
||||
cls.F_TYPES = {f.name: f.type.__args__[0] for f in dataclasses.fields(TTContext)}
|
||||
|
||||
@classmethod
|
||||
def inputs(cls) -> list[Input]:
|
||||
def get_input(input_name, optional):
|
||||
match input_name:
|
||||
case "context":
|
||||
return Context.Input(input_name, optional=optional)
|
||||
case _:
|
||||
f_type = cls.F_TYPES[input_name]
|
||||
if f_type in (IO.Boolean, IO.Int, IO.Float, IO.String,):
|
||||
return f_type.Input(input_name, optional=optional, force_input=True)
|
||||
else:
|
||||
return f_type.Input(input_name, optional=optional)
|
||||
|
||||
input_schema = []
|
||||
|
||||
for f_name in cls.REQUIRED:
|
||||
input_schema.append(get_input(f_name, optional=False))
|
||||
|
||||
for f_name in cls.OPTIONAL:
|
||||
input_schema.append(get_input(f_name, optional=True))
|
||||
|
||||
return input_schema
|
||||
|
||||
@classmethod
|
||||
def outputs(cls) -> list[Output]:
|
||||
output_schema = []
|
||||
|
||||
for f_name in cls.RETURNS if len(cls.RETURNS) > 0 else cls.OPTIONAL:
|
||||
output_schema.append(
|
||||
Context.Output(display_name="CONTEXT")
|
||||
if f_name == "context"
|
||||
else cls.F_TYPES[f_name].Output(display_name=f_name.upper())
|
||||
)
|
||||
|
||||
return output_schema
|
||||
|
||||
@classmethod
|
||||
def node_output(cls, context) -> IO.NodeOutput:
|
||||
args = [context]
|
||||
for f_name in cls.RETURNS if len(cls.RETURNS) > 0 else cls.OPTIONAL:
|
||||
if f_name != "context":
|
||||
wf = context.get_attr("workflow_config")
|
||||
if hasattr(wf, f_name):
|
||||
args.append(wf.get_attr(f_name))
|
||||
else:
|
||||
args.append(context.get_attr(f_name))
|
||||
|
||||
return IO.NodeOutput(*args)
|
||||
|
||||
@classmethod
|
||||
def display_name(cls):
|
||||
return re.sub(r'(?<=[a-z])(?=[A-Z])', ' ', cls.__name__.replace("TT_", "TT ").replace("Node", ""))
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
if len(cls.F_TYPES) == 0:
|
||||
cls.fill_field_types()
|
||||
|
||||
return IO.Schema(
|
||||
node_id=cls.__name__,
|
||||
display_name=cls.display_name(),
|
||||
category=CATEGORY,
|
||||
inputs=cls.inputs(),
|
||||
outputs=cls.outputs(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def handle(cls, context, **kwargs):
|
||||
workflow_config = kwargs.get("workflow_config")
|
||||
kwargs["workflow_config"] = (
|
||||
TTWorkflowSettings.create(**kwargs)
|
||||
if not workflow_config
|
||||
else workflow_config.update_attrs(**kwargs)
|
||||
)
|
||||
context.update_attrs(**kwargs)
|
||||
|
||||
return (context, kwargs,)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, context=None, **kwargs) -> IO.NodeOutput:
|
||||
if not context:
|
||||
context = TTContext.create(**kwargs)
|
||||
cls.handle(context, **kwargs)
|
||||
else:
|
||||
context, kwargs = cls.handle(context, **kwargs)
|
||||
|
||||
return cls.node_output(context)
|
||||
|
||||
|
||||
class TT_BaseContextNode(ContextNode):
|
||||
"""Base context node. Accepts model, clip, vae, latent and optional workflow_config."""
|
||||
REQUIRED = ("model", "clip", "vae", "latent")
|
||||
OPTIONAL = ("workflow_config",)
|
||||
RETURNS = ("context",)
|
||||
|
||||
|
||||
class TT_BaseContextFlux2Node(ContextNode):
|
||||
"""
|
||||
Base context node for FLUX2 workflows. Accepts model, clip, vae, latent, sampler, noise, and sigmas.
|
||||
Optionally accepts workflow_config. Returns context.
|
||||
"""
|
||||
REQUIRED = ("model", "clip", "vae", "latent", "sampler", "rnd_noise", "sigmas",)
|
||||
OPTIONAL = ("workflow_config",)
|
||||
RETURNS = ("context",)
|
||||
|
||||
|
||||
class TT_BaseContextPassthroughNode(ContextNode):
|
||||
"""Passes context through without modification."""
|
||||
REQUIRED = ("context",)
|
||||
OPTIONAL = ()
|
||||
RETURNS = ("context",)
|
||||
|
||||
|
||||
class TT_ContextNode(ContextNode):
|
||||
"""General-purpose context node for FLUX and SDXL workflows. All fields are optional."""
|
||||
REQUIRED = ()
|
||||
OPTIONAL = (
|
||||
"context", "workflow_config", "model", "clip", "vae", "latent", "positive", "negative", "image",
|
||||
"seed", "steps", "cfg",
|
||||
)
|
||||
RETURNS = ()
|
||||
|
||||
|
||||
class TT_ContextFlux2Node(ContextNode):
|
||||
"""Context node for FLUX2 workflows. All fields are optional."""
|
||||
REQUIRED = ()
|
||||
OPTIONAL = (
|
||||
"context", "workflow_config", "model", "clip", "vae", "latent", "conditioning", "guider", "rnd_noise",
|
||||
"sampler", "sigmas", "image", "mask", "seed", "steps", "cfg",
|
||||
)
|
||||
RETURNS = ()
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# V3 entrypoint — registers context nodes with ComfyUI
|
||||
# ==============================================================================
|
||||
|
||||
class ContextExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
TT_BaseContextNode,
|
||||
TT_BaseContextFlux2Node,
|
||||
TT_BaseContextPassthroughNode,
|
||||
TT_ContextNode,
|
||||
TT_ContextFlux2Node,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> ContextExtension:
|
||||
return ContextExtension()
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Re-exports for backward compatibility
|
||||
# ==============================================================================
|
||||
|
||||
__all__ = [
|
||||
"TT_BaseContextNode",
|
||||
"TT_BaseContextFlux2Node",
|
||||
"TT_BaseContextPassthroughNode",
|
||||
"TT_ContextNode",
|
||||
"TT_ContextFlux2Node",
|
||||
]
|
||||
@@ -0,0 +1,258 @@
|
||||
# (c) TenserTensor || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
||||
|
||||
import dataclasses
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import override
|
||||
|
||||
import torch
|
||||
|
||||
from comfy.samplers import KSampler, sampler_object, Sampler
|
||||
from comfy_api.latest import IO, ComfyExtension
|
||||
from nodes import MAX_RESOLUTION
|
||||
|
||||
|
||||
@dataclass
|
||||
class TTWorkflowSettings():
|
||||
ascore_negative: float | None = None
|
||||
ascore_positive: float | None = None
|
||||
cfg: int | None = None
|
||||
clip_g_negative: str | None = None
|
||||
clip_l_negative: str | None = None
|
||||
clip_g_positive: str | None = None
|
||||
clip_l_positive: str | None = None
|
||||
guidance: float | None = None
|
||||
height: int | None = None
|
||||
lora_triggers: str | None = None
|
||||
prompt: str | None = None
|
||||
sampler: Sampler | None = None
|
||||
sampler_name: str | None = None
|
||||
scheduler: str | None = None
|
||||
seed: int | None = None
|
||||
sigmas: torch.Tensor | None = None
|
||||
steps: int | None = None
|
||||
target_height: int | None = None
|
||||
target_width: int | None = None
|
||||
t5xxl_negative: str | None = None
|
||||
t5xxl_positive: str | None = None
|
||||
width: int | None = None
|
||||
|
||||
@classmethod
|
||||
def create(cls, **kwargs) -> TTWorkflowSettings:
|
||||
wf = cls()
|
||||
for field in dataclasses.fields(wf):
|
||||
setattr(wf, field.name, kwargs.get(field.name, None))
|
||||
|
||||
return wf
|
||||
|
||||
def update_attrs(self, **kwargs) -> TTWorkflowSettings:
|
||||
for key in kwargs.keys():
|
||||
if hasattr(self, key):
|
||||
val = kwargs.get(key)
|
||||
if val is not None:
|
||||
setattr(self, key, val)
|
||||
|
||||
return self
|
||||
|
||||
def set_attr(self, key, attr) -> None:
|
||||
if hasattr(self, key):
|
||||
setattr(self, key, attr)
|
||||
|
||||
def get_attr(self, key, default=None) -> None:
|
||||
if hasattr(self, key):
|
||||
val = getattr(self, key)
|
||||
return val if val is not None else default
|
||||
|
||||
|
||||
@IO.comfytype(io_type="TT_WORKFLOW_CONFIG")
|
||||
class WorkflowSettings:
|
||||
Type = TTWorkflowSettings
|
||||
|
||||
class Input(IO.Input):
|
||||
def __init__(self, id: str, **kwargs):
|
||||
super().__init__(id, **kwargs)
|
||||
|
||||
class Output(IO.Output):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
|
||||
def compute_empirical_mu(image_seq_len, num_steps):
|
||||
slope_high_res, intercept_high_res = 8.73809524e-05, 1.89833333
|
||||
slope_low_res, intercept_low_res = 0.00016927, 0.45666666
|
||||
HIGH_RES_THRESHOLD = 4300
|
||||
MIN_STEPS, MAX_STEPS = 10, 200
|
||||
mu_at_max_steps = slope_low_res * image_seq_len + intercept_low_res
|
||||
|
||||
if image_seq_len > HIGH_RES_THRESHOLD:
|
||||
return float(mu_at_max_steps)
|
||||
|
||||
mu_at_min_steps = slope_high_res * image_seq_len + intercept_high_res
|
||||
slope = (mu_at_max_steps - mu_at_min_steps) / (MAX_STEPS - MIN_STEPS)
|
||||
intercept = mu_at_max_steps - MAX_STEPS * slope
|
||||
mu = slope * num_steps + intercept
|
||||
|
||||
return float(mu)
|
||||
|
||||
|
||||
def generalized_time_snr_shift(t, mu, sigma):
|
||||
return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)
|
||||
|
||||
|
||||
def get_schedule(steps, width, height):
|
||||
image_seq_len = (width * height / (16 * 16))
|
||||
mu = compute_empirical_mu(image_seq_len, steps)
|
||||
timesteps = torch.linspace(1, 0, steps + 1)
|
||||
timesteps = generalized_time_snr_shift(timesteps, mu, 1.0)
|
||||
|
||||
return timesteps
|
||||
|
||||
|
||||
def build_sampler_sigmas(sampler_name, steps, width, height):
|
||||
sampler = sampler_object(sampler_name)
|
||||
sigmas = get_schedule(steps, width, height)
|
||||
|
||||
return sampler, sigmas
|
||||
|
||||
|
||||
class TT_Flux2WorkflowSettingsNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="TT_Flux2WorkflowSettingsNode",
|
||||
display_name="TT FLUX2 Workflow Settings",
|
||||
category="TenserTensor/Workflow",
|
||||
description="",
|
||||
inputs=[
|
||||
IO.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
|
||||
IO.Int.Input("steps", default=30, min=1, max=10_000),
|
||||
IO.Float.Input("cfg", default=3.0, min=0.0, max=100.0, step=0.1),
|
||||
IO.Float.Input("guidance", default=3.0, min=1.0, max=10.0, step=0.1),
|
||||
IO.Int.Input("width", default=1024, min=16, max=MAX_RESOLUTION, step=8),
|
||||
IO.Int.Input("height", default=1024, min=16, max=MAX_RESOLUTION, step=9),
|
||||
IO.Combo.Input("sampler_name", options=KSampler.SAMPLERS)
|
||||
],
|
||||
outputs=[
|
||||
WorkflowSettings.Output(display_name="WORKFLOW_CONFIG"),
|
||||
IO.Sampler.Output(display_name="SAMPLER"),
|
||||
IO.Sigmas.Output(display_name="SIGMAS"),
|
||||
IO.Int.Output(display_name="SEED"),
|
||||
IO.Int.Output(display_name="STEPS"),
|
||||
IO.Float.Output(display_name="CFG"),
|
||||
IO.Float.Output(display_name="GUIDANCE"),
|
||||
IO.Int.Output(display_name="WIDTH"),
|
||||
IO.Int.Output(display_name="HEIGHT"),
|
||||
]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, **kwargs) -> IO.NodeOutput:
|
||||
kwargs["sampler"], kwargs["sigmas"] = build_sampler_sigmas(
|
||||
kwargs.get("sampler_name"),
|
||||
kwargs.get("steps"),
|
||||
kwargs.get("width"),
|
||||
kwargs.get("height"),
|
||||
)
|
||||
|
||||
workflow_config = TTWorkflowSettings.create(**kwargs)
|
||||
args = {
|
||||
"workflow_config": workflow_config,
|
||||
"sampler": kwargs.get("sampler"),
|
||||
"sigmas": kwargs.get("sigmas"),
|
||||
"seed": kwargs.get("seed"),
|
||||
"steps": kwargs.get("steps"),
|
||||
"cfg": kwargs.get("cfg"),
|
||||
"guidance": kwargs.get("guidance"),
|
||||
"width": kwargs.get("width"),
|
||||
"height": kwargs.get("height"),
|
||||
}
|
||||
|
||||
return IO.NodeOutput(*args.values())
|
||||
|
||||
|
||||
class TT_Flux2WorkflowSettingsAdvancedNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="TT_Flux2WorkflowSettingsAdvancedNode",
|
||||
display_name="TT FLUX2 Workflow Settings (Advanced)",
|
||||
category="TenserTensor/Workflow",
|
||||
description="",
|
||||
inputs=[
|
||||
IO.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff),
|
||||
IO.Int.Input("steps", default=30, min=1, max=10_000),
|
||||
IO.Float.Input("cfg", default=3.0, min=0.0, max=100.0, step=0.1),
|
||||
IO.Combo.Input("sampler_name", options=KSampler.SAMPLERS),
|
||||
IO.Int.Input("width", default=1024, min=16, max=MAX_RESOLUTION, step=8),
|
||||
IO.Int.Input("height", default=1024, min=16, max=MAX_RESOLUTION, step=8),
|
||||
IO.String.Input("prompt", multiline=True, dynamic_prompts=True),
|
||||
IO.String.Input("lora_triggers", multiline=True, dynamic_prompts=True),
|
||||
IO.Float.Input("guidance", default=3.0, min=1.0, max=10.0, step=0.1),
|
||||
],
|
||||
outputs=[
|
||||
WorkflowSettings.Output(display_name="WORKFLOW_CONFIG"),
|
||||
IO.Sampler.Output(display_name="SAMPLER"),
|
||||
IO.Sigmas.Output(display_name="SIGMAS"),
|
||||
IO.Int.Output(display_name="SEED"),
|
||||
IO.Int.Output(display_name="STEPS"),
|
||||
IO.Float.Output(display_name="CFG"),
|
||||
IO.Int.Output(display_name="WIDTH"),
|
||||
IO.Int.Output(display_name="HEIGHT"),
|
||||
IO.String.Output(display_name="PROMPT"),
|
||||
IO.String.Output(display_name="LORA_TRIGGERS"),
|
||||
IO.Float.Output(display_name="GUIDANCE"),
|
||||
]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, **kwargs) -> IO.NodeOutput:
|
||||
kwargs["sampler"], kwargs["sigmas"] = build_sampler_sigmas(
|
||||
kwargs.get("sampler_name"),
|
||||
kwargs.get("steps"),
|
||||
kwargs.get("width"),
|
||||
kwargs.get("height"),
|
||||
)
|
||||
|
||||
workflow_config = TTWorkflowSettings.create(**kwargs)
|
||||
args = {
|
||||
"workflow_config": workflow_config,
|
||||
"sampler": kwargs.get("sampler"),
|
||||
"sigmas": kwargs.get("sigmas"),
|
||||
"seed": kwargs.get("seed"),
|
||||
"steps": kwargs.get("steps"),
|
||||
"cfg": kwargs.get("cfg"),
|
||||
"width": kwargs.get("width"),
|
||||
"height": kwargs.get("height"),
|
||||
"prompt": kwargs.get("prompt"),
|
||||
"lora_triggers": kwargs.get("lora_triggers"),
|
||||
"guidance": kwargs.get("guidance"),
|
||||
}
|
||||
|
||||
return IO.NodeOutput(*args.values())
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# V3 entrypoint — registers context nodes with ComfyUI
|
||||
# ==============================================================================
|
||||
|
||||
class ContextExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
TT_Flux2WorkflowSettingsNode,
|
||||
TT_Flux2WorkflowSettingsAdvancedNode,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> ContextExtension:
|
||||
return ContextExtension()
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Re-exports for backward compatibility
|
||||
# ==============================================================================
|
||||
|
||||
__all__ = [
|
||||
"TT_Flux2WorkflowSettingsNode",
|
||||
"TT_Flux2WorkflowSettingsAdvancedNode",
|
||||
]
|
||||
Reference in New Issue
Block a user