194 lines
6.6 KiB
Python
194 lines
6.6 KiB
Python
# (c) TenserTensor <tenser.tensor@proton.me> || Apache-2.0 (apache.org/licenses/LICENSE-2.0)
|
|
|
|
from typing import override
|
|
|
|
import folder_paths
|
|
from comfy import controlnet
|
|
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_files():
|
|
return folder_paths.get_filename_list("controlnet")
|
|
|
|
|
|
def load_control_net(filename):
|
|
cnet_path = folder_paths.get_full_path_or_raise("controlnet", filename)
|
|
|
|
print(f"CNET path: {cnet_path}")
|
|
|
|
cnet = controlnet.load_controlnet(cnet_path)
|
|
|
|
if cnet is None:
|
|
raise RuntimeError("ERROR: ControlNet file is invalid or corrupted.")
|
|
|
|
return cnet
|
|
|
|
|
|
def apply_control_net(**kwargs):
|
|
guider, image, loaded_cnet = (
|
|
kwargs.get("guider"),
|
|
kwargs.get("cond_hint"),
|
|
kwargs.get("loaded_cnet"),
|
|
)
|
|
conditioning = guider.get_conds()
|
|
|
|
cnet_cache = {}
|
|
cond_hint = image.movedim(-1, 1)
|
|
for tensor in conditioning:
|
|
cond_copy = tensor[1].copy()
|
|
prev_cnet = cond_copy.get("control", None)
|
|
if prev_cnet in cnet_cache:
|
|
applied_cnet = cnet_cache[prev_cnet]
|
|
else:
|
|
applied_cnet = loaded_cnet.copy()
|
|
applied_cnet.set_cond_hint(
|
|
cond_hint,
|
|
kwargs.get("strength"),
|
|
(kwargs.get("start_percent") / 100, kwargs.get("end_percent") / 100),
|
|
vae=kwargs.get("vae_opt"),
|
|
extra_concat=kwargs.get("extra_concat", None)
|
|
)
|
|
applied_cnet.set_previous_controlnet(prev_cnet)
|
|
cnet_cache[prev_cnet] = applied_cnet
|
|
|
|
cond_copy["control"] = applied_cnet
|
|
cond_copy['control_apply_to_uncond'] = False
|
|
updated_cond = [tensor[0], cond_copy]
|
|
|
|
guider.set_conds(updated_cond)
|
|
|
|
return guider
|
|
|
|
|
|
class TT_Flux2ApplyControlNetNode(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="TT_Flux2ApplyControlNetNode",
|
|
display_name="TT FLUX2 Apply ControlNet",
|
|
category=CATEGORY,
|
|
description="",
|
|
inputs=[
|
|
io.Guider.Input("guider"),
|
|
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_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),
|
|
],
|
|
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)
|
|
|
|
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
|
|
# ==============================================================================
|
|
|
|
class ControlNetNodesExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
TT_Flux2ApplyControlNetNode,
|
|
TT_Flux2ApplyControlNetAdvancedNode,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> ControlNetNodesExtension:
|
|
return ControlNetNodesExtension()
|
|
|
|
|
|
# ==============================================================================
|
|
# Re-exports for backward compatibility
|
|
# ==============================================================================
|
|
|
|
__all__ = [
|
|
"TT_Flux2ApplyControlNetNode",
|
|
"TT_Flux2ApplyControlNetAdvancedNode",
|
|
]
|