diff --git a/node_mappings.py b/node_mappings.py index 7b28084..54a49fa 100644 --- a/node_mappings.py +++ b/node_mappings.py @@ -28,6 +28,11 @@ try: except ImportError: print(f"{YELLOW}TenserTensor: {RED}ERROR: Context nodes unavailable{RESET}") +try: + from .nodes_detector import * +except ImportError: + print(f"{YELLOW}TenserTensor: {RED}ERROR: Detector nodes unavailable{RESET}") + try: from .nodes.image.image_preview_save import TT_ImagePreviewSave from .nodes.image.image_preview_upscale_save import TT_ImagePreviewUpscaleSave @@ -84,6 +89,7 @@ try: from .nodes.text_encoders.clip_text_encode_sdxl_context import TT_ClipTextEncodeSdxlContext from .nodes.text_encoders.clip_text_encode_flux2 import TT_ClipTextEncodeFlux2 from .nodes.text_encoders.clip_text_encode_flux2_context import TT_ClipTextEncodeFlux2Context + from .nodes_text_encoder import * except ImportError: print(f"{YELLOW}TenserTensor: {RED}ERROR: Text Encoder nodes unavailable{RESET}") @@ -141,6 +147,9 @@ NODE_CLASS_MAPPINGS = { "TT_ContextExtractImageNode": TT_ContextExtractImageNode, # ControlNet V3 Nodes "TT_Flux2ApplyControlNetNode": TT_Flux2ApplyControlNetNode, + "TT_Flux2ApplyControlNetAdvancedNode": TT_Flux2ApplyControlNetAdvancedNode, + # Detector V3 Nodes + "TT_CannyEdgeDetectorNode": TT_CannyEdgeDetectorNode, # Image "TT_ImagePreviewSave": TT_ImagePreviewSave, "TT_ImagePreviewUpscaleSave": TT_ImagePreviewUpscaleSave, @@ -183,6 +192,9 @@ NODE_CLASS_MAPPINGS = { "TT_ClipTextEncodeSdxlContext": TT_ClipTextEncodeSdxlContext, "TT_ClipTextEncodeFlux2": TT_ClipTextEncodeFlux2, "TT_ClipTextEncodeFlux2Context": TT_ClipTextEncodeFlux2Context, + # Text Encoder V3 Nodes + "TT_Flux2TextEncoderNode": TT_Flux2TextEncoderNode, + "TT_Flux2TextEncoderContextNode": TT_Flux2TextEncoderContextNode, # VAE "TT_VaeDecodeContext": TT_VaeDecodeContext, "TT_VaeDecodeTiled": TT_VaeDecodeTiled, @@ -226,6 +238,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "TT_ContextExtractImageNode": "TT Context Extract Image Flux2", # ControlNet V3 Nodes "TT_Flux2ApplyControlNetNode": "TT Flux2 Apply ControlNet", + "TT_Flux2ApplyControlNetAdvancedNode": "TT FLUX2 Apply ControlNet (Advanced)", + # Detector V3 Nodes + "TT_CannyEdgeDetectorNode": "TT_Canny Edge Detector", # Image "TT_ImagePreviewSave": "TT Image Preview / Save (Deprecated)", "TT_ImagePreviewUpscaleSave": "TT Image Preview / Upscale / Save (Deprecated)", @@ -268,6 +283,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "TT_ClipTextEncodeSdxlContext": "TT CLIP Text Encode SDXL (Context)", "TT_ClipTextEncodeFlux2": "TT CLIP Text Encode FLUX2", "TT_ClipTextEncodeFlux2Context": "TT CLIP Text Encode FLUX2 (Context)", + # Text Encoder V3 Nodes + "TT_Flux2TextEncoderNode":"", + "TT_Flux2TextEncoderContextNode":"", # VAE "TT_VaeDecodeContext": "TT VAE Decode (Context)", "TT_VaeDecodeTiled": "TT VAE Decode (Tiled)", diff --git a/nodes/text_encoders/clip_text_encode_flux2.py b/nodes/text_encoders/clip_text_encode_flux2.py index 1d05d31..998ba7c 100644 --- a/nodes/text_encoders/clip_text_encode_flux2.py +++ b/nodes/text_encoders/clip_text_encode_flux2.py @@ -19,7 +19,7 @@ class TT_ClipTextEncodeFlux2(): RETURN_TYPES = ("GUIDER",) RETURN_NAMES = ("GUIDER",) FUNCTION = "execute" - CATEGORY = "TenserTensor/Text Encoders/FLUX2" + CATEGORY = "TenserTensor/Deprecated/Text Encoders/FLUX2" def execute(self, model, clip, prompt, lora_triggers, guidance): return encode_prompts_flux2(model, clip, prompt, lora_triggers, guidance) diff --git a/nodes/text_encoders/clip_text_encode_flux2_context.py b/nodes/text_encoders/clip_text_encode_flux2_context.py index 16d3005..c99bb3b 100644 --- a/nodes/text_encoders/clip_text_encode_flux2_context.py +++ b/nodes/text_encoders/clip_text_encode_flux2_context.py @@ -15,7 +15,7 @@ class TT_ClipTextEncodeFlux2Context(): RETURN_TYPES = ("TT_CONTEXT", "GUIDER",) RETURN_NAMES = ("CONTEXT", "GUIDER",) FUNCTION = "execute" - CATEGORY = "TenserTensor/Text Encoders/FLUX2" + CATEGORY = "TenserTensor/Deprecated/Text Encoders/FLUX2" def execute(self, context): model = context.get("model", None) diff --git a/nodes_context.py b/nodes_context.py index a3580b6..ac6f859 100644 --- a/nodes_context.py +++ b/nodes_context.py @@ -112,8 +112,8 @@ class Context: super().__init__(id, **kwargs) class Output(IO.Output): - def __init__(self, **kwargs): - super().__init__(**kwargs) + def __init__(self, id: str, **kwargs): + super().__init__(id, **kwargs) class ContextNode(IO.ComfyNode): diff --git a/nodes_controlnet.py b/nodes_controlnet.py index 5e955d4..d66d42f 100644 --- a/nodes_controlnet.py +++ b/nodes_controlnet.py @@ -4,13 +4,13 @@ from typing import override import folder_paths from comfy import controlnet -from comfy_api.latest import io, ComfyExtension -from .nodes_image import SingleCondCFGGuider +from comfy_api.latest import io, ui, ComfyExtension +from .nodes_image import SingleCondCFGGuider, get_image_files, load_image CATEGORY = "TenserTensor/ControlNet" -def get_control_net_names(): +def get_control_net_files(): return folder_paths.get_filename_list("controlnet") @@ -26,7 +26,7 @@ def load_control_net(filename): def apply_control_net(**kwargs): guider, image, loaded_cnet = ( kwargs.get("guider"), - kwargs.get("image"), + kwargs.get("cond_hint"), kwargs.get("loaded_cnet"), ) conditioning = guider.get_conds() @@ -69,9 +69,10 @@ class TT_Flux2ApplyControlNetNode(io.ComfyNode): description="", inputs=[ io.Guider.Input("guider"), - io.Image.Input("image"), + io.Image.Input("reference_image"), + io.ControlNet.Input("control_net_opt", optional=True), io.Vae.Input("vae_opt", optional=True), - io.Combo.Input("control_net", options=get_control_net_names()), + io.Combo.Input("control_net", options=get_control_net_files()), io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.1), io.Float.Input("start_percent", default=0.0, min=0.0, max=100.0, step=0.1, advanced=True), io.Float.Input("end_percent", default=100.0, min=0.0, max=100.0, step=0.1, advanced=True), @@ -94,12 +95,73 @@ class TT_Flux2ApplyControlNetNode(io.ComfyNode): if strength == 0.0: return io.NodeOutput(guider) - kwargs["loaded_cnet"] = load_control_net(kwargs.get("control_net")) + loaded_cnet = kwargs.get("control_net_opt") + kwargs["loaded_cnet"] = ( + loaded_cnet + if loaded_cnet is not None + else load_control_net(kwargs.get("control_net")) + ) + kwargs["cond_hint"] = kwargs.get("reference_image") guider = apply_control_net(**kwargs) return io.NodeOutput(guider) +class TT_Flux2ApplyControlNetAdvancedNode(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="TT_Flux2ApplyControlNetAdvancedNode", + display_name="TT FLUX2 Apply ControlNet (Advanced)", + category=CATEGORY, + description="", + inputs=[ + io.Guider.Input("guider"), + io.ControlNet.Input("control_net_opt", optional=True), + io.Image.Input("reference_image_opt", optional=True), + io.Vae.Input("vae_opt", optional=True), + io.Combo.Input("control_net", options=get_control_net_files()), + io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.1), + io.Float.Input("start_percent", default=0.0, min=0.0, max=100.0, step=0.1, advanced=True), + io.Float.Input("end_percent", default=100.0, min=0.0, max=100.0, step=0.1, advanced=True), + io.Combo.Input("reference_image", options=get_image_files(), upload=io.UploadType.image) + ], + outputs=[ + io.Guider.Output("GUIDER"), + ] + ) + + @classmethod + def execute(cls, **kwargs) -> io.NodeOutput: + guider, strength, control_net = ( + kwargs.get("guider"), + kwargs.get("strength"), + kwargs.get("control_net")) + + if not isinstance(guider, SingleCondCFGGuider): + guider = SingleCondCFGGuider.from_cfg_guider(guider) + + if strength == 0.0: + return io.NodeOutput(guider) + + reference_image = kwargs.get("reference_image_opt") + kwargs["cond_hint"] = ( + reference_image + if reference_image is not None + else load_image(kwargs.get("reference_image"))[0] + ) + + loaded_cnet = kwargs.get("control_net_opt") + kwargs["loaded_cnet"] = ( + loaded_cnet + if loaded_cnet is not None + else load_control_net(kwargs.get("control_net")) + ) + guider = apply_control_net(**kwargs) + + return io.NodeOutput(guider, ui=ui.PreviewImage(kwargs["cond_hint"], cls=cls)) + + # ============================================================================== # V3 entrypoint — registers context nodes with ComfyUI # ============================================================================== @@ -109,6 +171,7 @@ class ControlNetNodesExtension(ComfyExtension): async def get_node_list(self) -> list[type[io.ComfyNode]]: return [ TT_Flux2ApplyControlNetNode, + TT_Flux2ApplyControlNetAdvancedNode, ] @@ -122,4 +185,5 @@ async def comfy_entrypoint() -> ControlNetNodesExtension: __all__ = [ "TT_Flux2ApplyControlNetNode", + "TT_Flux2ApplyControlNetAdvancedNode", ] diff --git a/nodes_detector.py b/nodes_detector.py new file mode 100644 index 0000000..dbdafe2 --- /dev/null +++ b/nodes_detector.py @@ -0,0 +1,72 @@ +# (c) TenserTensor || Apache-2.0 (apache.org/licenses/LICENSE-2.0) + + +from typing import override + +from kornia.filters import canny + +from comfy import model_management +from comfy_api.latest import io, ui, ComfyExtension +from .nodes_image import get_image_files, load_image + +CATEGORY = "TenserTensor/Detector" + + +def detect_edge(image, low_threshold, high_threshold): + torch_device = model_management.get_torch_device() + intermediate_device = model_management.intermediate_device() + timage = image.to(torch_device).movedim(-1, 1) + o_images = canny(timage, low_threshold, high_threshold) + + return o_images[1].to(intermediate_device).expand(-1, 3, -1, -1).movedim(1, -1) + + +class TT_CannyEdgeDetectorNode(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="TT_CannyEdgeDetectorNode", + display_name="TT Canny Edge Detector", + category=CATEGORY, + description="", + inputs=[ + io.Float.Input("low_threshold", default=0.4, min=0.01, max=0.99, step=0.01), + io.Float.Input("high_threshold", default=0.8, min=0.01, max=0.99, step=0.01), + io.Combo.Input("reference_image", options=get_image_files(), upload=io.UploadType.image) + ], + outputs=[ + io.Image.Output("IMAGE") + ] + ) + + @classmethod + def execute(cls, **kwargs) -> io.NodeOutput: + timage, _ = load_image(kwargs.get("reference_image")) + cn_image = detect_edge(timage, kwargs.get("low_threshold"), kwargs.get("high_threshold")) + + return io.NodeOutput(cn_image, ui=ui.PreviewImage(timage, cls=cls)) + + +# ============================================================================== +# V3 entrypoint — registers context nodes with ComfyUI +# ============================================================================== + +class DetectorNodesExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[io.ComfyNode]]: + return [ + TT_CannyEdgeDetectorNode, + ] + + +async def comfy_entrypoint() -> DetectorNodesExtension: + return DetectorNodesExtension() + + +# ============================================================================== +# Re-exports for backward compatibility +# ============================================================================== + +__all__ = [ + "TT_CannyEdgeDetectorNode", +] diff --git a/nodes_image.py b/nodes_image.py index 1ba055d..e95a1aa 100644 --- a/nodes_image.py +++ b/nodes_image.py @@ -471,12 +471,12 @@ class TT_GuiderImageReferenceNode(io.ComfyNode): guider = SingleCondCFGGuider.from_cfg_guider(guider) resized = resize_image_to_megapixels(timage, resize_method, megapixels, dimension_step) - samples = vae_encode(resized, vae) + latent = vae_encode(resized, vae) conditioning = guider.get_conds() if conditioning is not None: - conditioning = conditioning_set_values(conditioning, {"reference_latents": [samples]}, append=True) + conditioning = conditioning_set_values(conditioning, {"reference_latents": [latent["samples"]]}, append=True) guider.set_conds(conditioning) else: raise ValueError("ERROR: Guider has no conditioning") diff --git a/nodes_text_encoder.py b/nodes_text_encoder.py index b726455..2306d43 100644 --- a/nodes_text_encoder.py +++ b/nodes_text_encoder.py @@ -3,16 +3,96 @@ 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): - pass + @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): - pass + @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) # ============================================================================== diff --git a/pyproject.toml b/pyproject.toml index 6956253..323be90 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,9 +2,9 @@ [project] name = "ComfyUI-TenserTensor" description = "A collection of custom nodes for ComfyUI that streamline common workflows. Includes all-in-one model loaders (checkpoint/unet, LoRA, CLIP, VAE), GGUF quantized model support (FLUX, FLUX2, SDXL), workflow settings nodes, latent factory (aspect ratio, megapixels, FLUX/SDXL), context nodes, CLIP text encoders, KSamplers (context, advanced, two-stage - draft/refine), VAE encoders/decoders (context, tiled), and post-production tools (brightness, contrast, LUT color grading, film grain, upscaler)." -version = "1.5.11" +version = "1.6.12" license = { file = "LICENSE" } -dependencies = ["gguf"] +dependencies = ["gguf", "kornia"] [project.urls] Repository = "https://github.com/tenser-tensor/ComfyUI-TenserTensor" diff --git a/requirements.txt b/requirements.txt index 0c704fa..172ecd7 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1 +1,2 @@ gguf>=0.17.0 +kornia \ No newline at end of file