v3 nodes (part c)

This commit is contained in:
bigcat88 2025-08-03 09:50:51 +03:00
parent 13aaa66ec2
commit 99e868f7b2
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
8 changed files with 1038 additions and 815 deletions

View File

@ -1,12 +1,13 @@
import nodes from __future__ import annotations
import torch
import numpy as np import numpy as np
import torch
from einops import rearrange from einops import rearrange
from typing_extensions import override
import comfy.model_management import comfy.model_management
import nodes
from comfy_api.latest import ComfyExtension, io
MAX_RESOLUTION = nodes.MAX_RESOLUTION
CAMERA_DICT = { CAMERA_DICT = {
"base_T_norm": 1.5, "base_T_norm": 1.5,
@ -148,32 +149,48 @@ def get_camera_motion(angle, T, speed, n=81):
RT = np.stack(RT) RT = np.stack(RT)
return RT return RT
class WanCameraEmbedding:
class WanCameraEmbedding(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(cls): def define_schema(cls):
return { return io.Schema(
"required": { node_id="WanCameraEmbedding",
"camera_pose":(["Static","Pan Up","Pan Down","Pan Left","Pan Right","Zoom In","Zoom Out","Anti Clockwise (ACW)", "ClockWise (CW)"],{"default":"Static"}), category="camera",
"width": ("INT", {"default": 832, "min": 16, "max": MAX_RESOLUTION, "step": 16}), inputs=[
"height": ("INT", {"default": 480, "min": 16, "max": MAX_RESOLUTION, "step": 16}), io.Combo.Input(
"length": ("INT", {"default": 81, "min": 1, "max": MAX_RESOLUTION, "step": 4}), "camera_pose",
}, options=[
"optional":{ "Static",
"speed":("FLOAT",{"default":1.0, "min": 0, "max": 10.0, "step": 0.1}), "Pan Up",
"fx":("FLOAT",{"default":0.5, "min": 0, "max": 1, "step": 0.000000001}), "Pan Down",
"fy":("FLOAT",{"default":0.5, "min": 0, "max": 1, "step": 0.000000001}), "Pan Left",
"cx":("FLOAT",{"default":0.5, "min": 0, "max": 1, "step": 0.01}), "Pan Right",
"cy":("FLOAT",{"default":0.5, "min": 0, "max": 1, "step": 0.01}), "Zoom In",
} "Zoom Out",
"Anti Clockwise (ACW)",
"ClockWise (CW)",
],
default="Static",
),
io.Int.Input("width", default=832, min=16, max=nodes.MAX_RESOLUTION, step=16),
io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
io.Int.Input("length", default=81, min=1, max=nodes.MAX_RESOLUTION, step=4),
io.Float.Input("speed", default=1.0, min=0, max=10.0, step=0.1, optional=True),
io.Float.Input("fx", default=0.5, min=0, max=1, step=0.000000001, optional=True),
io.Float.Input("fy", default=0.5, min=0, max=1, step=0.000000001, optional=True),
io.Float.Input("cx", default=0.5, min=0, max=1, step=0.01, optional=True),
io.Float.Input("cy", default=0.5, min=0, max=1, step=0.01, optional=True),
],
outputs=[
io.WanCameraEmbedding.Output(display_name="camera_embedding"),
io.Int.Output(display_name="width"),
io.Int.Output(display_name="height"),
io.Int.Output(display_name="length"),
],
)
} @classmethod
def execute(cls, camera_pose, width, height, length, speed=1.0, fx=0.5, fy=0.5, cx=0.5, cy=0.5) -> io.NodeOutput:
RETURN_TYPES = ("WAN_CAMERA_EMBEDDING","INT","INT","INT")
RETURN_NAMES = ("camera_embedding","width","height","length")
FUNCTION = "run"
CATEGORY = "camera"
def run(self, camera_pose, width, height, length, speed=1.0, fx=0.5, fy=0.5, cx=0.5, cy=0.5):
""" """
Use Camera trajectory as extrinsic parameters to calculate Plücker embeddings (Sitzmannet al., 2021) Use Camera trajectory as extrinsic parameters to calculate Plücker embeddings (Sitzmannet al., 2021)
Adapted from https://github.com/aigc-apps/VideoX-Fun/blob/main/comfyui/comfyui_nodes.py Adapted from https://github.com/aigc-apps/VideoX-Fun/blob/main/comfyui/comfyui_nodes.py
@ -210,9 +227,17 @@ class WanCameraEmbedding:
control_camera_video = control_camera_video.contiguous().view(b, f // 4, 4, c, h, w).transpose(2, 3) control_camera_video = control_camera_video.contiguous().view(b, f // 4, 4, c, h, w).transpose(2, 3)
control_camera_video = control_camera_video.contiguous().view(b, f // 4, c * 4, h, w).transpose(1, 2) control_camera_video = control_camera_video.contiguous().view(b, f // 4, c * 4, h, w).transpose(1, 2)
return (control_camera_video, width, height, length) return io.NodeOutput(control_camera_video, width, height, length)
NODE_CLASS_MAPPINGS = { NODES_LIST: list[type[io.ComfyNode]] = [
"WanCameraEmbedding": WanCameraEmbedding, WanCameraEmbedding,
} ]
class CameraTrajectoryExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return NODES_LIST
async def comfy_entrypoint() -> CameraTrajectoryExtension:
return CameraTrajectoryExtension()

View File

@ -1,25 +1,41 @@
from __future__ import annotations
from kornia.filters import canny from kornia.filters import canny
from typing_extensions import override
import comfy.model_management import comfy.model_management
from comfy_api.latest import ComfyExtension, io
class Canny: class Canny(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return {"required": {"image": ("IMAGE",), return io.Schema(
"low_threshold": ("FLOAT", {"default": 0.4, "min": 0.01, "max": 0.99, "step": 0.01}), node_id="Canny",
"high_threshold": ("FLOAT", {"default": 0.8, "min": 0.01, "max": 0.99, "step": 0.01}) category="image/preprocessors",
}} inputs=[
io.Image.Input("image"),
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),
],
outputs=[io.Image.Output()],
)
RETURN_TYPES = ("IMAGE",) @classmethod
FUNCTION = "detect_edge" def execute(cls, image, low_threshold, high_threshold) -> io.NodeOutput:
CATEGORY = "image/preprocessors"
def detect_edge(self, image, low_threshold, high_threshold):
output = canny(image.to(comfy.model_management.get_torch_device()).movedim(-1, 1), low_threshold, high_threshold) output = canny(image.to(comfy.model_management.get_torch_device()).movedim(-1, 1), low_threshold, high_threshold)
img_out = output[1].to(comfy.model_management.intermediate_device()).repeat(1, 3, 1, 1).movedim(1, -1) img_out = output[1].to(comfy.model_management.intermediate_device()).repeat(1, 3, 1, 1).movedim(1, -1)
return (img_out,) return io.NodeOutput(img_out)
NODE_CLASS_MAPPINGS = {
"Canny": Canny, NODES_LIST: list[type[io.ComfyNode]] = [
} Canny,
]
class CannyExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return NODES_LIST
async def comfy_entrypoint() -> CannyExtension:
return CannyExtension()

View File

@ -1,4 +1,10 @@
from __future__ import annotations
import torch import torch
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
# https://github.com/WeichenFan/CFG-Zero-star # https://github.com/WeichenFan/CFG-Zero-star
def optimized_scale(positive, negative): def optimized_scale(positive, negative):
@ -16,17 +22,20 @@ def optimized_scale(positive, negative):
return st_star.reshape([positive.shape[0]] + [1] * (positive.ndim - 1)) return st_star.reshape([positive.shape[0]] + [1] * (positive.ndim - 1))
class CFGZeroStar: class CFGZeroStar(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls) -> io.Schema:
return {"required": {"model": ("MODEL",), return io.Schema(
}} node_id="CFGZeroStar",
RETURN_TYPES = ("MODEL",) category="advanced/guidance",
RETURN_NAMES = ("patched_model",) inputs=[
FUNCTION = "patch" io.Model.Input("model"),
CATEGORY = "advanced/guidance" ],
outputs=[io.Model.Output(display_name="patched_model")],
)
def patch(self, model): @classmethod
def execute(cls, model) -> io.NodeOutput:
m = model.clone() m = model.clone()
def cfg_zero_star(args): def cfg_zero_star(args):
guidance_scale = args['cond_scale'] guidance_scale = args['cond_scale']
@ -38,21 +47,24 @@ class CFGZeroStar:
return out + uncond_p * (alpha - 1.0) + guidance_scale * uncond_p * (1.0 - alpha) return out + uncond_p * (alpha - 1.0) + guidance_scale * uncond_p * (1.0 - alpha)
m.set_model_sampler_post_cfg_function(cfg_zero_star) m.set_model_sampler_post_cfg_function(cfg_zero_star)
return (m, ) return io.NodeOutput(m)
class CFGNorm: class CFGNorm(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls) -> io.Schema:
return {"required": {"model": ("MODEL",), return io.Schema(
"strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 100.0, "step": 0.01}), node_id="CFGNorm",
}} category="advanced/guidance",
RETURN_TYPES = ("MODEL",) inputs=[
RETURN_NAMES = ("patched_model",) io.Model.Input("model"),
FUNCTION = "patch" io.Float.Input("strength", default=1.0, min=0.0, max=100.0, step=0.01),
CATEGORY = "advanced/guidance" ],
EXPERIMENTAL = True outputs=[io.Model.Output(display_name="patched_model")],
is_experimental=True,
)
def patch(self, model, strength): @classmethod
def execute(cls, model, strength) -> io.NodeOutput:
m = model.clone() m = model.clone()
def cfg_norm(args): def cfg_norm(args):
cond_p = args['cond_denoised'] cond_p = args['cond_denoised']
@ -64,9 +76,18 @@ class CFGNorm:
return pred_text_ * scale * strength return pred_text_ * scale * strength
m.set_model_sampler_post_cfg_function(cfg_norm) m.set_model_sampler_post_cfg_function(cfg_norm)
return (m, ) return io.NodeOutput(m)
NODE_CLASS_MAPPINGS = {
"CFGZeroStar": CFGZeroStar, NODES_LIST: list[type[io.ComfyNode]] = [
"CFGNorm": CFGNorm, CFGNorm,
} CFGZeroStar,
]
class CfgExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return NODES_LIST
async def comfy_entrypoint() -> CfgExtension:
return CfgExtension()

View File

@ -1,43 +1,54 @@
from nodes import MAX_RESOLUTION from __future__ import annotations
class CLIPTextEncodeSDXLRefiner: from typing_extensions import override
import nodes
from comfy_api.latest import ComfyExtension, io
class CLIPTextEncodeSDXLRefiner(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return {"required": { return io.Schema(
"ascore": ("FLOAT", {"default": 6.0, "min": 0.0, "max": 1000.0, "step": 0.01}), node_id="CLIPTextEncodeSDXLRefiner",
"width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), category="advanced/conditioning",
"height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), inputs=[
"text": ("STRING", {"multiline": True, "dynamicPrompts": True}), "clip": ("CLIP", ), io.Float.Input("ascore", default=6.0, min=0.0, max=1000.0, step=0.01),
}} io.Int.Input("width", default=1024, min=0, max=nodes.MAX_RESOLUTION),
RETURN_TYPES = ("CONDITIONING",) io.Int.Input("height", default=1024, min=0, max=nodes.MAX_RESOLUTION),
FUNCTION = "encode" io.String.Input("text", multiline=True, dynamic_prompts=True),
io.Clip.Input("clip"),
],
outputs=[io.Conditioning.Output()],
)
CATEGORY = "advanced/conditioning" @classmethod
def execute(cls, ascore, width, height, text, clip) -> io.NodeOutput:
def encode(self, clip, ascore, width, height, text):
tokens = clip.tokenize(text) tokens = clip.tokenize(text)
return (clip.encode_from_tokens_scheduled(tokens, add_dict={"aesthetic_score": ascore, "width": width, "height": height}), ) return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens, add_dict={"aesthetic_score": ascore, "width": width, "height": height}))
class CLIPTextEncodeSDXL: class CLIPTextEncodeSDXL(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return {"required": { return io.Schema(
"clip": ("CLIP", ), node_id="CLIPTextEncodeSDXL",
"width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), category="advanced/conditioning",
"height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), inputs=[
"crop_w": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}), io.Clip.Input("clip"),
"crop_h": ("INT", {"default": 0, "min": 0, "max": MAX_RESOLUTION}), io.Int.Input("width", default=1024, min=0, max=nodes.MAX_RESOLUTION),
"target_width": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), io.Int.Input("height", default=1024, min=0, max=nodes.MAX_RESOLUTION),
"target_height": ("INT", {"default": 1024.0, "min": 0, "max": MAX_RESOLUTION}), io.Int.Input("crop_w", default=0, min=0, max=nodes.MAX_RESOLUTION),
"text_g": ("STRING", {"multiline": True, "dynamicPrompts": True}), io.Int.Input("crop_h", default=0, min=0, max=nodes.MAX_RESOLUTION),
"text_l": ("STRING", {"multiline": True, "dynamicPrompts": True}), io.Int.Input("target_width", default=1024, min=0, max=nodes.MAX_RESOLUTION),
}} io.Int.Input("target_height", default=1024, min=0, max=nodes.MAX_RESOLUTION),
RETURN_TYPES = ("CONDITIONING",) io.String.Input("text_g", multiline=True, dynamic_prompts=True),
FUNCTION = "encode" io.String.Input("text_l", multiline=True, dynamic_prompts=True),
],
outputs=[io.Conditioning.Output()],
)
CATEGORY = "advanced/conditioning" @classmethod
def execute(cls, clip, width, height, crop_w, crop_h, target_width, target_height, text_g, text_l) -> io.NodeOutput:
def encode(self, clip, width, height, crop_w, crop_h, target_width, target_height, text_g, text_l):
tokens = clip.tokenize(text_g) tokens = clip.tokenize(text_g)
tokens["l"] = clip.tokenize(text_l)["l"] tokens["l"] = clip.tokenize(text_l)["l"]
if len(tokens["l"]) != len(tokens["g"]): if len(tokens["l"]) != len(tokens["g"]):
@ -46,9 +57,18 @@ class CLIPTextEncodeSDXL:
tokens["l"] += empty["l"] tokens["l"] += empty["l"]
while len(tokens["l"]) > len(tokens["g"]): while len(tokens["l"]) > len(tokens["g"]):
tokens["g"] += empty["g"] tokens["g"] += empty["g"]
return (clip.encode_from_tokens_scheduled(tokens, add_dict={"width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}), ) return io.NodeOutput(clip.encode_from_tokens_scheduled(tokens, add_dict={"width": width, "height": height, "crop_w": crop_w, "crop_h": crop_h, "target_width": target_width, "target_height": target_height}))
NODE_CLASS_MAPPINGS = {
"CLIPTextEncodeSDXLRefiner": CLIPTextEncodeSDXLRefiner, NODES_LIST: list[type[io.ComfyNode]] = [
"CLIPTextEncodeSDXL": CLIPTextEncodeSDXL, CLIPTextEncodeSDXLRefiner,
} CLIPTextEncodeSDXL,
]
class ClipSdxlExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return NODES_LIST
async def comfy_entrypoint() -> ClipSdxlExtension:
return ClipSdxlExtension()

View File

@ -1,7 +1,14 @@
import torch from __future__ import annotations
import comfy.utils
from enum import Enum from enum import Enum
import torch
from typing_extensions import override
import comfy.utils
from comfy_api.latest import ComfyExtension, io
def resize_mask(mask, shape): def resize_mask(mask, shape):
return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1) return torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(shape[0], shape[1]), mode="bilinear").squeeze(1)
@ -101,24 +108,28 @@ def porter_duff_composite(src_image: torch.Tensor, src_alpha: torch.Tensor, dst_
return out_image, out_alpha return out_image, out_alpha
class PorterDuffImageComposite: class PorterDuffImageComposite(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return io.Schema(
"required": { node_id="PorterDuffImageComposite",
"source": ("IMAGE",), display_name="Porter-Duff Image Composite",
"source_alpha": ("MASK",), category="mask/compositing",
"destination": ("IMAGE",), inputs=[
"destination_alpha": ("MASK",), io.Image.Input("source"),
"mode": ([mode.name for mode in PorterDuffMode], {"default": PorterDuffMode.DST.name}), io.Mask.Input("source_alpha"),
}, io.Image.Input("destination"),
} io.Mask.Input("destination_alpha"),
io.Combo.Input("mode", options=[mode.name for mode in PorterDuffMode], default=PorterDuffMode.DST.name),
],
outputs=[
io.Image.Output(),
io.Mask.Output(),
],
)
RETURN_TYPES = ("IMAGE", "MASK") @classmethod
FUNCTION = "composite" def execute(cls, source: torch.Tensor, source_alpha: torch.Tensor, destination: torch.Tensor, destination_alpha: torch.Tensor, mode) -> io.NodeOutput:
CATEGORY = "mask/compositing"
def composite(self, source: torch.Tensor, source_alpha: torch.Tensor, destination: torch.Tensor, destination_alpha: torch.Tensor, mode):
batch_size = min(len(source), len(source_alpha), len(destination), len(destination_alpha)) batch_size = min(len(source), len(source_alpha), len(destination), len(destination_alpha))
out_images = [] out_images = []
out_alphas = [] out_alphas = []
@ -151,44 +162,47 @@ class PorterDuffImageComposite:
out_alphas.append(out_alpha.squeeze(2)) out_alphas.append(out_alpha.squeeze(2))
result = (torch.stack(out_images), torch.stack(out_alphas)) result = (torch.stack(out_images), torch.stack(out_alphas))
return result return io.NodeOutput(result)
class SplitImageWithAlpha(io.ComfyNode):
class SplitImageWithAlpha:
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return io.Schema(
"required": { node_id="SplitImageWithAlpha",
"image": ("IMAGE",), display_name="Split Image with Alpha",
} category="mask/compositing",
} inputs=[
io.Image.Input("image"),
],
outputs=[
io.Image.Output(),
io.Mask.Output(),
],
)
CATEGORY = "mask/compositing" @classmethod
RETURN_TYPES = ("IMAGE", "MASK") def execute(cls, image: torch.Tensor) -> io.NodeOutput:
FUNCTION = "split_image_with_alpha"
def split_image_with_alpha(self, image: torch.Tensor):
out_images = [i[:,:,:3] for i in image] out_images = [i[:,:,:3] for i in image]
out_alphas = [i[:,:,3] if i.shape[2] > 3 else torch.ones_like(i[:,:,0]) for i in image] out_alphas = [i[:,:,3] if i.shape[2] > 3 else torch.ones_like(i[:,:,0]) for i in image]
result = (torch.stack(out_images), 1.0 - torch.stack(out_alphas)) result = (torch.stack(out_images), 1.0 - torch.stack(out_alphas))
return result return io.NodeOutput(result)
class JoinImageWithAlpha(io.ComfyNode):
class JoinImageWithAlpha:
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls):
return { return io.Schema(
"required": { node_id="JoinImageWithAlpha",
"image": ("IMAGE",), display_name="Join Image with Alpha",
"alpha": ("MASK",), category="mask/compositing",
} inputs=[
} io.Image.Input("image"),
io.Mask.Input("alpha"),
],
outputs=[io.Image.Output()],
)
CATEGORY = "mask/compositing" @classmethod
RETURN_TYPES = ("IMAGE",) def execute(cls, image: torch.Tensor, alpha: torch.Tensor) -> io.NodeOutput:
FUNCTION = "join_image_with_alpha"
def join_image_with_alpha(self, image: torch.Tensor, alpha: torch.Tensor):
batch_size = min(len(image), len(alpha)) batch_size = min(len(image), len(alpha))
out_images = [] out_images = []
@ -196,19 +210,19 @@ class JoinImageWithAlpha:
for i in range(batch_size): for i in range(batch_size):
out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2)) out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
result = (torch.stack(out_images),) return io.NodeOutput(torch.stack(out_images))
return result
NODE_CLASS_MAPPINGS = { NODES_LIST: list[type[io.ComfyNode]] = [
"PorterDuffImageComposite": PorterDuffImageComposite, PorterDuffImageComposite,
"SplitImageWithAlpha": SplitImageWithAlpha, SplitImageWithAlpha,
"JoinImageWithAlpha": JoinImageWithAlpha, JoinImageWithAlpha,
} ]
class CompositingExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return NODES_LIST
NODE_DISPLAY_NAME_MAPPINGS = { async def comfy_entrypoint() -> CompositingExtension:
"PorterDuffImageComposite": "Porter-Duff Image Composite", return CompositingExtension()
"SplitImageWithAlpha": "Split Image with Alpha",
"JoinImageWithAlpha": "Join Image with Alpha",
}

View File

@ -1,15 +1,26 @@
from __future__ import annotations
from typing_extensions import override
from comfy_api.latest import ComfyExtension, io
class CLIPTextEncodeControlnet: class CLIPTextEncodeControlnet(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls) -> io.Schema:
return {"required": {"clip": ("CLIP", ), "conditioning": ("CONDITIONING", ), "text": ("STRING", {"multiline": True, "dynamicPrompts": True})}} return io.Schema(
RETURN_TYPES = ("CONDITIONING",) node_id="CLIPTextEncodeControlnet",
FUNCTION = "encode" category="_for_testing/conditioning",
inputs=[
io.Clip.Input("clip"),
io.Conditioning.Input("conditioning"),
io.String.Input("text", multiline=True, dynamic_prompts=True),
],
outputs=[io.Conditioning.Output()],
)
CATEGORY = "_for_testing/conditioning" @classmethod
def execute(cls, clip, conditioning, text) -> io.NodeOutput:
def encode(self, clip, conditioning, text):
tokens = clip.tokenize(text) tokens = clip.tokenize(text)
cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True) cond, pooled = clip.encode_from_tokens(tokens, return_pooled=True)
c = [] c = []
@ -18,32 +29,41 @@ class CLIPTextEncodeControlnet:
n[1]['cross_attn_controlnet'] = cond n[1]['cross_attn_controlnet'] = cond
n[1]['pooled_output_controlnet'] = pooled n[1]['pooled_output_controlnet'] = pooled
c.append(n) c.append(n)
return (c, ) return io.NodeOutput(c)
class T5TokenizerOptions: class T5TokenizerOptions(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls) -> io.Schema:
return { return io.Schema(
"required": { node_id="T5TokenizerOptions",
"clip": ("CLIP", ), category="_for_testing/conditioning",
"min_padding": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1}), inputs=[
"min_length": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1}), io.Clip.Input("clip"),
} io.Int.Input("min_padding", default=0, min=0, max=10000, step=1),
} io.Int.Input("min_length", default=0, min=0, max=10000, step=1),
],
outputs=[io.Clip.Output()],
)
CATEGORY = "_for_testing/conditioning" @classmethod
RETURN_TYPES = ("CLIP",) def execute(cls, clip, min_padding, min_length) -> io.NodeOutput:
FUNCTION = "set_options"
def set_options(self, clip, min_padding, min_length):
clip = clip.clone() clip = clip.clone()
for t5_type in ["t5xxl", "pile_t5xl", "t5base", "mt5xl", "umt5xxl"]: for t5_type in ["t5xxl", "pile_t5xl", "t5base", "mt5xl", "umt5xxl"]:
clip.set_tokenizer_option("{}_min_padding".format(t5_type), min_padding) clip.set_tokenizer_option("{}_min_padding".format(t5_type), min_padding)
clip.set_tokenizer_option("{}_min_length".format(t5_type), min_length) clip.set_tokenizer_option("{}_min_length".format(t5_type), min_length)
return (clip, ) return io.NodeOutput(clip)
NODE_CLASS_MAPPINGS = {
"CLIPTextEncodeControlnet": CLIPTextEncodeControlnet, NODES_LIST: list[type[io.ComfyNode]] = [
"T5TokenizerOptions": T5TokenizerOptions, CLIPTextEncodeControlnet,
} T5TokenizerOptions,
]
class CondExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return NODES_LIST
async def comfy_entrypoint() -> CondExtension:
return CondExtension()

View File

@ -1,25 +1,34 @@
import nodes from __future__ import annotations
import torch import torch
from typing_extensions import override
import comfy.latent_formats
import comfy.model_management import comfy.model_management
import comfy.utils import comfy.utils
import comfy.latent_formats import nodes
from comfy_api.latest import ComfyExtension, io
class EmptyCosmosLatentVideo: class EmptyCosmosLatentVideo(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls) -> io.Schema:
return {"required": { "width": ("INT", {"default": 1280, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), return io.Schema(
"height": ("INT", {"default": 704, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), node_id="EmptyCosmosLatentVideo",
"length": ("INT", {"default": 121, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 8}), category="latent/video",
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096})}} inputs=[
RETURN_TYPES = ("LATENT",) io.Int.Input("width", default=1280, min=16, max=nodes.MAX_RESOLUTION, step=16),
FUNCTION = "generate" io.Int.Input("height", default=704, min=16, max=nodes.MAX_RESOLUTION, step=16),
io.Int.Input("length", default=121, min=1, max=nodes.MAX_RESOLUTION, step=8),
io.Int.Input("batch_size", default=1, min=1, max=4096),
],
outputs=[io.Latent.Output()],
)
CATEGORY = "latent/video" @classmethod
def execute(cls, width, height, length, batch_size=1) -> io.NodeOutput:
def generate(self, width, height, length, batch_size=1):
latent = torch.zeros([batch_size, 16, ((length - 1) // 8) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device()) latent = torch.zeros([batch_size, 16, ((length - 1) // 8) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
return ({"samples": latent}, ) return io.NodeOutput({"samples": latent})
def vae_encode_with_padding(vae, image, width, height, length, padding=0): def vae_encode_with_padding(vae, image, width, height, length, padding=0):
@ -33,31 +42,31 @@ def vae_encode_with_padding(vae, image, width, height, length, padding=0):
return latent_temp[:, :, :latent_len] return latent_temp[:, :, :latent_len]
class CosmosImageToVideoLatent: class CosmosImageToVideoLatent(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls) -> io.Schema:
return {"required": {"vae": ("VAE", ), return io.Schema(
"width": ("INT", {"default": 1280, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), node_id="CosmosImageToVideoLatent",
"height": ("INT", {"default": 704, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), category="conditioning/inpaint",
"length": ("INT", {"default": 121, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 8}), inputs=[
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), io.Vae.Input("vae"),
}, io.Int.Input("width", default=1280, min=16, max=nodes.MAX_RESOLUTION, step=16),
"optional": {"start_image": ("IMAGE", ), io.Int.Input("height", default=704, min=16, max=nodes.MAX_RESOLUTION, step=16),
"end_image": ("IMAGE", ), io.Int.Input("length", default=121, min=1, max=nodes.MAX_RESOLUTION, step=8),
}} io.Int.Input("batch_size", default=1, min=1, max=4096),
io.Image.Input("start_image", optional=True),
io.Image.Input("end_image", optional=True),
],
outputs=[io.Latent.Output()],
)
@classmethod
RETURN_TYPES = ("LATENT",) def execute(cls, vae, width, height, length, batch_size, start_image=None, end_image=None) -> io.NodeOutput:
FUNCTION = "encode"
CATEGORY = "conditioning/inpaint"
def encode(self, vae, width, height, length, batch_size, start_image=None, end_image=None):
latent = torch.zeros([1, 16, ((length - 1) // 8) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device()) latent = torch.zeros([1, 16, ((length - 1) // 8) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
if start_image is None and end_image is None: if start_image is None and end_image is None:
out_latent = {} out_latent = {}
out_latent["samples"] = latent out_latent["samples"] = latent
return (out_latent,) return io.NodeOutput(out_latent)
mask = torch.ones([latent.shape[0], 1, ((length - 1) // 8) + 1, latent.shape[-2], latent.shape[-1]], device=comfy.model_management.intermediate_device()) mask = torch.ones([latent.shape[0], 1, ((length - 1) // 8) + 1, latent.shape[-2], latent.shape[-1]], device=comfy.model_management.intermediate_device())
@ -74,33 +83,34 @@ class CosmosImageToVideoLatent:
out_latent = {} out_latent = {}
out_latent["samples"] = latent.repeat((batch_size, ) + (1,) * (latent.ndim - 1)) out_latent["samples"] = latent.repeat((batch_size, ) + (1,) * (latent.ndim - 1))
out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1)) out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1))
return (out_latent,) return io.NodeOutput(out_latent)
class CosmosPredict2ImageToVideoLatent:
class CosmosPredict2ImageToVideoLatent(io.ComfyNode):
@classmethod @classmethod
def INPUT_TYPES(s): def define_schema(cls) -> io.Schema:
return {"required": {"vae": ("VAE", ), return io.Schema(
"width": ("INT", {"default": 848, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), node_id="CosmosPredict2ImageToVideoLatent",
"height": ("INT", {"default": 480, "min": 16, "max": nodes.MAX_RESOLUTION, "step": 16}), category="conditioning/inpaint",
"length": ("INT", {"default": 93, "min": 1, "max": nodes.MAX_RESOLUTION, "step": 4}), inputs=[
"batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), io.Vae.Input("vae"),
}, io.Int.Input("width", default=848, min=16, max=nodes.MAX_RESOLUTION, step=16),
"optional": {"start_image": ("IMAGE", ), io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16),
"end_image": ("IMAGE", ), io.Int.Input("length", default=93, min=1, max=nodes.MAX_RESOLUTION, step=4),
}} io.Int.Input("batch_size", default=1, min=1, max=4096),
io.Image.Input("start_image", optional=True),
io.Image.Input("end_image", optional=True),
],
outputs=[io.Latent.Output()],
)
@classmethod
RETURN_TYPES = ("LATENT",) def execute(cls, vae, width, height, length, batch_size, start_image=None, end_image=None) -> io.NodeOutput:
FUNCTION = "encode"
CATEGORY = "conditioning/inpaint"
def encode(self, vae, width, height, length, batch_size, start_image=None, end_image=None):
latent = torch.zeros([1, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device()) latent = torch.zeros([1, 16, ((length - 1) // 4) + 1, height // 8, width // 8], device=comfy.model_management.intermediate_device())
if start_image is None and end_image is None: if start_image is None and end_image is None:
out_latent = {} out_latent = {}
out_latent["samples"] = latent out_latent["samples"] = latent
return (out_latent,) return io.NodeOutput(out_latent)
mask = torch.ones([latent.shape[0], 1, ((length - 1) // 4) + 1, latent.shape[-2], latent.shape[-1]], device=comfy.model_management.intermediate_device()) mask = torch.ones([latent.shape[0], 1, ((length - 1) // 4) + 1, latent.shape[-2], latent.shape[-1]], device=comfy.model_management.intermediate_device())
@ -119,10 +129,19 @@ class CosmosPredict2ImageToVideoLatent:
latent = latent_format.process_out(latent) * mask + latent * (1.0 - mask) latent = latent_format.process_out(latent) * mask + latent * (1.0 - mask)
out_latent["samples"] = latent.repeat((batch_size, ) + (1,) * (latent.ndim - 1)) out_latent["samples"] = latent.repeat((batch_size, ) + (1,) * (latent.ndim - 1))
out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1)) out_latent["noise_mask"] = mask.repeat((batch_size, ) + (1,) * (mask.ndim - 1))
return (out_latent,) return io.NodeOutput(out_latent)
NODE_CLASS_MAPPINGS = {
"EmptyCosmosLatentVideo": EmptyCosmosLatentVideo, NODES_LIST: list[type[io.ComfyNode]] = [
"CosmosImageToVideoLatent": CosmosImageToVideoLatent, EmptyCosmosLatentVideo,
"CosmosPredict2ImageToVideoLatent": CosmosPredict2ImageToVideoLatent, CosmosImageToVideoLatent,
} CosmosPredict2ImageToVideoLatent,
]
class CosmosExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[io.ComfyNode]]:
return NODES_LIST
async def comfy_entrypoint() -> CosmosExtension:
return CosmosExtension()

File diff suppressed because it is too large Load Diff