119 lines
4.3 KiB
Python
119 lines
4.3 KiB
Python
# (c) TenserTensor <tenser.tensor@proton.me> || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
|
|
|
|
from typing import override
|
|
|
|
from comfy import samplers
|
|
from comfy_api.latest import io
|
|
from node_helpers import conditioning_set_values
|
|
from .utils import raise_unless, CommonTypes
|
|
|
|
CATEGORY = "TenserTensor/Text Encoder"
|
|
|
|
|
|
# ==============================================================================
|
|
# Helper classes — data structures and base types
|
|
# ==============================================================================
|
|
|
|
|
|
class AdaptiveCFGGuider(samplers.CFGGuider):
|
|
@classmethod
|
|
def from_cfg_guider(cls, guider):
|
|
obj = cls.__new__(cls)
|
|
obj.__dict__.update(guider.__dict__)
|
|
|
|
return obj
|
|
|
|
@override
|
|
def set_conds(self, positive, negative=None):
|
|
dict = {"positive": positive}
|
|
if negative is not None:
|
|
dict["negative"] = negative
|
|
|
|
self.inner_set_conds(dict)
|
|
|
|
def get_conds(self, key="positive"):
|
|
return [[c.get("cross_attn", None), c] for c in self.original_conds[key]]
|
|
|
|
|
|
# ==============================================================================
|
|
# Helper functions — pipeline utilities and loaders
|
|
# ==============================================================================
|
|
|
|
def ltxv_prepare_conditioning(**kwargs):
|
|
model = kwargs.get("model")
|
|
raise_unless(model, ValueError, "MODEL is required for text encoder")
|
|
text_encoder = kwargs.get("text_encoder")
|
|
raise_unless(text_encoder, ValueError, "TEXT ENCODER is required for text encoder")
|
|
use_prompts = kwargs.get("use_prompts")
|
|
positive_prompt, negative_prompt = (
|
|
use_prompts.get("positive_prompt"),
|
|
use_prompts.get("negative_prompt"),
|
|
)
|
|
|
|
generated_prompt_tokens = None
|
|
positive_tokens = generated_prompt_tokens or text_encoder.tokenize(positive_prompt)
|
|
negative_tokens = text_encoder.tokenize(negative_prompt)
|
|
fps = int(kwargs.get("frame_rate").replace("fps", ""))
|
|
positive = conditioning_set_values(text_encoder.encode_from_tokens_scheduled(positive_tokens), {"frame_rate": fps})
|
|
negative = conditioning_set_values(text_encoder.encode_from_tokens_scheduled(negative_tokens), {"frame_rate": fps})
|
|
|
|
guider = AdaptiveCFGGuider(model)
|
|
guider.set_conds(positive, negative)
|
|
guider.set_cfg(kwargs.get("cfg"))
|
|
|
|
return guider, positive, negative
|
|
|
|
|
|
# ==============================================================================
|
|
# Node classes — ComfyUI node definitions
|
|
# ==============================================================================
|
|
|
|
|
|
class TT_LtxvTextEncodeNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_LtxvTextEncodeNode",
|
|
display_name="TT LTXV Text Encode",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Model.Input("model"),
|
|
io.Clip.Input("text_encoder"),
|
|
io.Float.Input("cfg", default=3.5, min=0.0, max=20.0, step=0.1),
|
|
io.Combo.Input("frame_rate", options=CommonTypes.FRAME_RATES, default="24fps"),
|
|
io.DynamicCombo.Input("use_prompts", options=[
|
|
io.DynamicCombo.Option("Internal", [
|
|
io.String.Input("positive_prompt", multiline=True, dynamic_prompts=True),
|
|
io.String.Input("negative_prompt", multiline=True, dynamic_prompts=True),
|
|
]),
|
|
io.DynamicCombo.Option("External", [
|
|
io.String.Input("positive_prompt", force_input=True),
|
|
io.String.Input("negative_prompt", force_input=True),
|
|
]),
|
|
]),
|
|
],
|
|
outputs=[
|
|
io.Guider.Output("GUIDER"),
|
|
io.Conditioning.Output("POSITIVE"),
|
|
io.Conditioning.Output("NEGATIVE"),
|
|
],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, **kwargs) -> io.NodeOutput:
|
|
guider, positive, negative = ltxv_prepare_conditioning(**kwargs)
|
|
|
|
return io.NodeOutput(guider, positive, negative)
|
|
|
|
|
|
# ==============================================================================
|
|
# V3 entrypoint — registers context nodes with ComfyUI
|
|
# ==============================================================================
|
|
|
|
|
|
NODES = [
|
|
TT_LtxvTextEncodeNode,
|
|
]
|