123 lines
3.8 KiB
Python
123 lines
3.8 KiB
Python
# (c) TenserTensor <tenser.tensor@proton.me> || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
|
|
from typing import override
|
|
|
|
from comfy_api.latest import io, ComfyExtension
|
|
from .nodes_context import Context
|
|
from .nodes_image import SingleCondCFGGuider
|
|
|
|
CATEGORY = "TenserTensor/Text Encoder"
|
|
|
|
|
|
def encode_prompts_flux2(**kwargs):
|
|
clip = kwargs.get("clip")
|
|
model = kwargs.get("model")
|
|
|
|
if model is None:
|
|
raise ValueError("ERROR: MODEL is required for text encoder")
|
|
|
|
if clip is None:
|
|
raise ValueError("ERROR: CLIP is required for text encoder")
|
|
|
|
prompt, lora_triggers, guidance = (
|
|
kwargs.get("prompt"),
|
|
kwargs.get("lora_triggers"),
|
|
kwargs.get("guidance"),
|
|
)
|
|
|
|
full_prompt = f"{lora_triggers}, {prompt}" if lora_triggers.strip() else prompt
|
|
tokens = clip.tokenize(full_prompt)
|
|
conditioning = clip.encode_from_tokens_scheduled(tokens, add_dict={"guidance": guidance, })
|
|
guider = SingleCondCFGGuider(model)
|
|
guider.set_conds(conditioning)
|
|
|
|
return guider
|
|
|
|
|
|
class TT_Flux2TextEncoderNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_Flux2TextEncoderNode",
|
|
display_name="TT FLUX2 Text Encoder",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Clip.Input("clip"),
|
|
io.String.Input("prompt", multiline=True, placeholder="Prompt", dynamic_prompts=True),
|
|
io.String.Input("lora_triggers", multiline=True, placeholder="LoRA Triggers", dynamic_prompts=True),
|
|
io.Float.Input("guidance", default=3.5, min=1.0, max=10.0, step=0.1)
|
|
],
|
|
outputs=[
|
|
io.Guider.Output("GUIDER"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
guider = encode_prompts_flux2(**kwargs)
|
|
|
|
return io.NodeOutput(guider)
|
|
|
|
|
|
class TT_Flux2TextEncoderContextNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_Flux2TextEncoderContextNode",
|
|
display_name="TT FLUX2 Text Encoder (Context)",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
Context.Input("context")
|
|
],
|
|
outputs=[
|
|
Context.Output("CONTEXT"),
|
|
io.Guider.Output("GUIDER"),
|
|
]
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
context = kwargs.get("context")
|
|
|
|
args = {
|
|
"model": context.get_attr("model"),
|
|
"clip": context.get_attr("clip"),
|
|
"prompt": context.get_attr("prompt"),
|
|
"lora_triggers": context.get_attr("lora_triggers"),
|
|
"guidance": context.get_attr("guidance"),
|
|
}
|
|
guider = encode_prompts_flux2(**args)
|
|
context.set_attr("guider", guider)
|
|
|
|
return io.NodeOutput(context, guider)
|
|
|
|
|
|
# ==============================================================================
|
|
# V3 entrypoint — registers context nodes with ComfyUI
|
|
# ==============================================================================
|
|
|
|
class TextEncodeNodesExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
TT_Flux2TextEncoderNode,
|
|
TT_Flux2TextEncoderContextNode,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> TextEncodeNodesExtension:
|
|
return TextEncodeNodesExtension()
|
|
|
|
|
|
# ==============================================================================
|
|
# Re-exports for backward compatibility
|
|
# ==============================================================================
|
|
|
|
__all__ = [
|
|
"TT_Flux2TextEncoderNode",
|
|
"TT_Flux2TextEncoderContextNode",
|
|
]
|