359 lines
12 KiB
Python
359 lines
12 KiB
Python
# (c) TenserTensor <tenser.tensor@proton.me> || 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
|
|
guidance: IO.Float | None = None
|
|
guider: IO.Guider | None = None
|
|
image: IO.Image | None = None
|
|
latent: IO.Latent | None = None
|
|
lora_triggers: IO.String | None = None
|
|
mask: IO.Mask | None = None
|
|
model: IO.Model | None = None
|
|
negative: IO.Conditioning | None = None
|
|
positive: IO.Conditioning | None = None
|
|
prompt: IO.String | 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 attr is not None:
|
|
if hasattr(self.workflow_config, key):
|
|
self.set_workflow_config_attr(key, attr)
|
|
else:
|
|
setattr(self, key, attr)
|
|
|
|
def get_attr(self, key, default=None):
|
|
if hasattr(self.workflow_config, key):
|
|
val = self.get_workflow_config_attr(key, default)
|
|
else:
|
|
val = getattr(self, key)
|
|
|
|
return val if val is not None else default
|
|
|
|
def update_attrs(self, **kwargs):
|
|
for key in kwargs.keys():
|
|
val = kwargs.get(key, None)
|
|
if hasattr(self.workflow_config, key):
|
|
self.set_workflow_config_attr(key, val)
|
|
else:
|
|
self.set_attr(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, id: str, **kwargs):
|
|
super().__init__(id, **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.REQUIRED + cls.OPTIONAL:
|
|
output_schema.append(
|
|
Context.Output("CONTEXT")
|
|
if f_name == "context"
|
|
else cls.F_TYPES[f_name].Output(id=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.REQUIRED + cls.OPTIONAL:
|
|
if f_name != "context":
|
|
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 = context.get_attr("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", "rnd_noise",)
|
|
OPTIONAL = ("workflow_config", "sampler", "sigmas",)
|
|
RETURNS = ("context",)
|
|
|
|
|
|
class TT_ContextPassthroughNode(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 = ()
|
|
|
|
|
|
class TT_ContextSetGuiderNode(ContextNode):
|
|
"""Sets the guider field in the context object."""
|
|
REQUIRED = ("context", "guider",)
|
|
OPTIONAL = ()
|
|
RETURNS = ("context",)
|
|
|
|
|
|
class TT_ContextSetImageNode(ContextNode):
|
|
"""Sets the image field in the context object."""
|
|
REQUIRED = ("context", "image",)
|
|
OPTIONAL = ()
|
|
RETURNS = ("context",)
|
|
|
|
|
|
class TT_ContextSetLatentNode(ContextNode):
|
|
"""Sets the latent field in the context object."""
|
|
REQUIRED = ("context", "latent",)
|
|
OPTIONAL = ()
|
|
RETURNS = ("context",)
|
|
|
|
|
|
class TT_ContextExtractEncoderFlux2Node(ContextNode):
|
|
"""Extracts model, clip, and encoding parameters from context for use with a FLUX2 text encoder."""
|
|
REQUIRED = ("context",)
|
|
OPTIONAL = ()
|
|
RETURNS = ("context", "model", "clip", "prompt", "lora_triggers", "guidance",)
|
|
|
|
|
|
class TT_ContextExtractGuidedSamplerFlux2Node(ContextNode):
|
|
"""Extracts latent, guider, sigmas, sampler, and noise from context for use with a FLUX2 guided sampler."""
|
|
REQUIRED = ("context",)
|
|
OPTIONAL = ()
|
|
RETURNS = ("context", "latent", "guider", "sigmas", "sampler", "rnd_noise",)
|
|
|
|
|
|
class TT_ContextExtractVaeNode(ContextNode):
|
|
"""Extracts VAE from context for use with a VAE decoder or encoder node."""
|
|
REQUIRED = ("context",)
|
|
OPTIONAL = ()
|
|
RETURNS = ("context", "vae",)
|
|
|
|
|
|
class TT_ContextExtractImageNode(ContextNode):
|
|
"""Extracts Image from context for use with e.g. preview or postproduction node."""
|
|
REQUIRED = ("context",)
|
|
OPTIONAL = ()
|
|
RETURNS = ("context", "image",)
|
|
|
|
|
|
# ==============================================================================
|
|
# V3 entrypoint — registers context nodes with ComfyUI
|
|
# ==============================================================================
|
|
|
|
class ContextNodesExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
|
return [
|
|
TT_BaseContextNode,
|
|
TT_BaseContextFlux2Node,
|
|
TT_ContextPassthroughNode,
|
|
TT_ContextNode,
|
|
TT_ContextFlux2Node,
|
|
TT_ContextSetGuiderNode,
|
|
TT_ContextSetImageNode,
|
|
TT_ContextSetLatentNode,
|
|
TT_ContextExtractEncoderFlux2Node,
|
|
TT_ContextExtractGuidedSamplerFlux2Node,
|
|
TT_ContextExtractVaeNode,
|
|
TT_ContextExtractImageNode,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> ContextNodesExtension:
|
|
return ContextNodesExtension()
|
|
|
|
|
|
# ==============================================================================
|
|
# Re-exports for backward compatibility
|
|
# ==============================================================================
|
|
|
|
__all__ = [
|
|
"TT_BaseContextNode",
|
|
"TT_BaseContextFlux2Node",
|
|
"TT_ContextPassthroughNode",
|
|
"TT_ContextNode",
|
|
"TT_ContextFlux2Node",
|
|
"TT_ContextSetGuiderNode",
|
|
"TT_ContextSetImageNode",
|
|
"TT_ContextSetLatentNode",
|
|
"TT_ContextExtractEncoderFlux2Node",
|
|
"TT_ContextExtractGuidedSamplerFlux2Node",
|
|
"TT_ContextExtractVaeNode",
|
|
"TT_ContextExtractImageNode",
|
|
]
|