diff --git a/context_nodes.py b/context_nodes.py deleted file mode 100644 index 5b8af23..0000000 --- a/context_nodes.py +++ /dev/null @@ -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", -] diff --git a/lib/common.py b/lib/common.py index 91622ff..5d49a67 100644 --- a/lib/common.py +++ b/lib/common.py @@ -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) diff --git a/lib/workflow_classes.py b/lib/workflow_classes.py deleted file mode 100644 index df9c704..0000000 --- a/lib/workflow_classes.py +++ /dev/null @@ -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}'") diff --git a/node_mappings.py b/node_mappings.py index 47dceff..454a58c 100644 --- a/node_mappings.py +++ b/node_mappings.py @@ -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)", } diff --git a/nodes/workflow/flux2_workflow_settings.py b/nodes/workflow/flux2_workflow_settings.py index 897f6df..0486f86 100644 --- a/nodes/workflow/flux2_workflow_settings.py +++ b/nodes/workflow/flux2_workflow_settings.py @@ -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) diff --git a/nodes/workflow/flux2_workflow_settings_advanced.py b/nodes/workflow/flux2_workflow_settings_advanced.py index 16eb7bc..86b52bf 100644 --- a/nodes/workflow/flux2_workflow_settings_advanced.py +++ b/nodes/workflow/flux2_workflow_settings_advanced.py @@ -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) diff --git a/nodes_context.py b/nodes_context.py new file mode 100644 index 0000000..1c2a069 --- /dev/null +++ b/nodes_context.py @@ -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", +] diff --git a/nodes_workflow.py b/nodes_workflow.py new file mode 100644 index 0000000..01e2a76 --- /dev/null +++ b/nodes_workflow.py @@ -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", +]