custom_nodes

.gitattributes

@@ -15,3 +15,12 @@ disaster_girl.png filter=lfs diff=lfs merge=lfs -text
 istambul.jpg filter=lfs diff=lfs merge=lfs -text
 mona.png filter=lfs diff=lfs merge=lfs -text
 natasha.png filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.ttf filter=lfs diff=lfs merge=lfs -text
+*.whl filter=lfs diff=lfs merge=lfs -text
+*.task filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.woff2 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -5,7 +5,7 @@ __pycache__/
 !/input/example.png
 /models/
 /temp/
-/custom_nodes/
+#/custom_nodes/
 !custom_nodes/example_node.py.example
 extra_model_paths.yaml
 /.vs

custom_nodes/ComfyUI-tbox/.gitignore

@@ -0,0 +1,3 @@
+.idea
+**/__pycache__
+

custom_nodes/ComfyUI-tbox/README.md

@@ -0,0 +1,5 @@
+## Image Node:
+
+## Video Node
+
+## ControlNet PreProcessor

custom_nodes/ComfyUI-tbox/__init__.py

@@ -0,0 +1,82 @@
+import sys
+from pathlib import Path
+from .utils import here
+import platform
+
+sys.path.insert(0, str(Path(here, "src").resolve()))
+
+from .nodes.image.load_node import LoadImageNode
+from .nodes.image.save_node import SaveImageNode
+from .nodes.image.save_node import SaveImagesNode
+from .nodes.image.size_node import ImageResizeNode
+from .nodes.image.size_node import ImageSizeNode
+from .nodes.image.size_node import ConstrainImageNode
+from .nodes.image.watermark_node import WatermarkNode
+from .nodes.mask.mask_node import MaskAddNode
+from .nodes.video.load_node import LoadVideoNode
+from .nodes.video.save_node import SaveVideoNode
+from .nodes.video.info_node import VideoInfoNode
+from .nodes.video.batch_node import BatchManagerNode
+from .nodes.preprocessor.canny_node import Canny_Preprocessor
+from .nodes.preprocessor.lineart_node import LineArt_Preprocessor
+from .nodes.preprocessor.lineart_node import Lineart_Standard_Preprocessor
+from .nodes.preprocessor.midas_node import MIDAS_Depth_Map_Preprocessor
+from .nodes.preprocessor.dwpose_node import DWPose_Preprocessor, AnimalPose_Preprocessor
+from .nodes.preprocessor.densepose_node import DensePose_Preprocessor
+from .nodes.face.face_enhance_node import GFPGANNode
+from .nodes.other.vram_node import PurgeVRAMNode
+
+NODE_CLASS_MAPPINGS = {
+    "PurgeVRAMNode": PurgeVRAMNode,
+    "GFPGANNode": GFPGANNode,
+    "MaskAddNode": MaskAddNode,
+    "ImageLoader": LoadImageNode,
+    "ImageSaver": SaveImageNode,
+    "ImagesSaver": SaveImagesNode,
+    "ImageResize": ImageResizeNode,
+    "ImageSize": ImageSizeNode,
+    "WatermarkNode": WatermarkNode,
+    "VideoLoader": LoadVideoNode,
+    "VideoSaver": SaveVideoNode,
+    "VideoInfo": VideoInfoNode,
+    "BatchManager": BatchManagerNode,
+    "ConstrainImageNode": ConstrainImageNode,
+    "DensePosePreprocessor": DensePose_Preprocessor,
+    "DWPosePreprocessor": DWPose_Preprocessor,
+    "AnimalPosePreprocessor": AnimalPose_Preprocessor,
+    "MiDaSDepthPreprocessor": MIDAS_Depth_Map_Preprocessor,
+    "CannyPreprocessor": Canny_Preprocessor,
+    "LineArtPreprocessor": LineArt_Preprocessor,
+    "LineartStandardPreprocessor": Lineart_Standard_Preprocessor,
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "PurgeVRAMNode":"PurgeVRAMNode",
+    "GFPGANNode": "GFPGANNode",
+    "MaskAddNode": "MaskAddNode",
+    "ImageLoader": "Image Load",
+    "ImageSaver": "Image Save",
+    "ImagesSaver": "Image List Save",
+    "ImageResize": "Image Resize",
+    "ImageSize": "Image Size",
+    "WatermarkNode": "Watermark",
+    "VideoLoader": "Video Load",
+    "VideoSaver": "Video Save",
+    "VideoInfo": "Video Info", 
+    "BatchManager": "Batch Manager",
+    "ConstrainImageNode": "Image Constrain",
+    "DensePosePreprocessor": "DensePose Estimator",
+    "DWPosePreprocessor": "DWPose Estimator",
+    "AnimalPosePreprocessor": "AnimalPose Estimator",
+    "MiDaSDepthPreprocessor": "MiDaS Depth Estimator",
+    "CannyPreprocessor": "Canny Edge Estimator",
+    "LineArtPreprocessor": "Realistic Lineart",
+    "LineartStandardPreprocessor": "Standard Lineart",
+}
+
+
+if platform.system() == "Darwin":
+    WEB_DIRECTORY = "./web"
+    __all__ = ["WEB_DIRECTORY", "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
+else:
+    __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

custom_nodes/ComfyUI-tbox/config.yaml

@@ -0,0 +1,4 @@
+annotator_ckpts_path: "../../models/annotator"
+custom_temp_path: "../../temp"
+USE_SYMLINKS: False
+EP_list: ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CPUExecutionProvider"]

custom_nodes/ComfyUI-tbox/nodes/face/__init__.py

@@ -0,0 +1,8 @@
+
+import os
+import folder_paths
+
+
+model_path = os.path.join(folder_paths.models_dir, "facefusion")
+folder_paths.add_model_folder_path('facefusion', model_path)
+

custom_nodes/ComfyUI-tbox/nodes/face/face_enhance_node.py

@@ -0,0 +1,81 @@
+import os
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+from facefusion.gfpgan_onnx import GFPGANOnnx
+from facefusion.yoloface_onnx import YoloFaceOnnx
+from facefusion.affine import create_box_mask, warp_face_by_landmark, paste_back
+
+import folder_paths
+from ..utils import tensor2pil, pil2tensor
+    
+# class GFPGANProvider:
+#     @classmethod
+#     def INPUT_TYPES(s):
+#         return {
+#             "required": {
+#                 "model_name": ("IMAGE", ["gfpgan_1.4.onnx"]),
+#             },
+#         }
+    
+#     RETURN_TYPES = ("GFPGAN_MODEL",)
+#     RETURN_NAMES = ("model",)
+#     FUNCTION = "load_model"
+#     CATEGORY = "tbox/facefusion"
+
+#     def load_model(self, model_name):
+#         return (model_name,)
+
+
+class GFPGANNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "model_name": (['gfpgan_1.3', 'gfpgan_1.4'], {"default": 'gfpgan_1.4'}),
+                "device": (['CPU', 'CUDA', 'CoreML', 'ROCM'], {"default": 'CPU'}),
+                "weight": ("FLOAT", {"default": 0.8}),
+            }
+        }
+
+    RETURN_TYPES = ("IMAGE", )
+    FUNCTION = "process"
+    CATEGORY = "tbox/FaceFusion"
+
+    def process(self, images, model_name, device='CPU', weight=0.8):
+        providers = ['CPUExecutionProvider']
+        if device== 'CUDA':
+            providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
+        elif device == 'CoreML':
+            providers = ['CoreMLExecutionProvider', 'CPUExecutionProvider']
+        elif device == 'ROCM':
+            providers = ['ROCMExecutionProvider', 'CPUExecutionProvider']
+        
+        gfpgan_path = folder_paths.get_full_path("facefusion", f'{model_name}.onnx')
+        yolo_path = folder_paths.get_full_path("facefusion", 'yoloface_8n.onnx')
+        
+        detector = YoloFaceOnnx(model_path=yolo_path, providers=providers)
+        enhancer = GFPGANOnnx(model_path=gfpgan_path, providers=providers)
+        
+        image_list = []
+        for i, img in enumerate(images):
+            pil = tensor2pil(img)
+            image = np.ascontiguousarray(pil)
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+            output = image
+            face_list = detector.detect(image=image, conf=0.7)
+            for index, face in enumerate(face_list):
+                cropped, affine_matrix = warp_face_by_landmark(image, face.landmarks, enhancer.input_size)
+                box_mask = create_box_mask(enhancer.input_size, 0.3, (0,0,0,0))
+                crop_mask = np.minimum.reduce([box_mask]).clip(0, 1)
+                result = enhancer.run(cropped)
+                output = paste_back(output, result, crop_mask, affine_matrix)
+            image = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+            pil = Image.fromarray(image)
+            image_list.append(pil2tensor(pil))
+        image_list = torch.stack([tensor.squeeze() for tensor in image_list])
+        return (image_list,)
+
+    

custom_nodes/ComfyUI-tbox/nodes/image/load_node.py

@@ -0,0 +1,81 @@
+
+import os
+import cv2
+import torch
+import requests
+import itertools
+import folder_paths
+import psutil
+import numpy as np
+from comfy.utils import common_upscale
+from io import BytesIO
+from PIL import Image, ImageSequence, ImageOps
+
+
+
+def pil2tensor(img):
+    output_images = []
+    output_masks = []
+    for i in ImageSequence.Iterator(img):
+        i = ImageOps.exif_transpose(i)
+        if i.mode == 'I':
+            i = i.point(lambda i: i * (1 / 255))
+        image = i.convert("RGB")
+        image = np.array(image).astype(np.float32) / 255.0
+        image = torch.from_numpy(image)[None,]
+        if 'A' in i.getbands():
+            mask = np.array(i.getchannel('A')).astype(np.float32) / 255.0
+            mask = 1. - torch.from_numpy(mask)
+        else:
+            mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
+        output_images.append(image)
+        output_masks.append(mask.unsqueeze(0))
+
+    if len(output_images) > 1:
+        output_image = torch.cat(output_images, dim=0)
+        output_mask = torch.cat(output_masks, dim=0)
+    else:
+        output_image = output_images[0]
+        output_mask = output_masks[0]
+
+    return (output_image, output_mask)
+
+
+def load_image(image_source):
+    if image_source.startswith('http'):
+        print(image_source)
+        response = requests.get(image_source)
+        img = Image.open(BytesIO(response.content))
+        file_name = image_source.split('/')[-1]
+    else:
+        img = Image.open(image_source)
+        file_name = os.path.basename(image_source)
+    return img, file_name
+
+
+class LoadImageNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "path": ("STRING", {"multiline": True, "dynamicPrompts": False})
+            }
+        }
+
+
+    RETURN_TYPES = ("IMAGE", "MASK")
+    FUNCTION = "load_image"
+    CATEGORY = "tbox/Image"
+
+    def load_image(self, path):
+        filepaht = path.split('\n')[0]
+        img, name = load_image(filepaht)
+        img_out, mask_out = pil2tensor(img)
+        return (img_out, mask_out)
+
+
+
+if __name__ == "__main__":
+    img, name = load_image("https://creativestorage.blob.core.chinacloudapi.cn/test/bird.png")
+    img_out, mask_out = pil2tensor(img)
+

custom_nodes/ComfyUI-tbox/nodes/image/save_node.py

@@ -0,0 +1,79 @@
+import os
+import time
+import numpy as np
+from PIL import Image, ImageSequence, ImageOps
+
+#from load_node import load_image, pil2tensor
+
+def save_image(img, filepath, format, quality):
+    try:
+        if format in ["jpg", "jpeg"]:
+            img.convert("RGB").save(filepath, format="JPEG", quality=quality, subsampling=0)
+        elif format == "webp":
+            img.save(filepath, format="WEBP", quality=quality, method=6)
+        elif format == "bmp":
+            img.save(filepath, format="BMP")
+        else:
+            img.save(filepath, format="PNG", optimize=True)
+    except Exception as e:
+        print(f"Error saving {filepath}: {str(e)}")
+        
+class SaveImageNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "path": ("STRING", {"multiline": True, "dynamicPrompts": False}),
+                "quality": ([100, 95, 90, 85, 80, 75, 70, 60, 50], {"default": 100}),
+            }
+        }
+    RETURN_TYPES = ()
+    FUNCTION = "save_image"
+    CATEGORY = "tbox/Image"
+    OUTPUT_NODE = True
+    
+    def save_image(self, images, path, quality):
+        filepaht = path.split('\n')[0]
+        format = os.path.splitext(filepaht)[1][1:]
+        image = images[0] 
+        img = Image.fromarray((255. * image.cpu().numpy()).astype(np.uint8))
+        save_image(img, filepaht, format, quality)
+        return {}
+        
+class SaveImagesNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "path": ("STRING", {"multiline": False, "dynamicPrompts": False}),
+                "prefix": ("STRING", {"default": "image"}),
+                "format": (["PNG", "JPG", "WEBP", "BMP"],),
+                "quality": ([100, 95, 90, 85, 80, 75, 70, 60, 50], {"default": 100}),
+            }
+        }
+    RETURN_TYPES = ()
+    FUNCTION = "save_image"
+    CATEGORY = "tbox/Image"
+    OUTPUT_NODE = True
+    
+    def save_image(self, images, path, prefix, format, quality):
+        format = format.lower()            
+        for i, image in enumerate(images):
+            img = Image.fromarray((255. * image.cpu().numpy()).astype(np.uint8))
+            filepath = self.generate_filename(path, prefix, i, format)
+            save_image(img, filepath, format, quality)
+        return {}
+
+    def IS_CHANGED(s, images):
+        return time.time()
+    
+    def generate_filename(self, save_dir, prefix, index, format):
+        base_filename = f"{prefix}_{index+1}.{format}"
+        filename = os.path.join(save_dir, base_filename)
+        counter = 1
+        while os.path.exists(filename):
+            filename = os.path.join(save_dir, f"{prefix}_{index+1}_{counter}.{format}")
+            counter += 1
+        return filename

custom_nodes/ComfyUI-tbox/nodes/image/size_node.py

@@ -0,0 +1,121 @@
+
+
+import torch
+import comfy.utils
+import numpy as np
+from PIL import Image, ImageSequence, ImageOps
+
+class ConstrainImageNode:
+    """
+    A node that constrains an image to a maximum and minimum size while maintaining aspect ratio.
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "max_width": ("INT", {"default": 1024, "min": 0}),
+                "max_height": ("INT", {"default": 1024, "min": 0}),
+                "min_width": ("INT", {"default": 0, "min": 0}),
+                "min_height": ("INT", {"default": 0, "min": 0}),
+                "crop_if_required": (["yes", "no"], {"default": "no"}),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "constrain_image"
+    CATEGORY = "tbox/Image"
+    OUTPUT_IS_LIST = (True,)
+
+    def constrain_image(self, images, max_width, max_height, min_width, min_height, crop_if_required):
+        crop_if_required = crop_if_required == "yes"
+        results = []
+        for image in images:
+            i = 255. * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)).convert("RGB")
+
+            current_width, current_height = img.size
+            aspect_ratio = current_width / current_height
+
+            constrained_width = max(min(current_width, min_width), max_width)
+            constrained_height = max(min(current_height, min_height), max_height)
+
+            if constrained_width / constrained_height > aspect_ratio:
+                constrained_width = max(int(constrained_height * aspect_ratio), min_width)
+                if crop_if_required:
+                    constrained_height = int(current_height / (current_width / constrained_width))
+            else:
+                constrained_height = max(int(constrained_width / aspect_ratio), min_height)
+                if crop_if_required:
+                    constrained_width = int(current_width / (current_height / constrained_height))
+
+            resized_image = img.resize((constrained_width, constrained_height), Image.LANCZOS)
+
+            if crop_if_required and (constrained_width > max_width or constrained_height > max_height):
+                left = max((constrained_width - max_width) // 2, 0)
+                top = max((constrained_height - max_height) // 2, 0)
+                right = min(constrained_width, max_width) + left
+                bottom = min(constrained_height, max_height) + top
+                resized_image = resized_image.crop((left, top, right, bottom))
+
+            resized_image = np.array(resized_image).astype(np.float32) / 255.0
+            resized_image = torch.from_numpy(resized_image)[None,]
+            results.append(resized_image)
+                
+        return (results,)
+    
+# https://github.com/bronkula/comfyui-fitsize
+class ImageSizeNode: 
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+            }
+        }
+    
+    RETURN_TYPES = ("INT", "INT", "INT")
+    RETURN_NAMES = ("width", "height", "count")
+    FUNCTION = "get_size"
+    CATEGORY = "tbox/Image"
+    def get_size(self, image):
+        print(f'shape of image:{image.shape}')
+        return (image.shape[2], image.shape[1], image[0])
+
+
+
+class ImageResizeNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE", ),
+                "method": (["nearest", "bilinear", "bicubic", "area", "nearest-exact", "lanczos"],),
+            },
+            "optional": {
+                "width": ("INT,FLOAT", { "default": 0.0, "step": 0.1 }),
+                "height": ("INT,FLOAT", { "default": 0.0, "step": 0.1 }),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+
+    FUNCTION = "resize"
+    CATEGORY = "tbox/Image"
+    
+    def resize(self, image, method, width, height):
+        print(f'shape of image:{image.shape}, resolution:{width}x{height} type: {type(width)}, {type(height)}')
+        if width == 0 and height == 0:
+            s = image
+        else:
+            samples = image.movedim(-1,1)
+            if width == 0:
+                width = max(1, round(samples.shape[3] * height / samples.shape[2]))
+            elif height == 0:
+                height = max(1, round(samples.shape[2] * width / samples.shape[3]))
+
+            s = comfy.utils.common_upscale(samples, width, height, method, True)
+            s = s.movedim(1,-1)
+        return (s,)
+    

custom_nodes/ComfyUI-tbox/nodes/image/watermark_node.py

@@ -0,0 +1,58 @@
+
+import torch
+import numpy as np
+from PIL import Image, ImageSequence, ImageOps
+from ..utils import tensor2pil, pil2tensor
+
+PADDING = 4
+
+class WatermarkNode:
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "logo_list": ("IMAGE",),
+            },
+            "optional": {
+                "logo_mask": ("MASK",),
+                "enabled": ("BOOLEAN", {"default": True}),}
+        }
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "watermark"
+    CATEGORY = "tbox/Image"
+    
+    def watermark(self, images, logo_list, logo_mask, enabled):
+        outputs = []
+        if enabled == False: 
+            return(images,)
+        print(f'logo shape: {logo_list.shape}')
+        print(f'images shape: {images.shape}')
+        logo = tensor2pil(logo_list[0]) 
+        if logo_mask is not None:
+            logo_mask = tensor2pil(logo_mask)
+        for i, image in enumerate(images):
+            img = tensor2pil(image) #Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+            dst = self.add_watermark2(img, logo, logo_mask, 85)
+            result = pil2tensor(dst)
+            outputs.append(result)
+        base_image = torch.stack([tensor.squeeze() for tensor in outputs])
+        return (base_image,)
+    
+    def add_watermark2(self, image, logo, logo_mask, opacity=None):
+        logo_width, logo_height = logo.size
+        image_width, image_height = image.size
+        if image_height <= logo_height + PADDING * 2 or image_width <= logo_width + PADDING * 2:
+            return image
+        y = image_height - logo_height - PADDING * 1
+        x = PADDING
+        logo = logo.convert('RGBA')
+        opacity = int(opacity / 100 * 255)
+        logo.putalpha(Image.new("L", logo.size, opacity))
+        if logo_mask is not None:
+            logo.putalpha(ImageOps.invert(logo_mask))
+
+        position = (x, y)
+        image.paste(logo, position, logo)
+        return image

custom_nodes/ComfyUI-tbox/nodes/mask/mask_node.py

@@ -0,0 +1,86 @@
+
+import torch
+import numpy as np
+
+class MaskSubNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mask": ("MASK",),  
+            },
+            "optional": {
+                "src1": ("MASK",),
+                "src2": ("MASK",),
+                "src3": ("MASK",),
+                "src4": ("MASK",),
+                "src5": ("MASK",),
+                "src6": ("MASK",),
+            }
+        }
+
+    CATEGORY = "mask"
+    RETURN_TYPES = ("MASK",)
+
+    FUNCTION = "sub"
+    CATEGORY = "tbox/Mask"
+    
+    def sub_mask(self, dst, src):
+        if src != None:
+            mask = src.reshape((-1, src.shape[-2], src.shape[-1]))
+            return dst - mask
+        return dst
+        
+    def add(self, mask, src1=None, src2=None, src3=None, src4=None, src5=None, src6=None):
+        print(f'mask shape: {mask.shape}')
+        output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone()
+        output[:, :, :] = self.sub_mask(output, src1)
+        output[:, :, :] = self.sub_mask(output, src2)
+        output[:, :, :] = self.sub_mask(output, src3)
+        output[:, :, :] = self.sub_mask(output, src4)
+        output[:, :, :] = self.sub_mask(output, src5)
+        output[:, :, :] = self.sub_mask(output, src6)
+        output = torch.clamp(output, 0.0, 1.0)
+        return (output, )  
+
+class MaskAddNode:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "mask": ("MASK",),  
+            },
+            "optional": {
+                "src1": ("MASK",),
+                "src2": ("MASK",),
+                "src3": ("MASK",),
+                "src4": ("MASK",),
+                "src5": ("MASK",),
+                "src6": ("MASK",),
+            }
+        }
+
+    CATEGORY = "mask"
+    RETURN_TYPES = ("MASK",)
+
+    FUNCTION = "add"
+    CATEGORY = "tbox/Mask"
+    
+    def add_mask(self, dst, src):
+        if src != None:
+            mask = src.reshape((-1, src.shape[-2], src.shape[-1]))
+            return dst + mask
+        return dst
+        
+    def add(self, mask, src1=None, src2=None, src3=None, src4=None, src5=None, src6=None):
+        print(f'mask shape: {mask.shape}')
+        output = mask.reshape((-1, mask.shape[-2], mask.shape[-1])).clone()
+        output[:, :, :] = self.add_mask(output, src1)
+        output[:, :, :] = self.add_mask(output, src2)
+        output[:, :, :] = self.add_mask(output, src3)
+        output[:, :, :] = self.add_mask(output, src4)
+        output[:, :, :] = self.add_mask(output, src5)
+        output[:, :, :] = self.add_mask(output, src6)
+        output = torch.clamp(output, 0.0, 1.0)
+        return (output, )  
+    

custom_nodes/ComfyUI-tbox/nodes/other/vram_node.py

@@ -0,0 +1,45 @@
+import gc
+import torch.cuda
+import comfy.model_management
+
+class AnyType(str):
+  """A special class that is always equal in not equal comparisons. Credit to pythongosssss"""
+  def __eq__(self, __value: object) -> bool:
+    return True
+  def __ne__(self, __value: object) -> bool:
+    return False
+
+any = AnyType("*")
+
+class PurgeVRAMNode:
+
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "anything": (any, {}),
+                "purge_cache": ("BOOLEAN", {"default": True}),
+                "purge_models": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+            }
+        }
+
+    RETURN_TYPES = ()
+    FUNCTION = "purge_vram"
+    CATEGORY = "tbox/other"
+    OUTPUT_NODE = True
+
+    def purge_vram(self, anything, purge_cache, purge_models):
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+        if purge_models:
+            comfy.model_management.unload_all_models()
+            comfy.model_management.soft_empty_cache()
+        return (None,)

custom_nodes/ComfyUI-tbox/nodes/preprocessor/canny_node.py

@@ -0,0 +1,22 @@
+from ..utils import common_annotator_call, create_node_input_types
+import comfy.model_management as model_management
+import nodes
+
+class Canny_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return create_node_input_types(
+            low_threshold=("INT", {"default": 100, "min": 0, "max": 255}),
+            high_threshold=("INT", {"default": 100, "min": 0, "max": 255}),
+            resolution=("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 64})
+        )
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+
+    CATEGORY = "tbox/ControlNet Preprocessors"
+
+    def execute(self, image, low_threshold=100, high_threshold=200, resolution=512, **kwargs):
+        from canny import CannyDetector
+
+        return (common_annotator_call(CannyDetector(), image, low_threshold=low_threshold, high_threshold=high_threshold, resolution=resolution), )

custom_nodes/ComfyUI-tbox/nodes/preprocessor/densepose_node.py

@@ -0,0 +1,22 @@
+from ..utils import common_annotator_call, create_node_input_types
+import comfy.model_management as model_management
+
+class DensePose_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return create_node_input_types(
+            model=(["densepose_r50_fpn_dl.torchscript", "densepose_r101_fpn_dl.torchscript"], {"default": "densepose_r50_fpn_dl.torchscript"}),
+            cmap=(["Viridis (MagicAnimate)", "Parula (CivitAI)"], {"default": "Viridis (MagicAnimate)"})
+        )
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+
+    CATEGORY = "tbox/ControlNet Preprocessors"
+
+    def execute(self, image, model, cmap, resolution=512):
+        from densepose import DenseposeDetector
+        model = DenseposeDetector \
+                    .from_pretrained(filename=model) \
+                    .to(model_management.get_torch_device())
+        return (common_annotator_call(model, image, cmap="viridis" if "Viridis" in cmap else "parula", resolution=resolution), )

custom_nodes/ComfyUI-tbox/nodes/preprocessor/dwpose_node.py

@@ -0,0 +1,158 @@
+from ..utils import common_annotator_call, create_node_input_types
+import comfy.model_management as model_management
+import numpy as np
+import warnings
+from dwpose import DwposeDetector, AnimalposeDetector
+import os
+import json
+
+
+DWPOSE_MODEL_NAME = "yzd-v/DWPose"
+#Trigger startup caching for onnxruntime
+GPU_PROVIDERS = ["CUDAExecutionProvider", "DirectMLExecutionProvider", "OpenVINOExecutionProvider", "ROCMExecutionProvider", "CoreMLExecutionProvider"]
+
+def check_ort_gpu():
+    try:
+        import onnxruntime as ort
+        for provider in GPU_PROVIDERS:
+            if provider in ort.get_available_providers():
+                return True
+        return False
+    except:
+        return False
+
+if not os.environ.get("DWPOSE_ONNXRT_CHECKED"):
+    if check_ort_gpu():
+        print("DWPose: Onnxruntime with acceleration providers detected")
+    else:
+        warnings.warn("DWPose: Onnxruntime not found or doesn't come with acceleration providers, switch to OpenCV with CPU device. DWPose might run very slowly")
+        os.environ['AUX_ORT_PROVIDERS'] = ''
+    os.environ["DWPOSE_ONNXRT_CHECKED"] = '1'
+
+class DWPose_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        input_types = create_node_input_types(
+            detect_hand=(["enable", "disable"], {"default": "enable"}),
+            detect_body=(["enable", "disable"], {"default": "enable"}),
+            detect_face=(["enable", "disable"], {"default": "enable"})
+        )
+        input_types["optional"] = {
+            **input_types["optional"],
+            "bbox_detector": (
+                ["yolox_l.torchscript.pt", "yolox_l.onnx", "yolo_nas_l_fp16.onnx", "yolo_nas_m_fp16.onnx", "yolo_nas_s_fp16.onnx"],
+                {"default": "yolox_l.onnx"}
+            ),
+            "pose_estimator": (["dw-ll_ucoco_384_bs5.torchscript.pt", "dw-ll_ucoco_384.onnx", "dw-ll_ucoco.onnx"], {"default": "dw-ll_ucoco_384_bs5.torchscript.pt"})
+        }
+        return input_types
+
+    RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
+    FUNCTION = "estimate_pose"
+
+    CATEGORY = "tbox/ControlNet Preprocessors"
+
+    def estimate_pose(self, image, detect_hand, detect_body, detect_face, resolution=512, bbox_detector="yolox_l.onnx", pose_estimator="dw-ll_ucoco_384.onnx", **kwargs):
+        if bbox_detector == "yolox_l.onnx":
+            yolo_repo = DWPOSE_MODEL_NAME
+        elif "yolox" in bbox_detector:
+            yolo_repo = "hr16/yolox-onnx"
+        elif "yolo_nas" in bbox_detector:
+            yolo_repo = "hr16/yolo-nas-fp16"
+        else:
+            raise NotImplementedError(f"Download mechanism for {bbox_detector}")
+
+        if pose_estimator == "dw-ll_ucoco_384.onnx":
+            pose_repo = DWPOSE_MODEL_NAME
+        elif pose_estimator.endswith(".onnx"):
+            pose_repo = "hr16/UnJIT-DWPose"
+        elif pose_estimator.endswith(".torchscript.pt"):
+            pose_repo = "hr16/DWPose-TorchScript-BatchSize5"
+        else:
+            raise NotImplementedError(f"Download mechanism for {pose_estimator}")
+
+        model = DwposeDetector.from_pretrained(
+            pose_repo,
+            yolo_repo,
+            det_filename=bbox_detector, pose_filename=pose_estimator,
+            torchscript_device=model_management.get_torch_device()
+        )
+        detect_hand = detect_hand == "enable"
+        detect_body = detect_body == "enable"
+        detect_face = detect_face == "enable"
+        self.openpose_dicts = []
+        def func(image, **kwargs):
+            pose_img, openpose_dict = model(image, **kwargs)
+            self.openpose_dicts.append(openpose_dict)
+            return pose_img
+
+        out = common_annotator_call(func, image, include_hand=detect_hand, include_face=detect_face, include_body=detect_body, image_and_json=True, resolution=resolution)
+        del model
+        return {
+            'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
+            "result": (out, self.openpose_dicts)
+        }
+
+class AnimalPose_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return create_node_input_types(
+            bbox_detector = (
+                ["yolox_l.torchscript.pt", "yolox_l.onnx", "yolo_nas_l_fp16.onnx", "yolo_nas_m_fp16.onnx", "yolo_nas_s_fp16.onnx"],
+                {"default": "yolox_l.torchscript.pt"}
+            ),
+            pose_estimator = (["rtmpose-m_ap10k_256_bs5.torchscript.pt", "rtmpose-m_ap10k_256.onnx"], {"default": "rtmpose-m_ap10k_256_bs5.torchscript.pt"})
+        )
+
+    RETURN_TYPES = ("IMAGE", "POSE_KEYPOINT")
+    FUNCTION = "estimate_pose"
+
+    CATEGORY = "tbox/ControlNet Preprocessors"
+
+    def estimate_pose(self, image, resolution=512, bbox_detector="yolox_l.onnx", pose_estimator="rtmpose-m_ap10k_256.onnx", **kwargs):
+        if bbox_detector == "yolox_l.onnx":
+            yolo_repo = DWPOSE_MODEL_NAME
+        elif "yolox" in bbox_detector:
+            yolo_repo = "hr16/yolox-onnx"
+        elif "yolo_nas" in bbox_detector:
+            yolo_repo = "hr16/yolo-nas-fp16"
+        else:
+            raise NotImplementedError(f"Download mechanism for {bbox_detector}")
+
+        if pose_estimator == "dw-ll_ucoco_384.onnx":
+            pose_repo = DWPOSE_MODEL_NAME
+        elif pose_estimator.endswith(".onnx"):
+            pose_repo = "hr16/UnJIT-DWPose"
+        elif pose_estimator.endswith(".torchscript.pt"):
+            pose_repo = "hr16/DWPose-TorchScript-BatchSize5"
+        else:
+            raise NotImplementedError(f"Download mechanism for {pose_estimator}")
+
+        model = AnimalposeDetector.from_pretrained(
+            pose_repo,
+            yolo_repo,
+            det_filename=bbox_detector, pose_filename=pose_estimator,
+            torchscript_device=model_management.get_torch_device()
+        )
+
+        self.openpose_dicts = []
+        def func(image, **kwargs):
+            pose_img, openpose_dict = model(image, **kwargs)
+            self.openpose_dicts.append(openpose_dict)
+            return pose_img
+
+        out = common_annotator_call(func, image, image_and_json=True, resolution=resolution)
+        del model
+        return {
+            'ui': { "openpose_json": [json.dumps(self.openpose_dicts, indent=4)] },
+            "result": (out, self.openpose_dicts)
+        }
+
+NODE_CLASS_MAPPINGS = {
+    "DWPreprocessor": DWPose_Preprocessor,
+    "AnimalPosePreprocessor": AnimalPose_Preprocessor
+}
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "DWPreprocessor": "DWPose Estimator",
+    "AnimalPosePreprocessor": "AnimalPose Estimator (AP10K)"
+}

custom_nodes/ComfyUI-tbox/nodes/preprocessor/lineart_node.py

@@ -0,0 +1,44 @@
+from ..utils import common_annotator_call, create_node_input_types
+import comfy.model_management as model_management
+import nodes
+
+class LineArt_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return create_node_input_types(
+            coarse=(["disable", "enable"], {"default": "enable"}),
+            resolution=("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 64})
+        )
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+
+    CATEGORY = "tbox/ControlNet Preprocessors"
+    
+
+    def execute(self, image, resolution=512, **kwargs):
+        from lineart import LineartDetector
+
+        model = LineartDetector.from_pretrained().to(model_management.get_torch_device())
+        out = common_annotator_call(model, image, resolution=resolution, coarse = kwargs["coarse"] == "enable")
+        del model
+        return (out, )
+    
+class Lineart_Standard_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return create_node_input_types(
+            guassian_sigma=("FLOAT", {"default":6.0, "max": 100.0}),
+            intensity_threshold=("INT", {"default": 8, "max": 16}),
+            resolution=("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 64})
+        )
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+
+    CATEGORY = "tbox/ControlNet Preprocessors"
+    
+
+    def execute(self, image, guassian_sigma=6, intensity_threshold=8, resolution=512, **kwargs):
+        from lineart import LineartStandardDetector
+        return (common_annotator_call(LineartStandardDetector(), image, guassian_sigma=guassian_sigma, intensity_threshold=intensity_threshold, resolution=resolution), )

custom_nodes/ComfyUI-tbox/nodes/preprocessor/midas_node.py

@@ -0,0 +1,25 @@
+from ..utils import common_annotator_call, create_node_input_types
+import comfy.model_management as model_management
+import numpy as np
+
+class MIDAS_Depth_Map_Preprocessor:
+    @classmethod
+    def INPUT_TYPES(s):
+        return create_node_input_types(
+            a =  ("FLOAT", {"default": np.pi * 2.0, "min": 0.0, "max": np.pi * 5.0, "step": 0.05}),
+            bg_threshold = ("FLOAT", {"default": 0.1, "min": 0, "max": 1, "step": 0.05})
+        )
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "execute"
+
+    CATEGORY = "tbox/ControlNet Preprocessors"
+
+    def execute(self, image, a, bg_threshold, resolution=512, **kwargs):
+        from midas import MidasDetector
+
+        # Ref: https://github.com/lllyasviel/ControlNet/blob/main/gradio_depth2image.py
+        model = MidasDetector.from_pretrained().to(model_management.get_torch_device())
+        out = common_annotator_call(model, image, resolution=resolution, a=a, bg_th=bg_threshold)
+        del model
+        return (out, )

custom_nodes/ComfyUI-tbox/nodes/utils.py

@@ -0,0 +1,165 @@
+import hashlib
+import os
+import torch
+import nodes
+import server
+import folder_paths
+import numpy as np
+from typing import Iterable
+from PIL import Image
+
+BIGMIN = -(2**53-1)
+BIGMAX = (2**53-1)
+
+DIMMAX = 8192
+
+
+def tensor_to_int(tensor, bits):
+    #TODO: investigate benefit of rounding by adding 0.5 before clip/cast
+    tensor = tensor.cpu().numpy() * (2**bits-1)
+    return np.clip(tensor, 0, (2**bits-1))
+def tensor_to_shorts(tensor):
+    return tensor_to_int(tensor, 16).astype(np.uint16)
+def tensor_to_bytes(tensor):
+    return tensor_to_int(tensor, 8).astype(np.uint8)
+def tensor2pil(x):
+    return Image.fromarray(np.clip(255. * x.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
+def pil2tensor(image: Image.Image) -> torch.Tensor:
+    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
+
+
+def is_url(url):
+    return url.split("://")[0] in ["http", "https"]
+
+def strip_path(path):
+    #This leaves whitespace inside quotes and only a single "
+    #thus ' ""test"' -> '"test'
+    #consider path.strip(string.whitespace+"\"")
+    #or weightier re.fullmatch("[\\s\"]*(.+?)[\\s\"]*", path).group(1)
+    path = path.strip()
+    if path.startswith("\""):
+        path = path[1:]
+    if path.endswith("\""):
+        path = path[:-1]
+    return path
+def hash_path(path):
+    if path is None:
+        return "input"
+    if is_url(path):
+        return "url"
+    return calculate_file_hash(strip_path(path))
+
+# modified from https://stackoverflow.com/questions/22058048/hashing-a-file-in-python
+def calculate_file_hash(filename: str, hash_every_n: int = 1):
+    #Larger video files were taking >.5 seconds to hash even when cached,
+    #so instead the modified time from the filesystem is used as a hash
+    h = hashlib.sha256()
+    h.update(filename.encode())
+    h.update(str(os.path.getmtime(filename)).encode())
+    return h.hexdigest()
+
+def is_safe_path(path):
+    if "VHS_STRICT_PATHS" not in os.environ:
+        return True
+    basedir = os.path.abspath('.')
+    try:
+        common_path = os.path.commonpath([basedir, path])
+    except:
+        #Different drive on windows
+        return False
+    return common_path == basedir
+
+def validate_path(path, allow_none=False, allow_url=True):
+    if path is None:
+        return allow_none
+    if is_url(path):
+        #Probably not feasible to check if url resolves here
+        if not allow_url:
+            return "URLs are unsupported for this path"
+        return is_safe_path(path)
+    if not os.path.isfile(strip_path(path)):
+        return "Invalid file path: {}".format(path)
+    return is_safe_path(path)
+
+def common_annotator_call(model, tensor_image, input_batch=False, show_pbar=False, **kwargs):
+    if "detect_resolution" in kwargs:
+        del kwargs["detect_resolution"] #Prevent weird case?
+
+    if "resolution" in kwargs:
+        detect_resolution = kwargs["resolution"] if type(kwargs["resolution"]) == int and kwargs["resolution"] >= 64 else 512
+        del kwargs["resolution"]
+    else:
+        detect_resolution = 512
+
+    if input_batch:
+        np_images = np.asarray(tensor_image * 255., dtype=np.uint8)
+        np_results = model(np_images, output_type="np", detect_resolution=detect_resolution, **kwargs)
+        return torch.from_numpy(np_results.astype(np.float32) / 255.0)
+
+    batch_size = tensor_image.shape[0]
+
+    out_tensor = None
+    for i, image in enumerate(tensor_image):
+        np_image = np.asarray(image.cpu() * 255., dtype=np.uint8)
+        np_result = model(np_image, output_type="np", detect_resolution=detect_resolution, **kwargs)
+        out = torch.from_numpy(np_result.astype(np.float32) / 255.0)
+        if out_tensor is None:
+            out_tensor = torch.zeros(batch_size, *out.shape, dtype=torch.float32)
+        out_tensor[i] = out
+
+    return out_tensor
+
+def create_node_input_types(**extra_kwargs):
+    return {
+        "required": {
+            "image": ("IMAGE",)
+        },
+        "optional": {
+            **extra_kwargs,
+            "resolution": ("INT", {"default": 512, "min": 64, "max": nodes.MAX_RESOLUTION, "step": 64})
+        }
+    }
+
+
+prompt_queue = server.PromptServer.instance.prompt_queue
+def requeue_workflow_unchecked():
+    """Requeues the current workflow without checking for multiple requeues"""
+    currently_running = prompt_queue.currently_running
+    print(f'requeue_workflow_unchecked >>>>>> ')
+    (_, _, prompt, extra_data, outputs_to_execute) = next(iter(currently_running.values()))
+
+    #Ensure batch_managers are marked stale
+    prompt = prompt.copy()
+    for uid in prompt:
+        if prompt[uid]['class_type'] == 'BatchManager':
+            prompt[uid]['inputs']['requeue'] = prompt[uid]['inputs'].get('requeue',0)+1
+
+    #execution.py has guards for concurrency, but server doesn't.
+    #TODO: Check that this won't be an issue
+    number = -server.PromptServer.instance.number
+    server.PromptServer.instance.number += 1
+    prompt_id = str(server.uuid.uuid4())
+    prompt_queue.put((number, prompt_id, prompt, extra_data, outputs_to_execute))
+    print(f'requeue_workflow_unchecked <<<<<<<<<< prompt_id:{prompt_id}, number:{number}')
+
+requeue_guard = [None, 0, 0, {}]
+def requeue_workflow(requeue_required=(-1,True)):
+    assert(len(prompt_queue.currently_running) == 1)
+    global requeue_guard
+    (run_number, _, prompt, _, _) = next(iter(prompt_queue.currently_running.values()))
+    print(f'requeue_workflow >> run_number:{run_number}\n')
+    if requeue_guard[0] != run_number:
+        #Calculate a count of how many outputs are managed by a batch manager
+        managed_outputs=0
+        for bm_uid in prompt:
+            if prompt[bm_uid]['class_type'] == 'BatchManager':
+                for output_uid in prompt:
+                    if prompt[output_uid]['class_type'] in ["VideoSaver"]:
+                        for inp in prompt[output_uid]['inputs'].values():
+                            if inp == [bm_uid, 0]:
+                                managed_outputs+=1
+        requeue_guard = [run_number, 0, managed_outputs, {}]
+    requeue_guard[1] = requeue_guard[1]+1
+    requeue_guard[3][requeue_required[0]] = requeue_required[1]
+    if requeue_guard[1] == requeue_guard[2] and max(requeue_guard[3].values()):
+        requeue_workflow_unchecked()

custom_nodes/ComfyUI-tbox/nodes/video/batch_node.py

@@ -0,0 +1,69 @@
+
+import hashlib
+import os
+
+class BatchManagerNode:
+    def __init__(self, frames_per_batch=-1):
+        print("BatchNode init")
+        self.frames_per_batch = frames_per_batch
+        self.inputs = {}
+        self.outputs = {}
+        self.unique_id = None
+        self.has_closed_inputs = False
+        self.total_frames = float('inf')
+    
+    def reset(self):
+        print("BatchNode reset")
+        self.close_inputs()
+        for key in self.outputs:
+            if getattr(self.outputs[key][-1], "gi_suspended", False):
+                try:
+                    self.outputs[key][-1].send(None)
+                except StopIteration:
+                    pass
+        self.__init__(self.frames_per_batch)
+    def has_open_inputs(self):
+        return len(self.inputs) > 0
+    def close_inputs(self):
+        for key in self.inputs:
+            if getattr(self.inputs[key][-1], "gi_suspended", False):
+                try:
+                    self.inputs[key][-1].send(1)
+                except StopIteration:
+                    pass
+        self.inputs = {}
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+                "required": {"frames_per_batch": ("INT", {"default": 16, "min": 1, "max": 128, "step": 1})},
+                "hidden": {"prompt": "PROMPT", "unique_id": "UNIQUE_ID"},
+                }
+
+    RETURN_TYPES = ("BatchManager",)
+    RETURN_NAMES = ("meta_batch",)
+    CATEGORY = "tbox/Video"
+    FUNCTION = "update_batch"
+    
+    def update_batch(self, frames_per_batch, prompt=None, unique_id=None):
+        if unique_id is not None and prompt is not None:
+            requeue = prompt[unique_id]['inputs'].get('requeue', 0)
+        else:
+            requeue = 0
+        print(f'update_batch >> unique_id: {unique_id}; requeue: {requeue}')
+        if requeue == 0:
+            self.reset()
+            self.frames_per_batch = frames_per_batch
+            self.unique_id = unique_id
+        else:
+            num_batches = (self.total_frames+self.frames_per_batch-1)//frames_per_batch
+            print(f'Meta-Batch {requeue}/{num_batches}')
+        #onExecuted seems to not be called unless some message is sent
+        return (self,)
+    
+    @classmethod
+    def IS_CHANGED(self, frames_per_batch, prompt=None, unique_id=None):
+        print(f"BatchManagerNode >>>  IS_CHANGED : {result}")
+        random_bytes = os.urandom(32)
+        result = hashlib.sha256(random_bytes).hexdigest()
+        return result

custom_nodes/ComfyUI-tbox/nodes/video/ffmpeg.py

@@ -0,0 +1,129 @@
+import re
+import os
+import torch
+import shutil
+import subprocess
+import folder_paths
+from torch import Tensor
+from collections.abc import Mapping
+
+audio_extensions = ['mp3', 'mp4', 'wav', 'ogg']
+video_extensions = ['webm', 'mp4', 'mov']
+
+def ffmpeg_suitability(path):
+    try:
+        version = subprocess.run([path, "-version"], check=True,
+                                 capture_output=True).stdout.decode("utf-8")
+    except:
+        return 0
+    score = 0
+    #rough layout of the importance of various features
+    simple_criterion = [("libvpx", 20),("264",10), ("265",3),
+                        ("svtav1",5),("libopus", 1)]
+    for criterion in simple_criterion:
+        if version.find(criterion[0]) >= 0:
+            score += criterion[1]
+    #obtain rough compile year from copyright information
+    copyright_index = version.find('2000-2')
+    if copyright_index >= 0:
+        copyright_year = version[copyright_index+6:copyright_index+9]
+        if copyright_year.isnumeric():
+            score += int(copyright_year)
+    return score
+
+folder_paths.folder_names_and_paths["VHS_video_formats"] = (
+    [
+       os.path.join(os.path.dirname(os.path.abspath(__file__)), "video_formats"),
+    ],
+    [".json"]
+)
+
+
+if "VHS_FORCE_FFMPEG_PATH" in os.environ:
+    ffmpeg_path = os.environ.get("VHS_FORCE_FFMPEG_PATH")
+else:
+    ffmpeg_paths = []
+    try:
+        from imageio_ffmpeg import get_ffmpeg_exe
+        imageio_ffmpeg_path = get_ffmpeg_exe()
+        ffmpeg_paths.append(imageio_ffmpeg_path)
+    except:
+        if "VHS_USE_IMAGEIO_FFMPEG" in os.environ:
+            raise
+        print("Failed to import imageio_ffmpeg")
+    if "VHS_USE_IMAGEIO_FFMPEG" in os.environ:
+        ffmpeg_path = imageio_ffmpeg_path
+    else:
+        system_ffmpeg = shutil.which("ffmpeg")
+        if system_ffmpeg is not None:
+            ffmpeg_paths.append(system_ffmpeg)
+        if os.path.isfile("ffmpeg"):
+            ffmpeg_paths.append(os.path.abspath("ffmpeg"))
+        if os.path.isfile("ffmpeg.exe"):
+            ffmpeg_paths.append(os.path.abspath("ffmpeg.exe"))
+        if len(ffmpeg_paths) == 0:
+            print("No valid ffmpeg found.")
+            ffmpeg_path = None
+        elif len(ffmpeg_paths) == 1:
+            #Evaluation of suitability isn't required, can take sole option
+            #to reduce startup time
+            ffmpeg_path = ffmpeg_paths[0]
+        else:
+            ffmpeg_path = max(ffmpeg_paths, key=ffmpeg_suitability)
+            
+gifski_path = os.environ.get("VHS_GIFSKI", None)
+if gifski_path is None:
+    gifski_path = os.environ.get("JOV_GIFSKI", None)
+    if gifski_path is None:
+        gifski_path = shutil.which("gifski")
+        
+ytdl_path = os.environ.get("VHS_YTDL", None) or shutil.which('yt-dlp') \
+        or shutil.which('youtube-dl')
+
+def get_audio(file, start_time=0, duration=0):
+    args = [ffmpeg_path, "-i", file]
+    if start_time > 0:
+        args += ["-ss", str(start_time)]
+    if duration > 0:
+        args += ["-t", str(duration)]
+    try:
+        #TODO: scan for sample rate and maintain
+        res =  subprocess.run(args + ["-f", "f32le", "-"],
+                              capture_output=True, check=True)
+        audio = torch.frombuffer(bytearray(res.stdout), dtype=torch.float32)
+        match = re.search(', (\\d+) Hz, (\\w+), ',res.stderr.decode('utf-8'))
+    except subprocess.CalledProcessError as e:
+        raise Exception(f"VHS failed to extract audio from {file}:\n" \
+                + e.stderr.decode("utf-8"))
+    if match:
+        ar = int(match.group(1))
+        #NOTE: Just throwing an error for other channel types right now
+        #Will deal with issues if they come
+        ac = {"mono": 1, "stereo": 2}[match.group(2)]
+    else:
+        ar = 44100
+        ac = 2
+    audio = audio.reshape((-1,ac)).transpose(0,1).unsqueeze(0)
+    return {'waveform': audio, 'sample_rate': ar}
+
+class LazyAudioMap(Mapping):
+    def __init__(self, file, start_time, duration):
+        self.file = file
+        self.start_time=start_time
+        self.duration=duration
+        self._dict=None
+    def __getitem__(self, key):
+        if self._dict is None:
+            self._dict = get_audio(self.file, self.start_time, self.duration)
+        return self._dict[key]
+    def __iter__(self):
+        if self._dict is None:
+            self._dict = get_audio(self.file, self.start_time, self.duration)
+        return iter(self._dict)
+    def __len__(self):
+        if self._dict is None:
+            self._dict = get_audio(self.file, self.start_time, self.duration)
+        return len(self._dict)
+
+def lazy_get_audio(file, start_time=0, duration=0):
+    return LazyAudioMap(file, start_time, duration)

custom_nodes/ComfyUI-tbox/nodes/video/info_node.py

@@ -0,0 +1,39 @@
+
+
+class VideoInfoNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+                "required": {
+                    "video_info": ("VHS_VIDEOINFO",),
+                    }
+                }
+
+    CATEGORY = "tbox/Video"
+
+    RETURN_TYPES = ("FLOAT", "INT", "FLOAT", "INT", "INT", "FLOAT","INT", "FLOAT", "INT", "INT")
+    RETURN_NAMES = (
+        "source_fps",
+        "source_frame_count",
+        "source_duration",
+        "source_width",
+        "source_height",
+        "loaded_fps",
+        "loaded_frame_count",
+        "loaded_duration",
+        "loaded_width",
+        "loaded_height",
+    )
+    FUNCTION = "get_video_info"
+
+    def get_video_info(self, video_info):
+        keys = ["fps", "frame_count", "duration", "width", "height"]
+        
+        source_info = []
+        loaded_info = []
+
+        for key in keys:
+            source_info.append(video_info[f"source_{key}"])
+            loaded_info.append(video_info[f"loaded_{key}"])
+
+        return (*source_info, *loaded_info)

custom_nodes/ComfyUI-tbox/nodes/video/load_node.py

@@ -0,0 +1,261 @@
+
+import os
+import cv2
+import torch
+import requests
+import itertools
+import folder_paths
+import psutil
+import numpy as np
+from comfy.utils import common_upscale
+from io import BytesIO
+from PIL import Image, ImageSequence, ImageOps
+from .ffmpeg import lazy_get_audio, video_extensions
+from ..utils import BIGMAX, DIMMAX, strip_path, validate_path
+
+
+
+
+def is_gif(filename) -> bool:
+    file_parts = filename.split('.')
+    return len(file_parts) > 1 and file_parts[-1] == "gif"
+
+def target_size(width, height, force_size, custom_width, custom_height, downscale_ratio=8) -> tuple[int, int]:
+    if force_size == "Disabled":
+        pass
+    elif force_size == "Custom Width" or force_size.endswith('x?'):
+        height *= custom_width/width
+        width = custom_width
+    elif force_size == "Custom Height" or force_size.startswith('?x'):
+        width *= custom_height/height
+        height = custom_height
+    else:
+        width = custom_width
+        height = custom_height
+    width = int(width/downscale_ratio + 0.5) * downscale_ratio
+    height = int(height/downscale_ratio + 0.5) * downscale_ratio
+    return (width, height)
+
+def cv_frame_generator(path, force_rate, frame_load_cap, skip_first_frames,
+                       select_every_nth, meta_batch=None, unique_id=None):
+    video_cap = cv2.VideoCapture(strip_path(path))
+    if not video_cap.isOpened():
+        raise ValueError(f"{path} could not be loaded with cv.")
+
+    # extract video metadata
+    fps = video_cap.get(cv2.CAP_PROP_FPS)
+    width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    total_frames = int(video_cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    duration = total_frames / fps
+
+    # set video_cap to look at start_index frame
+    total_frame_count = 0
+    total_frames_evaluated = -1
+    frames_added = 0
+    base_frame_time = 1 / fps
+    prev_frame = None
+
+    if force_rate == 0:
+        target_frame_time = base_frame_time
+    else:
+        target_frame_time = 1/force_rate
+
+    yield (width, height, fps, duration, total_frames, target_frame_time)
+    if total_frames > 0:
+        if force_rate != 0:
+            yieldable_frames = int(total_frames / fps * force_rate)
+        else:
+            yieldable_frames = total_frames
+        if frame_load_cap != 0:
+            yieldable_frames =  min(frame_load_cap, yieldable_frames)
+    else:
+        yieldable_frames = 0
+
+    if meta_batch is not None:
+        yield yieldable_frames
+
+    time_offset=target_frame_time - base_frame_time
+    while video_cap.isOpened():
+        if time_offset < target_frame_time:
+            is_returned = video_cap.grab()
+            # if didn't return frame, video has ended
+            if not is_returned:
+                break
+            time_offset += base_frame_time
+        if time_offset < target_frame_time:
+            continue
+        time_offset -= target_frame_time
+        # if not at start_index, skip doing anything with frame
+        total_frame_count += 1
+        if total_frame_count <= skip_first_frames:
+            continue
+        else:
+            total_frames_evaluated += 1
+
+        # if should not be selected, skip doing anything with frame
+        if total_frames_evaluated%select_every_nth != 0:
+            continue
+
+        # opencv loads images in BGR format (yuck), so need to convert to RGB for ComfyUI use
+        # follow up: can videos ever have an alpha channel?
+        # To my testing: No. opencv has no support for alpha
+        unused, frame = video_cap.retrieve()
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        # convert frame to comfyui's expected format
+        # TODO: frame contains no exif information. Check if opencv2 has already applied
+        frame = np.array(frame, dtype=np.float32)
+        torch.from_numpy(frame).div_(255)
+        if prev_frame is not None:
+            inp  = yield prev_frame
+            if inp is not None:
+                #ensure the finally block is called
+                return
+        prev_frame = frame
+        frames_added += 1
+        
+        # if cap exists and we've reached it, stop processing frames
+        if frame_load_cap > 0 and frames_added >= frame_load_cap:
+            break
+
+    if meta_batch is not None:
+        meta_batch.inputs.pop(unique_id)
+        meta_batch.has_closed_inputs = True
+    if prev_frame is not None:
+        yield prev_frame
+        
+def batched(it, n):
+    while batch := tuple(itertools.islice(it, n)):
+        yield batch
+        
+def load_video_cv(path: str, force_rate: int, force_size: str,
+                  custom_width: int,custom_height: int, frame_load_cap: int,
+                  skip_first_frames: int, select_every_nth: int,
+                  meta_batch=None, unique_id=None,
+                  memory_limit_mb=None):
+
+    if meta_batch is None or unique_id not in meta_batch.inputs:
+        gen = cv_frame_generator(path, force_rate, frame_load_cap, skip_first_frames,
+                                 select_every_nth, meta_batch, unique_id)
+        (width, height, fps, duration, total_frames, target_frame_time) = next(gen)
+
+        if meta_batch is not None:
+            meta_batch.inputs[unique_id] = (gen, width, height, fps, duration, total_frames, target_frame_time)
+            yieldable_frames = next(gen)
+            if yieldable_frames:
+                meta_batch.total_frames = min(meta_batch.total_frames, yieldable_frames)
+    else:
+        (gen, width, height, fps, duration, total_frames, target_frame_time) = meta_batch.inputs[unique_id]
+
+    print(f'[{width}x{height}]@{fps} - duration:{duration}, total_frames: {total_frames}')
+    
+    memory_limit = memory_limit_mb
+    if memory_limit_mb is not None:
+        memory_limit *= 2 ** 20
+    else:
+        #TODO: verify if garbage collection should be performed here.
+        #leaves ~128 MB unreserved for safety
+        try:
+            memory_limit = (psutil.virtual_memory().available + psutil.swap_memory().free) - 2 ** 27
+        except:
+            print("Failed to calculate available memory. Memory load limit has been disabled")
+            
+    if memory_limit is not None:
+        #TODO: use better estimate for when vae is not None
+        #Consider completely ignoring for load_latent case?
+        max_loadable_frames = int(memory_limit//(width*height*3*(.1)))
+      
+        if meta_batch is not None:
+            if meta_batch.frames_per_batch > max_loadable_frames:
+                raise RuntimeError(f"Meta Batch set to {meta_batch.frames_per_batch} frames but only {max_loadable_frames} can fit in memory")
+            gen = itertools.islice(gen, meta_batch.frames_per_batch)
+        else:
+            original_gen = gen
+            gen = itertools.islice(gen, max_loadable_frames)
+        
+    downscale_ratio = 8
+    frames_per_batch = (1920 * 1080 * 16) // (width * height) or 1
+    if force_size != "Disabled":
+        new_size = target_size(width, height, force_size, custom_width, custom_height, downscale_ratio)
+        if new_size[0] != width or new_size[1] != height:
+            def rescale(frame):
+                s = torch.from_numpy(np.fromiter(frame, np.dtype((np.float32, (height, width, 3)))))
+                s = s.movedim(-1,1)
+                s = common_upscale(s, new_size[0], new_size[1], "lanczos", "center")
+                return s.movedim(1,-1).numpy()
+            gen = itertools.chain.from_iterable(map(rescale, batched(gen, frames_per_batch)))
+    else:
+        new_size = width, height
+
+    #Some minor wizardry to eliminate a copy and reduce max memory by a factor of ~2
+    images = torch.from_numpy(np.fromiter(gen, np.dtype((np.float32, (new_size[1], new_size[0], 3)))))
+    if meta_batch is None and memory_limit is not None:
+        try:
+            next(original_gen)
+            raise RuntimeError(f"Memory limit hit after loading {len(images)} frames. Stopping execution.")
+        except StopIteration:
+            pass
+    if len(images) == 0:
+        raise RuntimeError("No frames generated")
+
+    #Setup lambda for lazy audio capture
+    audio = lazy_get_audio(path, skip_first_frames * target_frame_time,
+                               frame_load_cap*target_frame_time*select_every_nth)
+    #Adjust target_frame_time for select_every_nth
+    target_frame_time *= select_every_nth
+    video_info = {
+        "source_fps": fps,
+        "source_frame_count": total_frames,
+        "source_duration": duration,
+        "source_width": width,
+        "source_height": height,
+        "loaded_fps": 1/target_frame_time,
+        "loaded_frame_count": len(images),
+        "loaded_duration": len(images) * target_frame_time,
+        "loaded_width": new_size[0],
+        "loaded_height": new_size[1],
+    }
+
+    return (images, len(images), audio, video_info)
+
+
+
+class LoadVideoNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "path": ("STRING", {"default": "/Users/wadahana/Desktop/live-motion2.mp4", "multiline": True, "vhs_path_extensions": video_extensions}),
+                "force_rate": ("INT", {"default": 0, "min": 0, "max": 60, "step": 1}),
+                "force_size": (["Disabled", "Custom Height", "Custom Width", "Custom", "256x?", "?x256", "256x256", "512x?", "?x512", "512x512"],),
+                "custom_width": ("INT", {"default": 512, "min": 0, "max": DIMMAX, "step": 8}),
+                "custom_height": ("INT", {"default": 512, "min": 0, "max": DIMMAX, "step": 8}),
+                "frame_load_cap": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
+                "skip_first_frames": ("INT", {"default": 0, "min": 0, "max": BIGMAX, "step": 1}),
+                "select_every_nth": ("INT", {"default": 1, "min": 1, "max": BIGMAX, "step": 1}),
+            },
+             "optional": {
+                "meta_batch": ("BatchManager",),
+            },
+            "hidden": {
+                "unique_id": "UNIQUE_ID"
+            },
+        }
+
+    CATEGORY = "tbox/Video"
+
+    RETURN_TYPES = ("IMAGE", "INT", "AUDIO", "VHS_VIDEOINFO")
+    RETURN_NAMES = ("IMAGE", "frame_count", "audio", "video_info")
+
+    FUNCTION = "load_video"
+
+    def load_video(self, **kwargs):
+        if kwargs['path'] is None :
+            raise Exception("video is not a valid path: " + kwargs['path'])
+        
+        kwargs['path'] = kwargs['path'].split('\n')[0]
+        if validate_path(kwargs['path']) != True:
+            raise Exception("video is not a valid path: " + kwargs['path'])
+        # if is_url(kwargs['video']):
+        #     kwargs['video'] = try_download_video(kwargs['video']) or kwargs['video']
+        return load_video_cv(**kwargs)

custom_nodes/ComfyUI-tbox/nodes/video/save_node.py

@@ -0,0 +1,415 @@
+import os
+import re
+import cv2
+import sys
+import json
+import torch
+import datetime
+import itertools
+import subprocess
+import folder_paths
+import numpy as np
+from string import Template
+from pathlib import Path
+from PIL import Image, ExifTags
+from PIL.PngImagePlugin import PngInfo
+from .ffmpeg import ffmpeg_path, gifski_path
+from ..utils import tensor_to_bytes, tensor_to_shorts, requeue_workflow
+
+
+
+def gen_format_widgets(video_format):
+    for k in video_format:
+        if k.endswith("_pass"):
+            for i in range(len(video_format[k])):
+                if isinstance(video_format[k][i], list):
+                    item = [video_format[k][i]]
+                    yield item
+                    video_format[k][i] = item[0]
+        else:
+            if isinstance(video_format[k], list):
+                item = [video_format[k]]
+                yield item
+                video_format[k] = item[0]
+               
+def get_format_widget_defaults(format_name):
+    video_format_path = folder_paths.get_full_path("VHS_video_formats", format_name + ".json")
+    with open(video_format_path, 'r') as stream:
+        video_format = json.load(stream)
+    results = {}
+    for w in gen_format_widgets(video_format):
+        if len(w[0]) > 2 and 'default' in w[0][2]:
+            default = w[0][2]['default']
+        else:
+            if type(w[0][1]) is list:
+                default = w[0][1][0]
+            else:
+                #NOTE: This doesn't respect max/min, but should be good enough as a fallback to a fallback to a fallback
+                default = {"BOOLEAN": False, "INT": 0, "FLOAT": 0, "STRING": ""}[w[0][1]]
+        results[w[0][0]] = default
+    return results
+ 
+def get_video_formats():
+    formats = []
+    for format_name in folder_paths.get_filename_list("VHS_video_formats"):
+        format_name = format_name[:-5]
+        formats.append("video/" + format_name)
+    return formats
+
+def gifski_process(args, video_format, file_path, env):
+    frame_data = yield
+    with subprocess.Popen(args + video_format['main_pass'] + ['-f', 'yuv4mpegpipe', '-'],
+                          stderr=subprocess.PIPE, stdin=subprocess.PIPE,
+                          stdout=subprocess.PIPE, env=env) as procff:
+        with subprocess.Popen([gifski_path] + video_format['gifski_pass']
+                              + ['-q', '-o', file_path, '-'], stderr=subprocess.PIPE,
+                              stdin=procff.stdout, stdout=subprocess.PIPE,
+                              env=env) as procgs:
+            try:
+                while frame_data is not None:
+                    procff.stdin.write(frame_data)
+                    frame_data = yield
+                procff.stdin.flush()
+                procff.stdin.close()
+                resff = procff.stderr.read()
+                resgs = procgs.stderr.read()
+                outgs = procgs.stdout.read()
+            except BrokenPipeError as e:
+                procff.stdin.close()
+                resff = procff.stderr.read()
+                resgs = procgs.stderr.read()
+                raise Exception("An error occurred while creating gifski output\n" \
+                        + "Make sure you are using gifski --version >=1.32.0\nffmpeg: " \
+                        + resff.decode("utf-8") + '\ngifski: ' + resgs.decode("utf-8"))
+    if len(resff) > 0:
+        print(resff.decode("utf-8"), end="", file=sys.stderr)
+    if len(resgs) > 0:
+        print(resgs.decode("utf-8"), end="", file=sys.stderr)
+    #should always be empty as the quiet flag is passed
+    if len(outgs) > 0:
+        print(outgs.decode("utf-8"))
+
+def ffmpeg_process(args, video_format, video_metadata, file_path, env):
+
+    res = None
+    frame_data = yield
+    total_frames_output = 0
+    if video_format.get('save_metadata', 'False') != 'False':
+        os.makedirs(folder_paths.get_temp_directory(), exist_ok=True)
+        metadata = json.dumps(video_metadata)
+        metadata_path = os.path.join(folder_paths.get_temp_directory(), "metadata.txt")
+        #metadata from file should  escape = ; # \ and newline
+        metadata = metadata.replace("\\","\\\\")
+        metadata = metadata.replace(";","\\;")
+        metadata = metadata.replace("#","\\#")
+        metadata = metadata.replace("=","\\=")
+        metadata = metadata.replace("\n","\\\n")
+        metadata = "comment=" + metadata
+        with open(metadata_path, "w") as f:
+            f.write(";FFMETADATA1\n")
+            f.write(metadata)
+        m_args = args[:1] + ["-i", metadata_path] + args[1:] + ["-metadata", "creation_time=now"]
+        print(f'ffmpeg: {m_args}')
+        with subprocess.Popen(m_args + [file_path], stderr=subprocess.PIPE,
+                              stdin=subprocess.PIPE, env=env) as proc:
+            try:
+                while frame_data is not None:
+                    proc.stdin.write(frame_data)
+                    #TODO: skip flush for increased speed
+                    frame_data = yield
+                    total_frames_output+=1
+                proc.stdin.flush()
+                proc.stdin.close()
+                res = proc.stderr.read()
+            except BrokenPipeError as e:
+                err = proc.stderr.read()
+                #Check if output file exists. If it does, the re-execution
+                #will also fail. This obscures the cause of the error
+                #and seems to never occur concurrent to the metadata issue
+                if os.path.exists(file_path):
+                    raise Exception("An error occurred in the ffmpeg subprocess:\n" \
+                            + err.decode("utf-8"))
+                #Res was not set
+                print(err.decode("utf-8"), end="", file=sys.stderr)
+                print("An error occurred when saving with metadata")
+    if res != b'':
+        with subprocess.Popen(args + [file_path], stderr=subprocess.PIPE,
+                              stdin=subprocess.PIPE, env=env) as proc:
+            try:
+                while frame_data is not None:
+                    proc.stdin.write(frame_data)
+                    frame_data = yield
+                    total_frames_output+=1
+                proc.stdin.flush()
+                proc.stdin.close()
+                res = proc.stderr.read()
+            except BrokenPipeError as e:
+                res = proc.stderr.read()
+                raise Exception("An error occurred in the ffmpeg subprocess:\n" \
+                        + res.decode("utf-8"))
+    yield total_frames_output
+    if len(res) > 0:
+        print(res.decode("utf-8"), end="", file=sys.stderr)
+        
+def to_pingpong(inp):
+    if not hasattr(inp, "__getitem__"):
+        inp = list(inp)
+    yield from inp
+    for i in range(len(inp)-2,0,-1):
+        yield inp[i]
+        
+def apply_format_widgets(format_name, kwargs):
+    video_format_path = folder_paths.get_full_path("VHS_video_formats", format_name + ".json")
+    with open(video_format_path, 'r') as stream:
+        video_format = json.load(stream)
+
+    for w in gen_format_widgets(video_format):
+        assert(w[0][0] in kwargs)
+        if len(w[0]) > 3:
+            w[0] = Template(w[0][3]).substitute(val=kwargs[w[0][0]])
+        else:
+            w[0] = str(kwargs[w[0][0]])
+    return video_format
+
+class SaveVideoNode:
+    @classmethod
+    def INPUT_TYPES(s):
+        ffmpeg_formats = get_video_formats()
+        return {
+            "required": {
+                "path": ("STRING", {"multiline": True, "dynamicPrompts": False}),
+                "format": (ffmpeg_formats,),
+                "quality": ([100, 95, 90, 85, 80, 75, 70, 60, 50], {"default": 100}),
+                "pingpong": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "images": ("IMAGE",),
+                "audio": ("AUDIO",),
+                "frame_rate": ("INT,FLOAT", { "default": 25.0, "step": 1.0, "min": 1.0, "max": 60.0 }),
+                "meta_batch": ("BatchManager",),
+            },
+            "hidden": {
+                "prompt": "PROMPT",
+                "unique_id": "UNIQUE_ID"
+            },
+        }
+
+    RETURN_TYPES = ()
+    CATEGORY = "tbox/Video"
+    FUNCTION = "save_video"
+    OUTPUT_NODE = True
+    
+    def save_video(
+        self,
+        path,
+        frame_rate=25,
+        images=None,
+        format="video/h264-mp4",
+        quality=85,
+        pingpong=False,
+        audio=None,
+        prompt=None,
+        meta_batch=None,
+        unique_id=None,
+        manual_format_widgets=None,
+    ):
+        if images is None:
+            return {}
+        if isinstance(images, torch.Tensor) and images.size(0) == 0:
+            return {}
+        
+        if frame_rate < 1:
+            frame_rate = 1
+        elif frame_rate > 120:
+            frame_rate = 120
+            
+        num_frames = len(images)
+
+        first_image = images[0]
+        images = iter(images)
+        
+        file_path = os.path.abspath(path.split('\n')[0])
+        output_dir = os.path.dirname(file_path)
+        filename = os.path.basename(file_path)
+        name, extension = os.path.splitext(filename)
+       
+        output_process = None
+
+        video_metadata = {}
+        if prompt is not None:
+            video_metadata["prompt"] = prompt
+            
+        if meta_batch is not None and unique_id in meta_batch.outputs:
+            (counter, output_process) = meta_batch.outputs[unique_id]
+        else:
+            counter = 0 
+            output_process = None    
+
+        format_type, format_ext = format.split("/")
+
+        # Use ffmpeg to save a video
+        if ffmpeg_path is None:
+            raise ProcessLookupError(f"ffmpeg is required for video outputs and could not be found.\nIn order to use video outputs, you must either:\n- Install imageio-ffmpeg with pip,\n- Place a ffmpeg executable in {os.path.abspath('')}, or\n- Install ffmpeg and add it to the system path.")
+
+        #Acquire additional format_widget values
+        kwargs = None
+        if manual_format_widgets is None:
+            if prompt is not None:
+                kwargs = prompt[unique_id]['inputs']
+            else:
+                manual_format_widgets = {}
+        
+        if kwargs is None:
+            kwargs = get_format_widget_defaults(format_ext)
+            missing = {}
+            for k in kwargs.keys():
+                if k in manual_format_widgets:
+                    kwargs[k] = manual_format_widgets[k]
+                else:
+                    missing[k] = kwargs[k]
+            if len(missing) > 0:
+                print("Extra format values were not provided, the following defaults will be used: " + str(kwargs) + "\nThis is likely due to usage of ComfyUI-to-python. These values can be manually set by supplying a manual_format_widgets argument")
+    
+        video_format = apply_format_widgets(format_ext, kwargs)
+        has_alpha = first_image.shape[-1] == 4
+        dim_alignment = video_format.get("dim_alignment", 8)
+        if (first_image.shape[1] % dim_alignment) or (first_image.shape[0] % dim_alignment):
+            #output frames must be padded
+            to_pad = (-first_image.shape[1] % dim_alignment,
+                        -first_image.shape[0] % dim_alignment)
+            padding = (to_pad[0]//2, to_pad[0] - to_pad[0]//2,
+                        to_pad[1]//2, to_pad[1] - to_pad[1]//2)
+            padfunc = torch.nn.ReplicationPad2d(padding)
+            def pad(image):
+                image = image.permute((2,0,1))#HWC to CHW
+                padded = padfunc(image.to(dtype=torch.float32))
+                return padded.permute((1,2,0))
+            images = map(pad, images)
+            new_dims = (-first_image.shape[1] % dim_alignment + first_image.shape[1],
+                        -first_image.shape[0] % dim_alignment + first_image.shape[0])
+            dimensions = f"{new_dims[0]}x{new_dims[1]}"
+            print(f"Output images were not of valid resolution and have had padding applied: {dimensions}")
+        else:
+            dimensions = f"{first_image.shape[1]}x{first_image.shape[0]}"
+
+        if pingpong:
+            if meta_batch is not None:
+                print("pingpong is incompatible with batched output")
+            images = to_pingpong(images)
+        
+        images = map(tensor_to_bytes, images)
+        if has_alpha:
+            i_pix_fmt = 'rgba'
+        else:
+            i_pix_fmt = 'rgb24'
+                
+        args = [ffmpeg_path, "-v", "error", "-f", "rawvideo", "-pix_fmt", i_pix_fmt,
+                "-s", dimensions, "-r", str(frame_rate), "-i", "-"]
+
+        images = map(lambda x: x.tobytes(), images)
+        env=os.environ.copy()
+        if  "environment" in video_format:
+            env.update(video_format["environment"])
+
+        if "pre_pass" in video_format:
+            images = [b''.join(images)]
+            os.makedirs(folder_paths.get_temp_directory(), exist_ok=True)
+            pre_pass_args = args[:13] + video_format['pre_pass']
+            try:
+                subprocess.run(pre_pass_args, input=images[0], env=env,
+                                capture_output=True, check=True)
+            except subprocess.CalledProcessError as e:
+                raise Exception("An error occurred in the ffmpeg prepass:\n" \
+                        + e.stderr.decode("utf-8"))
+        if "inputs_main_pass" in video_format:
+            args = args[:13] + video_format['inputs_main_pass'] + args[13:]
+
+        if output_process is None:
+            args += video_format['main_pass'] 
+            output_process = ffmpeg_process(args, video_format, video_metadata, file_path, env)
+            #Proceed to first yield
+            output_process.send(None)
+            if meta_batch is not None:
+                meta_batch.outputs[unique_id] = (0, output_process)
+
+        for image in images:
+            output_process.send(image)
+        if meta_batch is not None:
+            requeue_workflow((meta_batch.unique_id, not meta_batch.has_closed_inputs))
+        if meta_batch is None or meta_batch.has_closed_inputs:
+            #Close pipe and wait for termination.
+            try:
+                total_frames_output = output_process.send(None)
+                output_process.send(None)
+            except StopIteration:
+                pass
+            if meta_batch is not None:
+                meta_batch.outputs.pop(unique_id)
+                #if len(meta_batch.outputs) == 0:
+                #    meta_batch.reset()
+        else:
+            return {}
+
+        a_waveform = None
+        if audio is not None:
+            try:
+                #safely check if audio produced by VHS_LoadVideo actually exists
+                a_waveform = audio['waveform']
+            except:
+                print(f'save audio >> not waveform')    
+                pass
+        if a_waveform is not None:
+            # Create audio file if input was provided
+            output_file_with_audio = f"{name}-audio{extension}"
+            output_file_with_audio_path = os.path.join(output_dir, output_file_with_audio)
+            if "audio_pass" not in video_format:
+                print("Selected video format does not have explicit audio support")
+                video_format["audio_pass"] = ["-c:a", "libopus"]
+
+
+            # FFmpeg command with audio re-encoding
+            #TODO: expose audio quality options if format widgets makes it in
+            #Reconsider forcing apad/shortest
+            channels = audio['waveform'].size(1)
+            min_audio_dur = total_frames_output / frame_rate + 1
+            mux_args = [ffmpeg_path, "-v", "error", "-i", file_path,
+                        "-ar", str(audio['sample_rate']), "-ac", str(channels),
+                        "-y","-f", "f32le", "-i", "-", "-c:v", "copy"] \
+                        + video_format["audio_pass"] \
+                        + ["-af", "apad=whole_dur="+str(min_audio_dur),
+                            "-shortest", output_file_with_audio_path]
+
+            audio_data = audio['waveform'].squeeze(0).transpose(0,1) \
+                    .numpy().tobytes()
+            try:
+                res = subprocess.run(mux_args, input=audio_data,
+                                        env=env, capture_output=True, check=True)
+                if res.returncode == 0:
+                    self.replace_file(output_file_with_audio_path, file_path)
+            except subprocess.CalledProcessError as e:
+                raise Exception("An error occured in the ffmpeg subprocess:\n" \
+                        + e.stderr.decode("utf-8"))
+            if res.stderr:
+                print(res.stderr.decode("utf-8"), end="", file=sys.stderr)
+
+
+        return {}
+    
+    @classmethod
+    def VALIDATE_INPUTS(self, format, **kwargs):
+        return True
+
+    def replace_file(self, audio_path, file_path):
+        try:
+            # 删除 file_path 文件（如果存在）
+            if os.path.exists(file_path):
+                os.remove(file_path)
+                print(f"Deleted file: {file_path}")
+            else:
+                print(f"File not found, skipping deletion: {file_path}")
+            
+            # 将 output_file_with_audio_path 重命名为 file_path
+            os.rename(audio_path, file_path)
+            print(f"Renamed {audio_path} to {file_path}")
+        except Exception as e:
+            print(f"An error occurred: {e}")

custom_nodes/ComfyUI-tbox/nodes/video/video_formats/h264-mp4.json

@@ -0,0 +1,10 @@
+{
+    "main_pass":
+    [
+        "-y", "-c:v", "libx264",
+        "-pix_fmt", "yuv420p",
+        "-crf", "20"
+    ],
+    "audio_pass": ["-c:a", "aac"],
+    "extension": "mp4"
+}

custom_nodes/ComfyUI-tbox/nodes/video/video_formats/h265-mp4.json

@@ -0,0 +1,13 @@
+{
+    "main_pass":
+    [
+        "-y", "-c:v", "libx265",
+        "-vtag", "hvc1",
+        "-pix_fmt", "yuv420p10le",
+        "-crf", "22",
+        "-preset", "medium",
+        "-x265-params", "log-level=quiet"
+    ],
+    "audio_pass": ["-c:a", "aac"],
+    "extension": "mp4"
+}

custom_nodes/ComfyUI-tbox/nodes/video/video_formats/nvenc_h264-mp4.json

@@ -0,0 +1,12 @@
+{
+    "main_pass":
+    [
+        "-y", "-c:v", "h264_nvenc",
+        "-pix_fmt", "yuv420p",
+        "-qp", "20"
+    ],
+    "audio_pass": ["-c:a", "aac"],
+    "bitrate": ["bitrate","INT", {"default": 10, "min": 1, "max": 999, "step": 1 }],
+    "megabit": ["megabit","BOOLEAN", {"default": true}],
+    "extension": "mp4"
+}

custom_nodes/ComfyUI-tbox/nodes/video/video_formats/nvenc_h265-mp4.json

@@ -0,0 +1,10 @@
+{
+    "main_pass":
+    [
+        "-y", "-c:v", "hevc_nvenc",
+        "-vtag", "hvc1", 
+        "-qp", "22"
+    ],
+    "audio_pass": ["-c:a", "aac"],
+    "extension": "mp4"
+}

custom_nodes/ComfyUI-tbox/nodes/video/video_formats/webm.json

@@ -0,0 +1,11 @@
+{
+    "main_pass":
+    [
+        "-y",
+        "-crf", 20, 
+        "-pix_fmt", "yuv420p",
+        "-b:v", "0"
+    ],
+    "audio_pass": ["-c:a", "libvorbis"],
+    "extension": "webm"
+}

custom_nodes/ComfyUI-tbox/requirements.txt

@@ -0,0 +1,11 @@
+torch
+torchvision
+opencv-python-headless>=4.7.0.72
+scikit-learn
+scikit-image
+insightface
+ultralytics
+onnxruntime-gpu==1.18.0 
+onnxruntime==1.18.0
+imageio_ffmpeg
+pykalman

custom_nodes/ComfyUI-tbox/src/canny/__init__.py

@@ -0,0 +1,17 @@
+import warnings
+import cv2
+import numpy as np
+from PIL import Image
+from common import resize_image_with_pad, common_input_validate, HWC3
+
+class CannyDetector:
+    def __call__(self, input_image=None, low_threshold=100, high_threshold=200, detect_resolution=512, output_type=None, upscale_method="INTER_CUBIC", **kwargs):
+        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
+        detected_map, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
+        detected_map = cv2.Canny(detected_map, low_threshold, high_threshold)
+        detected_map = HWC3(remove_pad(detected_map))
+        
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
+            
+        return detected_map

custom_nodes/ComfyUI-tbox/src/common.py

@@ -0,0 +1,186 @@
+import cv2
+import numpy as np
+import torch
+import os
+import tempfile
+import warnings
+from contextlib import suppress
+from pathlib import Path
+from huggingface_hub import constants, hf_hub_download
+from ast import literal_eval
+
+TEMP_DIR = tempfile.gettempdir()
+ANNOTATOR_CKPTS_PATH = os.path.join(Path(__file__).parents[2], 'ckpts')
+USE_SYMLINKS = False
+
+
+BIGMIN = -(2**53-1)
+BIGMAX = (2**53-1)
+
+DIMMAX = 8192
+
+try:
+    ANNOTATOR_CKPTS_PATH = os.environ['AUX_ANNOTATOR_CKPTS_PATH']
+except:
+    warnings.warn("Custom pressesor model path not set successfully.")
+    pass
+
+try:
+    USE_SYMLINKS = literal_eval(os.environ['AUX_USE_SYMLINKS'])
+except:
+    warnings.warn("USE_SYMLINKS not set successfully. Using default value: False to download models.")
+    pass
+
+try:
+    TEMP_DIR = os.environ['AUX_TEMP_DIR']
+    if len(TEMP_DIR) >= 60:
+        warnings.warn(f"custom temp dir is too long. Using default")
+        TEMP_DIR = tempfile.gettempdir()
+except:
+    warnings.warn(f"custom temp dir not set successfully")
+    pass
+
+here = Path(__file__).parent.resolve()
+
+def safer_memory(x):
+    # Fix many MAC/AMD problems
+    return np.ascontiguousarray(x.copy()).copy()
+
+UPSCALE_METHODS = ["INTER_NEAREST", "INTER_LINEAR", "INTER_AREA", "INTER_CUBIC", "INTER_LANCZOS4"]
+def get_upscale_method(method_str):
+    assert method_str in UPSCALE_METHODS, f"Method {method_str} not found in {UPSCALE_METHODS}"
+    return getattr(cv2, method_str)
+
+def pad64(x):
+    return int(np.ceil(float(x) / 64.0) * 64 - x)
+
+def resize_image_with_pad(input_image, resolution, upscale_method = "", skip_hwc3=False, mode='edge'):
+    if skip_hwc3:
+        img = input_image
+    else:
+        img = HWC3(input_image)
+    H_raw, W_raw, _ = img.shape
+    if resolution == 0:
+        return img, lambda x: x
+    k = float(resolution) / float(min(H_raw, W_raw))
+    H_target = int(np.round(float(H_raw) * k))
+    W_target = int(np.round(float(W_raw) * k))
+    img = cv2.resize(img, (W_target, H_target), interpolation=get_upscale_method(upscale_method) if k > 1 else cv2.INTER_AREA)
+    H_pad, W_pad = pad64(H_target), pad64(W_target)
+    img_padded = np.pad(img, [[0, H_pad], [0, W_pad], [0, 0]], mode=mode)
+
+    def remove_pad(x):
+        return safer_memory(x[:H_target, :W_target, ...])
+
+    return safer_memory(img_padded), remove_pad
+    
+    
+def common_input_validate(input_image, output_type, **kwargs):
+    if "img" in kwargs:
+            warnings.warn("img is deprecated, please use `input_image=...` instead.", DeprecationWarning)
+            input_image = kwargs.pop("img")
+    
+    if "return_pil" in kwargs:
+            warnings.warn("return_pil is deprecated. Use output_type instead.", DeprecationWarning)
+            output_type = "pil" if kwargs["return_pil"] else "np"
+    
+    if type(output_type) is bool:
+        warnings.warn("Passing `True` or `False` to `output_type` is deprecated and will raise an error in future versions")
+        if output_type:
+            output_type = "pil"
+
+    if input_image is None:
+        raise ValueError("input_image must be defined.")
+
+    if not isinstance(input_image, np.ndarray):
+        input_image = np.array(input_image, dtype=np.uint8)
+        output_type = output_type or "pil"
+    else:
+        output_type = output_type or "np"
+    
+    return (input_image, output_type)
+
+def custom_hf_download(pretrained_model_or_path, filename, cache_dir=TEMP_DIR, ckpts_dir=ANNOTATOR_CKPTS_PATH, subfolder=str(""), use_symlinks=USE_SYMLINKS, repo_type="model"):
+
+    print(f'cache_dir: {cache_dir}')
+    print(f'ckpts_dir: {ckpts_dir}')
+    print(f'use_symlinks: {use_symlinks}')
+    local_dir = os.path.join(ckpts_dir, pretrained_model_or_path)
+    model_path = os.path.join(local_dir, *subfolder.split('/'), filename)
+
+    if len(str(model_path)) >= 255:
+        warnings.warn(f"Path {model_path} is too long, \n please change annotator_ckpts_path in config.yaml")
+
+    if not os.path.exists(model_path):
+        print(f"Failed to find {model_path}.\n Downloading from huggingface.co")
+        print(f"cacher folder is {cache_dir}, you can change it by custom_tmp_path in config.yaml")
+        if use_symlinks:
+            cache_dir_d = constants.HF_HUB_CACHE    # use huggingface newer env variables `HF_HUB_CACHE`
+            if cache_dir_d is None:
+                import platform
+                if platform.system() == "Windows":
+                    cache_dir_d = os.path.join(os.getenv("USERPROFILE"), ".cache", "huggingface", "hub")
+                else:
+                    cache_dir_d = os.path.join(os.getenv("HOME"), ".cache", "huggingface", "hub")
+            try:
+                # test_link
+                Path(cache_dir_d).mkdir(parents=True, exist_ok=True)
+                Path(ckpts_dir).mkdir(parents=True, exist_ok=True)
+                (Path(cache_dir_d) / f"linktest_{filename}.txt").touch()
+                # symlink instead of link avoid `invalid cross-device link` error.
+                os.symlink(os.path.join(cache_dir_d, f"linktest_{filename}.txt"), os.path.join(ckpts_dir, f"linktest_{filename}.txt"))
+                print("Using symlinks to download models. \n",\
+                      "Make sure you have enough space on your cache folder. \n",\
+                      "And do not purge the cache folder after downloading.\n",\
+                      "Otherwise, you will have to re-download the models every time you run the script.\n",\
+                      "You can use USE_SYMLINKS: False in config.yaml to avoid this behavior.")
+            except:
+                print("Maybe not able to create symlink. Disable using symlinks.")
+                use_symlinks = False
+                cache_dir_d = os.path.join(cache_dir, "ckpts", pretrained_model_or_path)
+            finally:    # always remove test link files
+                with suppress(FileNotFoundError):
+                    os.remove(os.path.join(ckpts_dir, f"linktest_{filename}.txt"))
+                    os.remove(os.path.join(cache_dir_d, f"linktest_{filename}.txt"))
+        else:
+            cache_dir_d = os.path.join(cache_dir, "ckpts", pretrained_model_or_path)
+
+        model_path = hf_hub_download(repo_id=pretrained_model_or_path,
+            cache_dir=cache_dir_d,
+            local_dir=local_dir,
+            subfolder=subfolder,
+            filename=filename,
+            local_dir_use_symlinks=use_symlinks,
+            resume_download=True,
+            etag_timeout=100,
+            repo_type=repo_type
+        )
+        if not use_symlinks:
+            try:
+                import shutil
+                shutil.rmtree(os.path.join(cache_dir, "ckpts"))
+            except Exception as e :
+                print(e)
+
+    print(f"model_path is {model_path}")
+
+    return model_path
+
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y

custom_nodes/ComfyUI-tbox/src/densepose/__init__.py

@@ -0,0 +1,67 @@
+import torchvision # Fix issue Unknown builtin op: torchvision::nms
+import cv2
+import numpy as np
+import torch
+import torch.nn as nn
+from einops import rearrange
+from PIL import Image
+
+from common import HWC3, resize_image_with_pad, common_input_validate, custom_hf_download
+from .densepose import DensePoseMaskedColormapResultsVisualizer, _extract_i_from_iuvarr, densepose_chart_predictor_output_to_result_with_confidences
+
+N_PART_LABELS = 24
+DENSEPOSE_MODEL_NAME = "LayerNorm/DensePose-TorchScript-with-hint-image"
+
+class DenseposeDetector:
+    def __init__(self, model):
+        self.dense_pose_estimation = model
+        self.device = "cpu"
+        self.result_visualizer = DensePoseMaskedColormapResultsVisualizer(
+            alpha=1, 
+            data_extractor=_extract_i_from_iuvarr, 
+            segm_extractor=_extract_i_from_iuvarr, 
+            val_scale = 255.0 / N_PART_LABELS
+        )
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path=DENSEPOSE_MODEL_NAME, filename="densepose_r50_fpn_dl.torchscript"):
+        torchscript_model_path = custom_hf_download(pretrained_model_or_path, filename)
+        densepose = torch.jit.load(torchscript_model_path, map_location="cpu")
+        return cls(densepose)
+
+    def to(self, device):
+        self.dense_pose_estimation.to(device)
+        self.device = device
+        return self
+    
+    def __call__(self, input_image, detect_resolution=512, output_type="pil", upscale_method="INTER_CUBIC", cmap="viridis", **kwargs):
+        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
+        input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
+        H, W  = input_image.shape[:2]
+
+        hint_image_canvas = np.zeros([H, W], dtype=np.uint8)
+        hint_image_canvas = np.tile(hint_image_canvas[:, :, np.newaxis], [1, 1, 3])
+
+        input_image = rearrange(torch.from_numpy(input_image).to(self.device), 'h w c -> c h w')
+
+        pred_boxes, corase_segm, fine_segm, u, v = self.dense_pose_estimation(input_image)
+
+        extractor = densepose_chart_predictor_output_to_result_with_confidences
+        densepose_results = [extractor(pred_boxes[i:i+1], corase_segm[i:i+1], fine_segm[i:i+1], u[i:i+1], v[i:i+1]) for i in range(len(pred_boxes))]
+
+        if cmap=="viridis":
+            self.result_visualizer.mask_visualizer.cmap = cv2.COLORMAP_VIRIDIS
+            hint_image = self.result_visualizer.visualize(hint_image_canvas, densepose_results)
+            hint_image = cv2.cvtColor(hint_image, cv2.COLOR_BGR2RGB)
+            hint_image[:, :, 0][hint_image[:, :, 0] == 0] = 68
+            hint_image[:, :, 1][hint_image[:, :, 1] == 0] = 1
+            hint_image[:, :, 2][hint_image[:, :, 2] == 0] = 84
+        else:
+            self.result_visualizer.mask_visualizer.cmap = cv2.COLORMAP_PARULA
+            hint_image = self.result_visualizer.visualize(hint_image_canvas, densepose_results)
+            hint_image = cv2.cvtColor(hint_image, cv2.COLOR_BGR2RGB)
+
+        detected_map = remove_pad(HWC3(hint_image))
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
+        return detected_map

custom_nodes/ComfyUI-tbox/src/densepose/densepose.py

@@ -0,0 +1,347 @@
+from typing import Tuple
+import math
+import numpy as np
+from enum import IntEnum
+from typing import List, Tuple, Union
+import torch
+from torch.nn import functional as F
+import logging
+import cv2
+
+Image = np.ndarray
+Boxes = torch.Tensor
+ImageSizeType = Tuple[int, int]
+_RawBoxType = Union[List[float], Tuple[float, ...], torch.Tensor, np.ndarray]
+IntTupleBox = Tuple[int, int, int, int]
+
+class BoxMode(IntEnum):
+    """
+    Enum of different ways to represent a box.
+    """
+
+    XYXY_ABS = 0
+    """
+    (x0, y0, x1, y1) in absolute floating points coordinates.
+    The coordinates in range [0, width or height].
+    """
+    XYWH_ABS = 1
+    """
+    (x0, y0, w, h) in absolute floating points coordinates.
+    """
+    XYXY_REL = 2
+    """
+    Not yet supported!
+    (x0, y0, x1, y1) in range [0, 1]. They are relative to the size of the image.
+    """
+    XYWH_REL = 3
+    """
+    Not yet supported!
+    (x0, y0, w, h) in range [0, 1]. They are relative to the size of the image.
+    """
+    XYWHA_ABS = 4
+    """
+    (xc, yc, w, h, a) in absolute floating points coordinates.
+    (xc, yc) is the center of the rotated box, and the angle a is in degrees ccw.
+    """
+
+    @staticmethod
+    def convert(box: _RawBoxType, from_mode: "BoxMode", to_mode: "BoxMode") -> _RawBoxType:
+        """
+        Args:
+            box: can be a k-tuple, k-list or an Nxk array/tensor, where k = 4 or 5
+            from_mode, to_mode (BoxMode)
+
+        Returns:
+            The converted box of the same type.
+        """
+        if from_mode == to_mode:
+            return box
+
+        original_type = type(box)
+        is_numpy = isinstance(box, np.ndarray)
+        single_box = isinstance(box, (list, tuple))
+        if single_box:
+            assert len(box) == 4 or len(box) == 5, (
+                "BoxMode.convert takes either a k-tuple/list or an Nxk array/tensor,"
+                " where k == 4 or 5"
+            )
+            arr = torch.tensor(box)[None, :]
+        else:
+            # avoid modifying the input box
+            if is_numpy:
+                arr = torch.from_numpy(np.asarray(box)).clone()
+            else:
+                arr = box.clone()
+
+        assert to_mode not in [BoxMode.XYXY_REL, BoxMode.XYWH_REL] and from_mode not in [
+            BoxMode.XYXY_REL,
+            BoxMode.XYWH_REL,
+        ], "Relative mode not yet supported!"
+
+        if from_mode == BoxMode.XYWHA_ABS and to_mode == BoxMode.XYXY_ABS:
+            assert (
+                arr.shape[-1] == 5
+            ), "The last dimension of input shape must be 5 for XYWHA format"
+            original_dtype = arr.dtype
+            arr = arr.double()
+
+            w = arr[:, 2]
+            h = arr[:, 3]
+            a = arr[:, 4]
+            c = torch.abs(torch.cos(a * math.pi / 180.0))
+            s = torch.abs(torch.sin(a * math.pi / 180.0))
+            # This basically computes the horizontal bounding rectangle of the rotated box
+            new_w = c * w + s * h
+            new_h = c * h + s * w
+
+            # convert center to top-left corner
+            arr[:, 0] -= new_w / 2.0
+            arr[:, 1] -= new_h / 2.0
+            # bottom-right corner
+            arr[:, 2] = arr[:, 0] + new_w
+            arr[:, 3] = arr[:, 1] + new_h
+
+            arr = arr[:, :4].to(dtype=original_dtype)
+        elif from_mode == BoxMode.XYWH_ABS and to_mode == BoxMode.XYWHA_ABS:
+            original_dtype = arr.dtype
+            arr = arr.double()
+            arr[:, 0] += arr[:, 2] / 2.0
+            arr[:, 1] += arr[:, 3] / 2.0
+            angles = torch.zeros((arr.shape[0], 1), dtype=arr.dtype)
+            arr = torch.cat((arr, angles), axis=1).to(dtype=original_dtype)
+        else:
+            if to_mode == BoxMode.XYXY_ABS and from_mode == BoxMode.XYWH_ABS:
+                arr[:, 2] += arr[:, 0]
+                arr[:, 3] += arr[:, 1]
+            elif from_mode == BoxMode.XYXY_ABS and to_mode == BoxMode.XYWH_ABS:
+                arr[:, 2] -= arr[:, 0]
+                arr[:, 3] -= arr[:, 1]
+            else:
+                raise NotImplementedError(
+                    "Conversion from BoxMode {} to {} is not supported yet".format(
+                        from_mode, to_mode
+                    )
+                )
+
+        if single_box:
+            return original_type(arr.flatten().tolist())
+        if is_numpy:
+            return arr.numpy()
+        else:
+            return arr
+
+class MatrixVisualizer:
+    """
+    Base visualizer for matrix data
+    """
+
+    def __init__(
+        self,
+        inplace=True,
+        cmap=cv2.COLORMAP_PARULA,
+        val_scale=1.0,
+        alpha=0.7,
+        interp_method_matrix=cv2.INTER_LINEAR,
+        interp_method_mask=cv2.INTER_NEAREST,
+    ):
+        self.inplace = inplace
+        self.cmap = cmap
+        self.val_scale = val_scale
+        self.alpha = alpha
+        self.interp_method_matrix = interp_method_matrix
+        self.interp_method_mask = interp_method_mask
+
+    def visualize(self, image_bgr, mask, matrix, bbox_xywh):
+        self._check_image(image_bgr)
+        self._check_mask_matrix(mask, matrix)
+        if self.inplace:
+            image_target_bgr = image_bgr
+        else:
+            image_target_bgr = image_bgr * 0
+        x, y, w, h = [int(v) for v in bbox_xywh]
+        if w <= 0 or h <= 0:
+            return image_bgr
+        mask, matrix = self._resize(mask, matrix, w, h)
+        mask_bg = np.tile((mask == 0)[:, :, np.newaxis], [1, 1, 3])
+        matrix_scaled = matrix.astype(np.float32) * self.val_scale
+        _EPSILON = 1e-6
+        if np.any(matrix_scaled > 255 + _EPSILON):
+            logger = logging.getLogger(__name__)
+            logger.warning(
+                f"Matrix has values > {255 + _EPSILON} after " f"scaling, clipping to [0..255]"
+            )
+        matrix_scaled_8u = matrix_scaled.clip(0, 255).astype(np.uint8)
+        matrix_vis = cv2.applyColorMap(matrix_scaled_8u, self.cmap)
+        matrix_vis[mask_bg] = image_target_bgr[y : y + h, x : x + w, :][mask_bg]
+        image_target_bgr[y : y + h, x : x + w, :] = (
+            image_target_bgr[y : y + h, x : x + w, :] * (1.0 - self.alpha) + matrix_vis * self.alpha
+        )
+        return image_target_bgr.astype(np.uint8)
+
+    def _resize(self, mask, matrix, w, h):
+        if (w != mask.shape[1]) or (h != mask.shape[0]):
+            mask = cv2.resize(mask, (w, h), self.interp_method_mask)
+        if (w != matrix.shape[1]) or (h != matrix.shape[0]):
+            matrix = cv2.resize(matrix, (w, h), self.interp_method_matrix)
+        return mask, matrix
+
+    def _check_image(self, image_rgb):
+        assert len(image_rgb.shape) == 3
+        assert image_rgb.shape[2] == 3
+        assert image_rgb.dtype == np.uint8
+
+    def _check_mask_matrix(self, mask, matrix):
+        assert len(matrix.shape) == 2
+        assert len(mask.shape) == 2
+        assert mask.dtype == np.uint8
+
+class DensePoseResultsVisualizer:
+    def visualize(
+        self,
+        image_bgr: Image,
+        results,
+    ) -> Image:
+        context = self.create_visualization_context(image_bgr)
+        for i, result in enumerate(results):
+            boxes_xywh, labels, uv = result
+            iuv_array = torch.cat(
+                (labels[None].type(torch.float32), uv * 255.0)
+            ).type(torch.uint8)
+            self.visualize_iuv_arr(context, iuv_array.cpu().numpy(), boxes_xywh)
+        image_bgr = self.context_to_image_bgr(context)
+        return image_bgr
+
+    def create_visualization_context(self, image_bgr: Image):
+        return image_bgr
+
+    def visualize_iuv_arr(self, context, iuv_arr: np.ndarray, bbox_xywh) -> None:
+        pass
+
+    def context_to_image_bgr(self, context):
+        return context
+
+    def get_image_bgr_from_context(self, context):
+        return context
+
+class DensePoseMaskedColormapResultsVisualizer(DensePoseResultsVisualizer):
+    def __init__(
+        self,
+        data_extractor,
+        segm_extractor,
+        inplace=True,
+        cmap=cv2.COLORMAP_PARULA,
+        alpha=0.7,
+        val_scale=1.0,
+        **kwargs,
+    ):
+        self.mask_visualizer = MatrixVisualizer(
+            inplace=inplace, cmap=cmap, val_scale=val_scale, alpha=alpha
+        )
+        self.data_extractor = data_extractor
+        self.segm_extractor = segm_extractor
+
+    def context_to_image_bgr(self, context):
+        return context
+
+    def visualize_iuv_arr(self, context, iuv_arr: np.ndarray, bbox_xywh) -> None:
+        image_bgr = self.get_image_bgr_from_context(context)
+        matrix = self.data_extractor(iuv_arr)
+        segm = self.segm_extractor(iuv_arr)
+        mask = np.zeros(matrix.shape, dtype=np.uint8)
+        mask[segm > 0] = 1
+        image_bgr = self.mask_visualizer.visualize(image_bgr, mask, matrix, bbox_xywh)
+
+
+def _extract_i_from_iuvarr(iuv_arr):
+    return iuv_arr[0, :, :]
+
+
+def _extract_u_from_iuvarr(iuv_arr):
+    return iuv_arr[1, :, :]
+
+
+def _extract_v_from_iuvarr(iuv_arr):
+    return iuv_arr[2, :, :]
+
+def make_int_box(box: torch.Tensor) -> IntTupleBox:
+    int_box = [0, 0, 0, 0]
+    int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
+    return int_box[0], int_box[1], int_box[2], int_box[3]
+
+def densepose_chart_predictor_output_to_result_with_confidences(
+    boxes: Boxes,
+    coarse_segm,
+    fine_segm,
+    u, v
+
+):
+    boxes_xyxy_abs = boxes.clone()
+    boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
+    box_xywh = make_int_box(boxes_xywh_abs[0])
+
+    labels = resample_fine_and_coarse_segm_tensors_to_bbox(fine_segm, coarse_segm, box_xywh).squeeze(0)
+    uv = resample_uv_tensors_to_bbox(u, v, labels, box_xywh)
+    confidences = []
+    return box_xywh, labels, uv
+
+def resample_fine_and_coarse_segm_tensors_to_bbox(
+    fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
+):
+    """
+    Resample fine and coarse segmentation tensors to the given
+    bounding box and derive labels for each pixel of the bounding box
+
+    Args:
+        fine_segm: float tensor of shape [1, C, Hout, Wout]
+        coarse_segm: float tensor of shape [1, K, Hout, Wout]
+        box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
+            corner coordinates, width (W) and height (H)
+    Return:
+        Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    # coarse segmentation
+    coarse_segm_bbox = F.interpolate(
+        coarse_segm,
+        (h, w),
+        mode="bilinear",
+        align_corners=False,
+    ).argmax(dim=1)
+    # combined coarse and fine segmentation
+    labels = (
+        F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
+        * (coarse_segm_bbox > 0).long()
+    )
+    return labels
+
+def resample_uv_tensors_to_bbox(
+    u: torch.Tensor,
+    v: torch.Tensor,
+    labels: torch.Tensor,
+    box_xywh_abs: IntTupleBox,
+) -> torch.Tensor:
+    """
+    Resamples U and V coordinate estimates for the given bounding box
+
+    Args:
+        u (tensor [1, C, H, W] of float): U coordinates
+        v (tensor [1, C, H, W] of float): V coordinates
+        labels (tensor [H, W] of long): labels obtained by resampling segmentation
+            outputs for the given bounding box
+        box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
+    Return:
+       Resampled U and V coordinates - a tensor [2, H, W] of float
+    """
+    x, y, w, h = box_xywh_abs
+    w = max(int(w), 1)
+    h = max(int(h), 1)
+    u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
+    v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
+    uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
+    for part_id in range(1, u_bbox.size(1)):
+        uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
+        uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
+    return uv
+

custom_nodes/ComfyUI-tbox/src/dwpose/LICENSE

@@ -0,0 +1,108 @@
+OPENPOSE: MULTIPERSON KEYPOINT DETECTION
+SOFTWARE LICENSE AGREEMENT
+ACADEMIC OR NON-PROFIT ORGANIZATION NONCOMMERCIAL RESEARCH USE ONLY
+
+BY USING OR DOWNLOADING THE SOFTWARE, YOU ARE AGREEING TO THE TERMS OF THIS LICENSE AGREEMENT.  IF YOU DO NOT AGREE WITH THESE TERMS, YOU MAY NOT USE OR DOWNLOAD THE SOFTWARE.
+
+This is a license agreement ("Agreement") between your academic institution or non-profit organization or self (called "Licensee" or "You" in this Agreement) and Carnegie Mellon University (called "Licensor" in this Agreement).  All rights not specifically granted to you in this Agreement are reserved for Licensor. 
+
+RESERVATION OF OWNERSHIP AND GRANT OF LICENSE: 
+Licensor retains exclusive ownership of any copy of the Software (as defined below) licensed under this Agreement and hereby grants to Licensee a personal, non-exclusive, 
+non-transferable license to use the Software for noncommercial research purposes, without the right to sublicense, pursuant to the terms and conditions of this Agreement.  As used in this Agreement, the term "Software" means (i) the actual copy of all or any portion of code for program routines made accessible to Licensee by Licensor pursuant to this Agreement, inclusive of backups, updates, and/or merged copies permitted hereunder or subsequently supplied by Licensor,  including all or any file structures, programming instructions, user interfaces and screen formats and sequences as well as any and all documentation and instructions related to it, and (ii) all or any derivatives and/or modifications created or made by You to any of the items specified in (i).
+
+CONFIDENTIALITY: Licensee acknowledges that the Software is proprietary to Licensor, and as such, Licensee agrees to receive all such materials in confidence and use the Software only in accordance with the terms of this Agreement.  Licensee agrees to use reasonable effort to protect the Software from unauthorized use, reproduction, distribution, or publication.
+
+COPYRIGHT: The Software is owned by Licensor and is protected by United 
+States copyright laws and applicable international treaties and/or conventions.
+
+PERMITTED USES:  The Software may be used for your own noncommercial internal research purposes. You understand and agree that Licensor is not obligated to implement any suggestions and/or feedback you might provide regarding the Software, but to the extent Licensor does so, you are not entitled to any compensation related thereto.
+
+DERIVATIVES: You may create derivatives of or make modifications to the Software, however, You agree that all and any such derivatives and modifications will be owned by Licensor and become a part of the Software licensed to You under this Agreement.  You may only use such derivatives and modifications for your own noncommercial internal research purposes, and you may not otherwise use, distribute or copy such derivatives and modifications in violation of this Agreement.
+
+BACKUPS:  If Licensee is an organization, it may make that number of copies of the Software necessary for internal noncommercial use at a single site within its organization provided that all information appearing in or on the original labels, including the copyright and trademark notices are copied onto the labels of the copies.
+
+USES NOT PERMITTED:  You may not distribute, copy or use the Software except as explicitly permitted herein. Licensee has not been granted any trademark license as part of this Agreement and may not use the name or mark “OpenPose", "Carnegie Mellon" or any renditions thereof without the prior written permission of Licensor.
+
+You may not sell, rent, lease, sublicense, lend, time-share or transfer, in whole or in part, or provide third parties access to prior or present versions (or any parts thereof) of the Software.
+
+ASSIGNMENT: You may not assign this Agreement or your rights hereunder without the prior written consent of Licensor. Any attempted assignment without such consent shall be null and void.
+
+TERM: The term of the license granted by this Agreement is from Licensee's acceptance of this Agreement by downloading the Software or by using the Software until terminated as provided below.
+
+The Agreement automatically terminates without notice if you fail to comply with any provision of this Agreement.  Licensee may terminate this Agreement by ceasing using the Software.  Upon any termination of this Agreement, Licensee will delete any and all copies of the Software. You agree that all provisions which operate to protect the proprietary rights of Licensor shall remain in force should breach occur and that the obligation of confidentiality described in this Agreement is binding in perpetuity and, as such, survives the term of the Agreement.
+
+FEE: Provided Licensee abides completely by the terms and conditions of this Agreement, there is no fee due to Licensor for Licensee's use of the Software in accordance with this Agreement.
+
+DISCLAIMER OF WARRANTIES:  THE SOFTWARE IS PROVIDED "AS-IS" WITHOUT WARRANTY OF ANY KIND INCLUDING ANY WARRANTIES OF PERFORMANCE OR MERCHANTABILITY OR FITNESS FOR A PARTICULAR USE OR PURPOSE OR OF NON-INFRINGEMENT.  LICENSEE BEARS ALL RISK RELATING TO QUALITY AND PERFORMANCE OF THE SOFTWARE AND RELATED MATERIALS.
+
+SUPPORT AND MAINTENANCE: No Software support or training by the Licensor is provided as part of this Agreement.  
+
+EXCLUSIVE REMEDY AND LIMITATION OF LIABILITY: To the maximum extent permitted under applicable law, Licensor shall not be liable for direct, indirect, special, incidental, or consequential damages or lost profits related to Licensee's use of and/or inability to use the Software, even if Licensor is advised of the possibility of such damage.
+
+EXPORT REGULATION: Licensee agrees to comply with any and all applicable 
+U.S. export control laws, regulations, and/or other laws related to embargoes and sanction programs administered by the Office of Foreign Assets Control.
+
+SEVERABILITY: If any provision(s) of this Agreement shall be held to be invalid, illegal, or unenforceable by a court or other tribunal of competent jurisdiction, the validity, legality and enforceability of the remaining provisions shall not in any way be affected or impaired thereby.
+
+NO IMPLIED WAIVERS: No failure or delay by Licensor in enforcing any right or remedy under this Agreement shall be construed as a waiver of any future or other exercise of such right or remedy by Licensor.
+
+GOVERNING LAW: This Agreement shall be construed and enforced in accordance with the laws of the Commonwealth of Pennsylvania without reference to conflict of laws principles.  You consent to the personal jurisdiction of the courts of this County and waive their rights to venue outside of Allegheny County, Pennsylvania.
+
+ENTIRE AGREEMENT AND AMENDMENTS: This Agreement constitutes the sole and entire agreement between Licensee and Licensor as to the matter set forth herein and supersedes any previous agreements, understandings, and arrangements between the parties relating hereto.
+
+
+
+************************************************************************
+
+THIRD-PARTY SOFTWARE NOTICES AND INFORMATION
+
+This project incorporates material from the project(s) listed below (collectively, "Third Party Code").  This Third Party Code is licensed to you under their original license terms set forth below.  We reserves all other rights not expressly granted, whether by implication, estoppel or otherwise.
+ 
+1.	Caffe, version 1.0.0, (https://github.com/BVLC/caffe/)
+
+COPYRIGHT
+
+All contributions by the University of California:
+Copyright (c) 2014-2017 The Regents of the University of California (Regents)
+All rights reserved.
+
+All other contributions:
+Copyright (c) 2014-2017, the respective contributors
+All rights reserved.
+
+Caffe uses a shared copyright model: each contributor holds copyright over
+their contributions to Caffe. The project versioning records all such
+contribution and copyright details. If a contributor wants to further mark
+their specific copyright on a particular contribution, they should indicate
+their copyright solely in the commit message of the change when it is
+committed.
+
+LICENSE
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met: 
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer. 
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution. 
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+CONTRIBUTION AGREEMENT
+
+By contributing to the BVLC/caffe repository through pull-request, comment,
+or otherwise, the contributor releases their content to the
+license and copyright terms herein.
+
+************END OF THIRD-PARTY SOFTWARE NOTICES AND INFORMATION**********

custom_nodes/ComfyUI-tbox/src/dwpose/__init__.py

@@ -0,0 +1,328 @@
+# Openpose
+# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
+# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
+# 3rd Edited by ControlNet
+# 4th Edited by ControlNet (added face and correct hands)
+# 5th Edited by ControlNet (Improved JSON serialization/deserialization, and lots of bug fixs)
+# This preprocessor is licensed by CMU for non-commercial use only.
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+
+import json
+import torch
+import numpy as np
+from . import util
+from .body import Body, BodyResult, Keypoint
+from .hand import Hand
+from .face import Face
+from .types import PoseResult, HandResult, FaceResult, AnimalPoseResult
+#from huggingface_hub import hf_hub_download
+from .wholebody import Wholebody
+import warnings
+from common import HWC3, resize_image_with_pad, common_input_validate, custom_hf_download
+import cv2
+from PIL import Image
+from .animalpose import AnimalPoseImage
+
+from typing import Tuple, List, Callable, Union, Optional
+
+
+def draw_animalposes(animals: list[list[Keypoint]], H: int, W: int) -> np.ndarray:
+    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+    for animal_pose in animals:
+        canvas = draw_animalpose(canvas, animal_pose)
+    return canvas
+
+
+def draw_animalpose(canvas: np.ndarray, keypoints: list[Keypoint]) -> np.ndarray:
+    # order of the keypoints for AP10k and a standardized list of colors for limbs
+    keypointPairsList = [
+        (1, 2),
+        (2, 3),
+        (1, 3),
+        (3, 4),
+        (4, 9),
+        (9, 10),
+        (10, 11),
+        (4, 6),
+        (6, 7),
+        (7, 8),
+        (4, 5),
+        (5, 15),
+        (15, 16),
+        (16, 17),
+        (5, 12),
+        (12, 13),
+        (13, 14),
+    ]
+    colorsList = [
+        (255, 255, 255),
+        (100, 255, 100),
+        (150, 255, 255),
+        (100, 50, 255),
+        (50, 150, 200),
+        (0, 255, 255),
+        (0, 150, 0),
+        (0, 0, 255),
+        (0, 0, 150),
+        (255, 50, 255),
+        (255, 0, 255),
+        (255, 0, 0),
+        (150, 0, 0),
+        (255, 255, 100),
+        (0, 150, 0),
+        (255, 255, 0),
+        (150, 150, 150),
+    ]  # 16 colors needed
+
+    for ind, (i, j) in enumerate(keypointPairsList):
+        p1 = keypoints[i - 1]
+        p2 = keypoints[j - 1]
+
+        if p1 is not None and p2 is not None:
+            cv2.line(
+                canvas,
+                (int(p1.x), int(p1.y)),
+                (int(p2.x), int(p2.y)),
+                colorsList[ind],
+                5,
+            )
+    return canvas
+
+
+def draw_poses(poses: List[PoseResult], H, W, draw_body=True, draw_hand=True, draw_face=True):
+    """
+    Draw the detected poses on an empty canvas.
+
+    Args:
+        poses (List[PoseResult]): A list of PoseResult objects containing the detected poses.
+        H (int): The height of the canvas.
+        W (int): The width of the canvas.
+        draw_body (bool, optional): Whether to draw body keypoints. Defaults to True.
+        draw_hand (bool, optional): Whether to draw hand keypoints. Defaults to True.
+        draw_face (bool, optional): Whether to draw face keypoints. Defaults to True.
+
+    Returns:
+        numpy.ndarray: A 3D numpy array representing the canvas with the drawn poses.
+    """
+    canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+
+    for pose in poses:
+        if draw_body:
+            canvas = util.draw_bodypose(canvas, pose.body.keypoints)
+
+        if draw_hand:
+            canvas = util.draw_handpose(canvas, pose.left_hand)
+            canvas = util.draw_handpose(canvas, pose.right_hand)
+
+        if draw_face:
+            canvas = util.draw_facepose(canvas, pose.face)
+
+    return canvas
+
+
+def decode_json_as_poses(
+    pose_json: dict,
+) -> Tuple[List[PoseResult], List[AnimalPoseResult], int, int]:
+    """Decode the json_string complying with the openpose JSON output format
+    to poses that controlnet recognizes.
+    https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md
+
+    Args:
+        json_string: The json string to decode.
+
+    Returns:
+        human_poses
+        animal_poses
+        canvas_height
+        canvas_width
+    """
+    height = pose_json["canvas_height"]
+    width = pose_json["canvas_width"]
+
+    def chunks(lst, n):
+        """Yield successive n-sized chunks from lst."""
+        for i in range(0, len(lst), n):
+            yield lst[i : i + n]
+
+    def decompress_keypoints(
+        numbers: Optional[List[float]],
+    ) -> Optional[List[Optional[Keypoint]]]:
+        if not numbers:
+            return None
+
+        assert len(numbers) % 3 == 0
+
+        def create_keypoint(x, y, c):
+            if c < 1.0:
+                return None
+            keypoint = Keypoint(x, y)
+            return keypoint
+
+        return [create_keypoint(x, y, c) for x, y, c in chunks(numbers, n=3)]
+
+    return (
+        [
+            PoseResult(
+                body=BodyResult(
+                    keypoints=decompress_keypoints(pose.get("pose_keypoints_2d"))
+                ),
+                left_hand=decompress_keypoints(pose.get("hand_left_keypoints_2d")),
+                right_hand=decompress_keypoints(pose.get("hand_right_keypoints_2d")),
+                face=decompress_keypoints(pose.get("face_keypoints_2d")),
+            )
+            for pose in pose_json.get("people", [])
+        ],
+        [decompress_keypoints(pose) for pose in pose_json.get("animals", [])],
+        height,
+        width,
+    )
+
+
+def encode_poses_as_dict(poses: List[PoseResult], canvas_height: int, canvas_width: int) -> str:
+    """ Encode the pose as a dict following openpose JSON output format:
+    https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/doc/02_output.md
+    """
+    def compress_keypoints(keypoints: Union[List[Keypoint], None]) -> Union[List[float], None]:
+        if not keypoints:
+            return None
+        
+        return [
+            value
+            for keypoint in keypoints
+            for value in (
+                [float(keypoint.x), float(keypoint.y), 1.0]
+                if keypoint is not None
+                else [0.0, 0.0, 0.0]
+            )
+        ]
+
+    return {
+        'people': [
+            {
+                'pose_keypoints_2d': compress_keypoints(pose.body.keypoints),
+                "face_keypoints_2d": compress_keypoints(pose.face),
+                "hand_left_keypoints_2d": compress_keypoints(pose.left_hand),
+                "hand_right_keypoints_2d":compress_keypoints(pose.right_hand),
+            }
+            for pose in poses
+        ],
+        'canvas_height': canvas_height,
+        'canvas_width': canvas_width,
+    }
+
+global_cached_dwpose = Wholebody()
+
+class DwposeDetector:
+    """
+    A class for detecting human poses in images using the Dwpose model.
+
+    Attributes:
+        model_dir (str): Path to the directory where the pose models are stored.
+    """
+    def __init__(self, dw_pose_estimation):
+        self.dw_pose_estimation = dw_pose_estimation
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path, pretrained_det_model_or_path=None, det_filename=None, pose_filename=None, torchscript_device="cuda"):
+        global global_cached_dwpose
+        pretrained_det_model_or_path = pretrained_det_model_or_path or pretrained_model_or_path
+        det_filename = det_filename or "yolox_l.onnx"
+        pose_filename = pose_filename or "dw-ll_ucoco_384.onnx"
+        det_model_path = custom_hf_download(pretrained_det_model_or_path, det_filename)
+        pose_model_path = custom_hf_download(pretrained_model_or_path, pose_filename)
+        
+        print(f"\nDWPose: Using {det_filename} for bbox detection and {pose_filename} for pose estimation")
+        if global_cached_dwpose.det is None or global_cached_dwpose.det_filename != det_filename:
+            t = Wholebody(det_model_path, None, torchscript_device=torchscript_device)
+            t.pose = global_cached_dwpose.pose
+            t.pose_filename = global_cached_dwpose.pose
+            global_cached_dwpose = t
+        
+        if global_cached_dwpose.pose is None or global_cached_dwpose.pose_filename != pose_filename:
+            t = Wholebody(None, pose_model_path, torchscript_device=torchscript_device)
+            t.det = global_cached_dwpose.det
+            t.det_filename = global_cached_dwpose.det_filename
+            global_cached_dwpose = t
+        return cls(global_cached_dwpose)
+
+    def detect_poses(self, oriImg) -> List[PoseResult]:
+        with torch.no_grad():
+            keypoints_info = self.dw_pose_estimation(oriImg.copy())
+            return Wholebody.format_result(keypoints_info)
+    
+    def __call__(self, input_image, detect_resolution=512, include_body=True, include_hand=False, include_face=False, hand_and_face=None, output_type="pil", image_and_json=False, upscale_method="INTER_CUBIC", **kwargs):
+        if hand_and_face is not None:
+            warnings.warn("hand_and_face is deprecated. Use include_hand and include_face instead.", DeprecationWarning)
+            include_hand = hand_and_face
+            include_face = hand_and_face
+
+        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
+        poses = self.detect_poses(input_image)
+        
+        canvas = draw_poses(poses, input_image.shape[0], input_image.shape[1], draw_body=include_body, draw_hand=include_hand, draw_face=include_face)
+        canvas, remove_pad = resize_image_with_pad(canvas, detect_resolution, upscale_method)
+        detected_map = HWC3(remove_pad(canvas))
+
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
+        
+        if image_and_json:
+            return (detected_map, encode_poses_as_dict(poses, input_image.shape[0], input_image.shape[1]))
+        
+        return detected_map
+
+global_cached_animalpose = AnimalPoseImage()
+class AnimalposeDetector:
+    """
+    A class for detecting animal poses in images using the RTMPose AP10k model.
+
+    Attributes:
+        model_dir (str): Path to the directory where the pose models are stored.
+    """
+    def __init__(self, animal_pose_estimation):
+        self.animal_pose_estimation = animal_pose_estimation
+    
+    @classmethod
+    def from_pretrained(cls, pretrained_model_or_path, pretrained_det_model_or_path=None, det_filename="yolox_l.onnx", pose_filename="dw-ll_ucoco_384.onnx", torchscript_device="cuda"):
+        global global_cached_animalpose
+        det_model_path = custom_hf_download(pretrained_det_model_or_path, det_filename)
+        pose_model_path = custom_hf_download(pretrained_model_or_path, pose_filename)
+        
+        print(f"\nAnimalPose: Using {det_filename} for bbox detection and {pose_filename} for pose estimation")
+        if global_cached_animalpose.det is None or global_cached_animalpose.det_filename != det_filename:
+            t = AnimalPoseImage(det_model_path, None, torchscript_device=torchscript_device)
+            t.pose = global_cached_animalpose.pose
+            t.pose_filename = global_cached_animalpose.pose
+            global_cached_animalpose = t
+        
+        if global_cached_animalpose.pose is None or global_cached_animalpose.pose_filename != pose_filename:
+            t = AnimalPoseImage(None, pose_model_path, torchscript_device=torchscript_device)
+            t.det = global_cached_animalpose.det
+            t.det_filename = global_cached_animalpose.det_filename
+            global_cached_animalpose = t
+        return cls(global_cached_animalpose)
+    
+    def __call__(self, input_image, detect_resolution=512, output_type="pil", image_and_json=False, upscale_method="INTER_CUBIC", **kwargs):
+        input_image, output_type = common_input_validate(input_image, output_type, **kwargs)
+        input_image, remove_pad = resize_image_with_pad(input_image, detect_resolution, upscale_method)
+        result = self.animal_pose_estimation(input_image)
+        if result is None:
+            detected_map = np.zeros_like(input_image)
+            openpose_dict = {
+                'version': 'ap10k',
+                'animals': [],
+                'canvas_height': input_image.shape[0],
+                'canvas_width': input_image.shape[1]
+            }
+        else:
+            detected_map, openpose_dict = result
+        detected_map = remove_pad(detected_map)
+        if output_type == "pil":
+            detected_map = Image.fromarray(detected_map)
+        
+        if image_and_json:
+            return (detected_map, openpose_dict)
+
+        return detected_map

custom_nodes/ComfyUI-tbox/src/dwpose/animalpose.py

@@ -0,0 +1,273 @@
+import numpy as np
+import cv2
+import os
+import cv2
+from .dw_onnx.cv_ox_det import inference_detector as inference_onnx_yolox
+from .dw_onnx.cv_ox_yolo_nas import inference_detector as inference_onnx_yolo_nas
+from .dw_onnx.cv_ox_pose import inference_pose as inference_onnx_pose
+
+from .dw_torchscript.jit_det import inference_detector as inference_jit_yolox
+from .dw_torchscript.jit_pose import inference_pose as inference_jit_pose
+from typing import List, Optional
+from .types import PoseResult, BodyResult, Keypoint
+from timeit import default_timer
+from .util import guess_onnx_input_shape_dtype, get_ort_providers, get_model_type, is_model_torchscript
+import json
+import torch
+
+def drawBetweenKeypoints(pose_img, keypoints, indexes, color, scaleFactor):
+    ind0 = indexes[0] - 1
+    ind1 = indexes[1] - 1
+    
+    point1 = (keypoints[ind0][0], keypoints[ind0][1])
+    point2 = (keypoints[ind1][0], keypoints[ind1][1])
+
+    thickness = int(5 // scaleFactor)
+
+
+    cv2.line(pose_img, (int(point1[0]), int(point1[1])), (int(point2[0]), int(point2[1])), color, thickness)
+
+
+def drawBetweenKeypointsList(pose_img, keypoints, keypointPairsList, colorsList, scaleFactor):
+    for ind, keypointPair in enumerate(keypointPairsList):
+        drawBetweenKeypoints(pose_img, keypoints, keypointPair, colorsList[ind], scaleFactor)
+
+def drawBetweenSetofKeypointLists(pose_img, keypoints_set, keypointPairsList, colorsList, scaleFactor):
+    for keypoints in keypoints_set:
+        drawBetweenKeypointsList(pose_img, keypoints, keypointPairsList, colorsList, scaleFactor)
+
+
+def padImg(img, size, blackBorder=True):
+    left, right, top, bottom = 0, 0, 0, 0
+
+    # pad x
+    if img.shape[1] < size[1]:
+        sidePadding = int((size[1] - img.shape[1]) // 2)
+        left = sidePadding
+        right = sidePadding
+
+        # pad extra on right if padding needed is an odd number
+        if img.shape[1] % 2 == 1:
+            right += 1
+
+    # pad y
+    if img.shape[0] < size[0]:
+        topBottomPadding = int((size[0] - img.shape[0]) // 2)
+        top = topBottomPadding
+        bottom = topBottomPadding
+        
+        # pad extra on bottom if padding needed is an odd number
+        if img.shape[0] % 2 == 1:
+            bottom += 1
+
+    if blackBorder:
+        paddedImg = cv2.copyMakeBorder(src=img, top=top, bottom=bottom, left=left, right=right, borderType=cv2.BORDER_CONSTANT, value=(0,0,0))
+    else:
+        paddedImg = cv2.copyMakeBorder(src=img, top=top, bottom=bottom, left=left, right=right, borderType=cv2.BORDER_REPLICATE)
+
+    return paddedImg
+
+def smartCrop(img, size, center):
+
+    width = img.shape[1]
+    height = img.shape[0]
+    xSize = size[1]
+    ySize = size[0]
+    xCenter = center[0]
+    yCenter = center[1]
+
+    if img.shape[0] > size[0] or img.shape[1] > size[1]:
+
+
+        leftMargin = xCenter - xSize//2
+        rightMargin = xCenter + xSize//2
+        upMargin = yCenter - ySize//2
+        downMargin = yCenter + ySize//2
+
+
+        if(leftMargin < 0):
+            xCenter += (-leftMargin)
+        if(rightMargin > width):
+            xCenter -= (rightMargin - width)
+
+        if(upMargin < 0):
+            yCenter -= -upMargin
+        if(downMargin > height):
+            yCenter -= (downMargin - height)
+
+
+        img = cv2.getRectSubPix(img, size, (xCenter, yCenter))
+
+
+
+    return img
+
+
+
+def calculateScaleFactor(img, size, poseSpanX, poseSpanY):
+
+    poseSpanX = max(poseSpanX, size[0])
+
+    scaleFactorX = 1
+
+
+    if poseSpanX > size[0]:
+        scaleFactorX = size[0] / poseSpanX
+
+    scaleFactorY = 1
+    if poseSpanY > size[1]:
+        scaleFactorY = size[1] / poseSpanY
+
+    scaleFactor = min(scaleFactorX, scaleFactorY)
+
+
+    return scaleFactor
+
+
+
+def scaleImg(img, size, poseSpanX, poseSpanY, scaleFactor):
+    scaledImg = img
+
+    scaledImg = cv2.resize(img, (0, 0), fx=scaleFactor, fy=scaleFactor) 
+
+    return scaledImg, scaleFactor
+
+class AnimalPoseImage:
+    def __init__(self, det_model_path: Optional[str] = None, pose_model_path: Optional[str] = None, torchscript_device="cuda"):
+        self.det_filename = det_model_path and os.path.basename(det_model_path)
+        self.pose_filename = pose_model_path and os.path.basename(pose_model_path)
+        self.det, self.pose = None, None
+        # return type: None ort cv2 torchscript
+        self.det_model_type = get_model_type("AnimalPose",self.det_filename)
+        self.pose_model_type = get_model_type("AnimalPose",self.pose_filename)
+        # Always loads to CPU to avoid building OpenCV.
+        cv2_device = 'cpu'
+        cv2_backend = cv2.dnn.DNN_BACKEND_OPENCV if cv2_device == 'cpu' else cv2.dnn.DNN_BACKEND_CUDA
+        # You need to manually build OpenCV through cmake to work with your GPU.
+        cv2_providers = cv2.dnn.DNN_TARGET_CPU if cv2_device == 'cpu' else cv2.dnn.DNN_TARGET_CUDA
+        ort_providers = get_ort_providers()
+
+        if self.det_model_type is None:
+            pass
+        elif self.det_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.det = ort.InferenceSession(det_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.det.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+        elif self.det_model_type == "cv2":
+            try:
+                self.det = cv2.dnn.readNetFromONNX(det_model_path)
+                self.det.setPreferableBackend(cv2_backend)
+                self.det.setPreferableTarget(cv2_providers)
+            except:
+                print("TopK operators may not work on your OpenCV, try use onnxruntime with CPUExecutionProvider")
+                try:
+                    import onnxruntime as ort
+                    self.det = ort.InferenceSession(det_model_path, providers=["CPUExecutionProvider"])
+                except:
+                    print(f"Failed to load {det_model_path}, you can use other models instead")
+        else:
+            self.det = torch.jit.load(det_model_path)
+            self.det.to(torchscript_device)
+
+        if self.pose_model_type is None:
+            pass
+        elif self.pose_model_type == "ort":
+            try:
+                import onnxruntime as ort
+                self.pose = ort.InferenceSession(pose_model_path, providers=ort_providers)
+            except:
+                print(f"Failed to load onnxruntime with {self.pose.get_providers()}.\nPlease change EP_list in the config.yaml and restart ComfyUI")
+                self.pose = ort.InferenceSession(pose_model_path, providers=["CPUExecutionProvider"])
+        elif self.pose_model_type == "cv2":
+            self.pose = cv2.dnn.readNetFromONNX(pose_model_path)
+            self.pose.setPreferableBackend(cv2_backend)
+            self.pose.setPreferableTarget(cv2_providers)
+        else:
+            self.pose = torch.jit.load(pose_model_path)
+            self.pose.to(torchscript_device)
+
+        if self.pose_filename is not None:
+            self.pose_input_size, _ = guess_onnx_input_shape_dtype(self.pose_filename)
+    
+    def __call__(self, oriImg):
+        import torch.utils.benchmark.utils.timer as torch_timer
+        detect_classes = list(range(14, 23 + 1)) #https://github.com/ultralytics/ultralytics/blob/main/ultralytics/cfg/datasets/coco.yaml
+
+        if is_model_torchscript(self.det):
+            det_start = torch_timer.timer()
+            det_result = inference_jit_yolox(self.det, oriImg, detect_classes=detect_classes)
+            print(f"AnimalPose: Bbox {((torch_timer.timer() - det_start) * 1000):.2f}ms")
+        else:
+            det_start = default_timer()
+            det_onnx_dtype = np.float32 if "yolox" in self.det_filename else np.uint8
+            if "yolox" in self.det_filename:
+                det_result = inference_onnx_yolox(self.det, oriImg, detect_classes=detect_classes, dtype=det_onnx_dtype)
+            else:
+                #FP16 and INT8 YOLO NAS accept uint8 input
+                det_result = inference_onnx_yolo_nas(self.det, oriImg, detect_classes=detect_classes, dtype=det_onnx_dtype)
+            print(f"AnimalPose: Bbox {((default_timer() - det_start) * 1000):.2f}ms")
+        if (det_result is None) or (det_result.shape[0] == 0):
+            openpose_dict = {
+                'version': 'ap10k',
+                'animals': [],
+                'canvas_height': oriImg.shape[0],
+                'canvas_width': oriImg.shape[1]
+            }
+            return np.zeros_like(oriImg), openpose_dict
+        
+        if is_model_torchscript(self.pose):
+            pose_start = torch_timer.timer()
+            keypoint_sets, scores = inference_jit_pose(self.pose, det_result, oriImg, self.pose_input_size)
+            print(f"AnimalPose: Pose {((torch_timer.timer() - pose_start) * 1000):.2f}ms on {det_result.shape[0]} animals\n")
+        else:
+            pose_start = default_timer()
+            _, pose_onnx_dtype = guess_onnx_input_shape_dtype(self.pose_filename)
+            keypoint_sets, scores = inference_onnx_pose(self.pose, det_result, oriImg, self.pose_input_size, dtype=pose_onnx_dtype)
+            print(f"AnimalPose: Pose {((default_timer() - pose_start) * 1000):.2f}ms on {det_result.shape[0]} animals\n")
+
+        animal_kps_scores = []
+        pose_img = np.zeros((oriImg.shape[0], oriImg.shape[1], 3), dtype = np.uint8)
+        for (idx, keypoints) in enumerate(keypoint_sets):
+            # don't use keypoints that go outside the frame in calculations for the center
+            interorKeypoints = keypoints[((keypoints[:,0] > 0) & (keypoints[:,0] < oriImg.shape[1])) & ((keypoints[:,1] > 0) & (keypoints[:,1] < oriImg.shape[0]))]
+
+            xVals = interorKeypoints[:,0]
+            yVals = interorKeypoints[:,1]
+
+            minX = np.amin(xVals)
+            minY = np.amin(yVals)
+            maxX = np.amax(xVals)
+            maxY = np.amax(yVals)
+
+            poseSpanX = maxX - minX
+            poseSpanY = maxY - minY
+
+            # find mean center
+
+            xSum = np.sum(xVals)
+            ySum = np.sum(yVals)
+
+            xCenter = xSum // xVals.shape[0]
+            yCenter = ySum // yVals.shape[0]
+            center_of_keypoints = (xCenter,yCenter)
+
+            # order of the keypoints for AP10k and a standardized list of colors for limbs
+            keypointPairsList = [(1,2), (2,3), (1,3), (3,4), (4,9), (9,10), (10,11), (4,6), (6,7), (7,8), (4,5), (5,15), (15,16), (16,17), (5,12), (12,13), (13,14)]
+            colorsList = [(255,255,255), (100,255,100), (150,255,255), (100,50,255), (50,150,200), (0,255,255), (0,150,0), (0,0,255), (0,0,150), (255,50,255), (255,0,255), (255,0,0), (150,0,0), (255,255,100), (0,150,0), (255,255,0), (150,150,150)] # 16 colors needed
+
+            drawBetweenKeypointsList(pose_img, keypoints, keypointPairsList, colorsList, scaleFactor=1.0)
+            score = scores[idx, ..., None]
+            score[score > 1.0] = 1.0
+            score[score < 0.0] = 0.0
+            animal_kps_scores.append(np.concatenate((keypoints, score), axis=-1))
+        
+        openpose_dict = {
+            'version': 'ap10k',
+            'animals': [keypoints.tolist() for keypoints in animal_kps_scores],
+            'canvas_height': oriImg.shape[0],
+            'canvas_width': oriImg.shape[1]
+        }
+        return pose_img, openpose_dict

custom_nodes/ComfyUI-tbox/src/dwpose/body.py

@@ -0,0 +1,261 @@
+import cv2
+import numpy as np
+import math
+import time
+from scipy.ndimage.filters import gaussian_filter
+import matplotlib.pyplot as plt
+import matplotlib
+import torch
+from torchvision import transforms
+from typing import NamedTuple, List, Union
+
+from . import util
+from .model import bodypose_model
+from .types import Keypoint, BodyResult
+
+class Body(object):
+    def __init__(self, model_path):
+        self.model = bodypose_model()
+        # if torch.cuda.is_available():
+        #     self.model = self.model.cuda()
+            # print('cuda')
+        model_dict = util.transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImg):
+        # scale_search = [0.5, 1.0, 1.5, 2.0]
+        scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre1 = 0.1
+        thre2 = 0.05
+        multiplier = [x * boxsize / oriImg.shape[0] for x in scale_search]
+        heatmap_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 19))
+        paf_avg = np.zeros((oriImg.shape[0], oriImg.shape[1], 38))
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = util.smart_resize_k(oriImg, fx=scale, fy=scale)
+            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
+            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+            # data = data.permute([2, 0, 1]).unsqueeze(0).float()
+            with torch.no_grad():
+                data = data.to(self.cn_device)
+                Mconv7_stage6_L1, Mconv7_stage6_L2 = self.model(data)
+            Mconv7_stage6_L1 = Mconv7_stage6_L1.cpu().numpy()
+            Mconv7_stage6_L2 = Mconv7_stage6_L2.cpu().numpy()
+
+            # extract outputs, resize, and remove padding
+            # heatmap = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[1]].data), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = np.transpose(np.squeeze(Mconv7_stage6_L2), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
+            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            heatmap = util.smart_resize(heatmap, (oriImg.shape[0], oriImg.shape[1]))
+
+            # paf = np.transpose(np.squeeze(net.blobs[output_blobs.keys()[0]].data), (1, 2, 0))  # output 0 is PAFs
+            paf = np.transpose(np.squeeze(Mconv7_stage6_L1), (1, 2, 0))  # output 0 is PAFs
+            paf = util.smart_resize_k(paf, fx=stride, fy=stride)
+            paf = paf[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            paf = util.smart_resize(paf, (oriImg.shape[0], oriImg.shape[1]))
+
+            heatmap_avg += heatmap_avg + heatmap / len(multiplier)
+            paf_avg += + paf / len(multiplier)
+
+        all_peaks = []
+        peak_counter = 0
+
+        for part in range(18):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+
+            map_left = np.zeros(one_heatmap.shape)
+            map_left[1:, :] = one_heatmap[:-1, :]
+            map_right = np.zeros(one_heatmap.shape)
+            map_right[:-1, :] = one_heatmap[1:, :]
+            map_up = np.zeros(one_heatmap.shape)
+            map_up[:, 1:] = one_heatmap[:, :-1]
+            map_down = np.zeros(one_heatmap.shape)
+            map_down[:, :-1] = one_heatmap[:, 1:]
+
+            peaks_binary = np.logical_and.reduce(
+                (one_heatmap >= map_left, one_heatmap >= map_right, one_heatmap >= map_up, one_heatmap >= map_down, one_heatmap > thre1))
+            peaks = list(zip(np.nonzero(peaks_binary)[1], np.nonzero(peaks_binary)[0]))  # note reverse
+            peaks_with_score = [x + (map_ori[x[1], x[0]],) for x in peaks]
+            peak_id = range(peak_counter, peak_counter + len(peaks))
+            peaks_with_score_and_id = [peaks_with_score[i] + (peak_id[i],) for i in range(len(peak_id))]
+
+            all_peaks.append(peaks_with_score_and_id)
+            peak_counter += len(peaks)
+
+        # find connection in the specified sequence, center 29 is in the position 15
+        limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \
+                   [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \
+                   [1, 16], [16, 18], [3, 17], [6, 18]]
+        # the middle joints heatmap correpondence
+        mapIdx = [[31, 32], [39, 40], [33, 34], [35, 36], [41, 42], [43, 44], [19, 20], [21, 22], \
+                  [23, 24], [25, 26], [27, 28], [29, 30], [47, 48], [49, 50], [53, 54], [51, 52], \
+                  [55, 56], [37, 38], [45, 46]]
+
+        connection_all = []
+        special_k = []
+        mid_num = 10
+
+        for k in range(len(mapIdx)):
+            score_mid = paf_avg[:, :, [x - 19 for x in mapIdx[k]]]
+            candA = all_peaks[limbSeq[k][0] - 1]
+            candB = all_peaks[limbSeq[k][1] - 1]
+            nA = len(candA)
+            nB = len(candB)
+            indexA, indexB = limbSeq[k]
+            if (nA != 0 and nB != 0):
+                connection_candidate = []
+                for i in range(nA):
+                    for j in range(nB):
+                        vec = np.subtract(candB[j][:2], candA[i][:2])
+                        norm = math.sqrt(vec[0] * vec[0] + vec[1] * vec[1])
+                        norm = max(0.001, norm)
+                        vec = np.divide(vec, norm)
+
+                        startend = list(zip(np.linspace(candA[i][0], candB[j][0], num=mid_num), \
+                                            np.linspace(candA[i][1], candB[j][1], num=mid_num)))
+
+                        vec_x = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 0] \
+                                          for I in range(len(startend))])
+                        vec_y = np.array([score_mid[int(round(startend[I][1])), int(round(startend[I][0])), 1] \
+                                          for I in range(len(startend))])
+
+                        score_midpts = np.multiply(vec_x, vec[0]) + np.multiply(vec_y, vec[1])
+                        score_with_dist_prior = sum(score_midpts) / len(score_midpts) + min(
+                            0.5 * oriImg.shape[0] / norm - 1, 0)
+                        criterion1 = len(np.nonzero(score_midpts > thre2)[0]) > 0.8 * len(score_midpts)
+                        criterion2 = score_with_dist_prior > 0
+                        if criterion1 and criterion2:
+                            connection_candidate.append(
+                                [i, j, score_with_dist_prior, score_with_dist_prior + candA[i][2] + candB[j][2]])
+
+                connection_candidate = sorted(connection_candidate, key=lambda x: x[2], reverse=True)
+                connection = np.zeros((0, 5))
+                for c in range(len(connection_candidate)):
+                    i, j, s = connection_candidate[c][0:3]
+                    if (i not in connection[:, 3] and j not in connection[:, 4]):
+                        connection = np.vstack([connection, [candA[i][3], candB[j][3], s, i, j]])
+                        if (len(connection) >= min(nA, nB)):
+                            break
+
+                connection_all.append(connection)
+            else:
+                special_k.append(k)
+                connection_all.append([])
+
+        # last number in each row is the total parts number of that person
+        # the second last number in each row is the score of the overall configuration
+        subset = -1 * np.ones((0, 20))
+        candidate = np.array([item for sublist in all_peaks for item in sublist])
+
+        for k in range(len(mapIdx)):
+            if k not in special_k:
+                partAs = connection_all[k][:, 0]
+                partBs = connection_all[k][:, 1]
+                indexA, indexB = np.array(limbSeq[k]) - 1
+
+                for i in range(len(connection_all[k])):  # = 1:size(temp,1)
+                    found = 0
+                    subset_idx = [-1, -1]
+                    for j in range(len(subset)):  # 1:size(subset,1):
+                        if subset[j][indexA] == partAs[i] or subset[j][indexB] == partBs[i]:
+                            subset_idx[found] = j
+                            found += 1
+
+                    if found == 1:
+                        j = subset_idx[0]
+                        if subset[j][indexB] != partBs[i]:
+                            subset[j][indexB] = partBs[i]
+                            subset[j][-1] += 1
+                            subset[j][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
+                    elif found == 2:  # if found 2 and disjoint, merge them
+                        j1, j2 = subset_idx
+                        membership = ((subset[j1] >= 0).astype(int) + (subset[j2] >= 0).astype(int))[:-2]
+                        if len(np.nonzero(membership == 2)[0]) == 0:  # merge
+                            subset[j1][:-2] += (subset[j2][:-2] + 1)
+                            subset[j1][-2:] += subset[j2][-2:]
+                            subset[j1][-2] += connection_all[k][i][2]
+                            subset = np.delete(subset, j2, 0)
+                        else:  # as like found == 1
+                            subset[j1][indexB] = partBs[i]
+                            subset[j1][-1] += 1
+                            subset[j1][-2] += candidate[partBs[i].astype(int), 2] + connection_all[k][i][2]
+
+                    # if find no partA in the subset, create a new subset
+                    elif not found and k < 17:
+                        row = -1 * np.ones(20)
+                        row[indexA] = partAs[i]
+                        row[indexB] = partBs[i]
+                        row[-1] = 2
+                        row[-2] = sum(candidate[connection_all[k][i, :2].astype(int), 2]) + connection_all[k][i][2]
+                        subset = np.vstack([subset, row])
+        # delete some rows of subset which has few parts occur
+        deleteIdx = []
+        for i in range(len(subset)):
+            if subset[i][-1] < 4 or subset[i][-2] / subset[i][-1] < 0.4:
+                deleteIdx.append(i)
+        subset = np.delete(subset, deleteIdx, axis=0)
+
+        # subset: n*20 array, 0-17 is the index in candidate, 18 is the total score, 19 is the total parts
+        # candidate: x, y, score, id
+        return candidate, subset
+    
+    @staticmethod
+    def format_body_result(candidate: np.ndarray, subset: np.ndarray) -> List[BodyResult]:
+        """
+        Format the body results from the candidate and subset arrays into a list of BodyResult objects.
+        
+        Args:
+            candidate (np.ndarray): An array of candidates containing the x, y coordinates, score, and id
+                for each body part.
+            subset (np.ndarray): An array of subsets containing indices to the candidate array for each
+                person detected. The last two columns of each row hold the total score and total parts
+                of the person.
+
+        Returns:
+            List[BodyResult]: A list of BodyResult objects, where each object represents a person with
+                detected keypoints, total score, and total parts.
+        """
+        return [
+            BodyResult(
+                keypoints=[
+                    Keypoint(
+                        x=candidate[candidate_index][0],
+                        y=candidate[candidate_index][1],
+                        score=candidate[candidate_index][2],
+                        id=candidate[candidate_index][3]
+                    ) if candidate_index != -1 else None
+                    for candidate_index in person[:18].astype(int)
+                ],
+                total_score=person[18],
+                total_parts=person[19]
+            )
+            for person in subset
+        ]
+    
+
+if __name__ == "__main__":
+    body_estimation = Body('../model/body_pose_model.pth')
+
+    test_image = '../images/ski.jpg'
+    oriImg = cv2.imread(test_image)  # B,G,R order
+    candidate, subset = body_estimation(oriImg)
+    bodies = body_estimation.format_body_result(candidate, subset)
+
+    canvas = oriImg
+    for body in bodies:
+        canvas = util.draw_bodypose(canvas, body)
+        
+    plt.imshow(canvas[:, :, [2, 1, 0]])
+    plt.show()

custom_nodes/ComfyUI-tbox/src/dwpose/dw_onnx/__init__.py

@@ -0,0 +1 @@
+#Dummy file ensuring this package will be recognized

custom_nodes/ComfyUI-tbox/src/dwpose/dw_onnx/cv_ox_det.py

@@ -0,0 +1,129 @@
+import cv2
+import numpy as np
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= nms_thr)[0]
+        order = order[inds + 1]
+
+    return keep
+
+def multiclass_nms(boxes, scores, nms_thr, score_thr):
+    """Multiclass NMS implemented in Numpy. Class-aware version."""
+    final_dets = []
+    num_classes = scores.shape[1]
+    for cls_ind in range(num_classes):
+        cls_scores = scores[:, cls_ind]
+        valid_score_mask = cls_scores > score_thr
+        if valid_score_mask.sum() == 0:
+            continue
+        else:
+            valid_scores = cls_scores[valid_score_mask]
+            valid_boxes = boxes[valid_score_mask]
+            keep = nms(valid_boxes, valid_scores, nms_thr)
+            if len(keep) > 0:
+                cls_inds = np.ones((len(keep), 1)) * cls_ind
+                dets = np.concatenate(
+                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
+                )
+                final_dets.append(dets)
+    if len(final_dets) == 0:
+        return None
+    return np.concatenate(final_dets, 0)
+
+def demo_postprocess(outputs, img_size, p6=False):
+    grids = []
+    expanded_strides = []
+    strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
+
+    hsizes = [img_size[0] // stride for stride in strides]
+    wsizes = [img_size[1] // stride for stride in strides]
+
+    for hsize, wsize, stride in zip(hsizes, wsizes, strides):
+        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
+        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+        grids.append(grid)
+        shape = grid.shape[:2]
+        expanded_strides.append(np.full((*shape, 1), stride))
+
+    grids = np.concatenate(grids, 1)
+    expanded_strides = np.concatenate(expanded_strides, 1)
+    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
+    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
+
+    return outputs
+
+def preprocess(img, input_size, swap=(2, 0, 1)):
+    if len(img.shape) == 3:
+        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+    else:
+        padded_img = np.ones(input_size, dtype=np.uint8) * 114
+
+    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+    resized_img = cv2.resize(
+        img,
+        (int(img.shape[1] * r), int(img.shape[0] * r)),
+        interpolation=cv2.INTER_LINEAR,
+    ).astype(np.uint8)
+    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+    padded_img = padded_img.transpose(swap)
+    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+    return padded_img, r
+
+def inference_detector(session, oriImg, detect_classes=[0], dtype=np.float32):
+    input_shape = (640,640)
+    img, ratio = preprocess(oriImg, input_shape)
+
+    input = img[None, :, :, :]
+    input = input.astype(dtype)
+    if "InferenceSession" in type(session).__name__:
+        input_name = session.get_inputs()[0].name
+        output = session.run(None, {input_name: input})
+    else:
+        outNames = session.getUnconnectedOutLayersNames()
+        session.setInput(input)
+        output = session.forward(outNames)
+
+    predictions = demo_postprocess(output[0], input_shape)[0]
+
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4:5] * predictions[:, 5:]
+
+    boxes_xyxy = np.ones_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
+    boxes_xyxy /= ratio
+    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
+    if dets is None:
+        return None
+    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+    isscore = final_scores>0.3
+    iscat = np.isin(final_cls_inds, detect_classes)
+    isbbox = [ i and j for (i, j) in zip(isscore, iscat)]
+    final_boxes = final_boxes[isbbox]
+    return final_boxes

custom_nodes/ComfyUI-tbox/src/dwpose/dw_onnx/cv_ox_pose.py

@@ -0,0 +1,363 @@
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+
+def preprocess(
+    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Do preprocessing for DWPose model inference.
+
+    Args:
+        img (np.ndarray): Input image in shape.
+        input_size (tuple): Input image size in shape (w, h).
+
+    Returns:
+        tuple:
+        - resized_img (np.ndarray): Preprocessed image.
+        - center (np.ndarray): Center of image.
+        - scale (np.ndarray): Scale of image.
+    """
+    # get shape of image
+    img_shape = img.shape[:2]
+    out_img, out_center, out_scale = [], [], []
+    if len(out_bbox) == 0:
+        out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
+    for i in range(len(out_bbox)):
+        x0 = out_bbox[i][0]
+        y0 = out_bbox[i][1]
+        x1 = out_bbox[i][2]
+        y1 = out_bbox[i][3]
+        bbox = np.array([x0, y0, x1, y1])
+
+        # get center and scale
+        center, scale = bbox_xyxy2cs(bbox, padding=1.25)
+
+        # do affine transformation
+        resized_img, scale = top_down_affine(input_size, scale, center, img)
+
+        # normalize image
+        mean = np.array([123.675, 116.28, 103.53])
+        std = np.array([58.395, 57.12, 57.375])
+        resized_img = (resized_img - mean) / std
+
+        out_img.append(resized_img)
+        out_center.append(center)
+        out_scale.append(scale)
+
+    return out_img, out_center, out_scale
+
+
+def inference(sess, img, dtype=np.float32):
+    """Inference DWPose model. Processing all image segments at once to take advantage of GPU's parallelism ability if onnxruntime is installed
+
+    Args:
+        sess : ONNXRuntime session.
+        img : Input image in shape.
+
+    Returns:
+        outputs : Output of DWPose model.
+    """
+    all_out = []
+    # build input
+    input = np.stack(img, axis=0).transpose(0, 3, 1, 2)
+    input = input.astype(dtype)
+    if "InferenceSession" in type(sess).__name__:
+        input_name = sess.get_inputs()[0].name
+        all_outputs = sess.run(None, {input_name: input})
+        for batch_idx in range(len(all_outputs[0])):
+            outputs = [all_outputs[i][batch_idx:batch_idx+1,...] for i in range(len(all_outputs))]
+            all_out.append(outputs)
+        return all_out
+
+    #OpenCV doesn't support batch processing sadly
+    for i in range(len(img)):
+        input = img[i].transpose(2, 0, 1)
+        input = input[None, :, :, :]
+
+        outNames = sess.getUnconnectedOutLayersNames()
+        sess.setInput(input)
+        outputs = sess.forward(outNames)
+        all_out.append(outputs)
+
+    return all_out
+
+def postprocess(outputs: List[np.ndarray],
+                model_input_size: Tuple[int, int],
+                center: Tuple[int, int],
+                scale: Tuple[int, int],
+                simcc_split_ratio: float = 2.0
+                ) -> Tuple[np.ndarray, np.ndarray]:
+    """Postprocess for DWPose model output.
+
+    Args:
+        outputs (np.ndarray): Output of RTMPose model.
+        model_input_size (tuple): RTMPose model Input image size.
+        center (tuple): Center of bbox in shape (x, y).
+        scale (tuple): Scale of bbox in shape (w, h).
+        simcc_split_ratio (float): Split ratio of simcc.
+
+    Returns:
+        tuple:
+        - keypoints (np.ndarray): Rescaled keypoints.
+        - scores (np.ndarray): Model predict scores.
+    """
+    all_key = []
+    all_score = []
+    for i in range(len(outputs)):
+        # use simcc to decode
+        simcc_x, simcc_y = outputs[i]
+        keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
+
+        # rescale keypoints
+        keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
+        all_key.append(keypoints[0])
+        all_score.append(scores[0])
+
+    return np.array(all_key), np.array(all_score)
+
+
+def bbox_xyxy2cs(bbox: np.ndarray,
+                 padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]:
+    """Transform the bbox format from (x,y,w,h) into (center, scale)
+
+    Args:
+        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
+            as (left, top, right, bottom)
+        padding (float): BBox padding factor that will be multilied to scale.
+            Default: 1.0
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
+            (n, 2)
+        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
+            (n, 2)
+    """
+    # convert single bbox from (4, ) to (1, 4)
+    dim = bbox.ndim
+    if dim == 1:
+        bbox = bbox[None, :]
+
+    # get bbox center and scale
+    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
+    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
+    scale = np.hstack([x2 - x1, y2 - y1]) * padding
+
+    if dim == 1:
+        center = center[0]
+        scale = scale[0]
+
+    return center, scale
+
+
+def _fix_aspect_ratio(bbox_scale: np.ndarray,
+                      aspect_ratio: float) -> np.ndarray:
+    """Extend the scale to match the given aspect ratio.
+
+    Args:
+        scale (np.ndarray): The image scale (w, h) in shape (2, )
+        aspect_ratio (float): The ratio of ``w/h``
+
+    Returns:
+        np.ndarray: The reshaped image scale in (2, )
+    """
+    w, h = np.hsplit(bbox_scale, [1])
+    bbox_scale = np.where(w > h * aspect_ratio,
+                          np.hstack([w, w / aspect_ratio]),
+                          np.hstack([h * aspect_ratio, h]))
+    return bbox_scale
+
+
+def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
+    """Rotate a point by an angle.
+
+    Args:
+        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
+        angle_rad (float): rotation angle in radian
+
+    Returns:
+        np.ndarray: Rotated point in shape (2, )
+    """
+    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
+    rot_mat = np.array([[cs, -sn], [sn, cs]])
+    return rot_mat @ pt
+
+
+def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """To calculate the affine matrix, three pairs of points are required. This
+    function is used to get the 3rd point, given 2D points a & b.
+
+    The 3rd point is defined by rotating vector `a - b` by 90 degrees
+    anticlockwise, using b as the rotation center.
+
+    Args:
+        a (np.ndarray): The 1st point (x,y) in shape (2, )
+        b (np.ndarray): The 2nd point (x,y) in shape (2, )
+
+    Returns:
+        np.ndarray: The 3rd point.
+    """
+    direction = a - b
+    c = b + np.r_[-direction[1], direction[0]]
+    return c
+
+
+def get_warp_matrix(center: np.ndarray,
+                    scale: np.ndarray,
+                    rot: float,
+                    output_size: Tuple[int, int],
+                    shift: Tuple[float, float] = (0., 0.),
+                    inv: bool = False) -> np.ndarray:
+    """Calculate the affine transformation matrix that can warp the bbox area
+    in the input image to the output size.
+
+    Args:
+        center (np.ndarray[2, ]): Center of the bounding box (x, y).
+        scale (np.ndarray[2, ]): Scale of the bounding box
+            wrt [width, height].
+        rot (float): Rotation angle (degree).
+        output_size (np.ndarray[2, ] | list(2,)): Size of the
+            destination heatmaps.
+        shift (0-100%): Shift translation ratio wrt the width/height.
+            Default (0., 0.).
+        inv (bool): Option to inverse the affine transform direction.
+            (inv=False: src->dst or inv=True: dst->src)
+
+    Returns:
+        np.ndarray: A 2x3 transformation matrix
+    """
+    shift = np.array(shift)
+    src_w = scale[0]
+    dst_w = output_size[0]
+    dst_h = output_size[1]
+
+    # compute transformation matrix
+    rot_rad = np.deg2rad(rot)
+    src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad)
+    dst_dir = np.array([0., dst_w * -0.5])
+
+    # get four corners of the src rectangle in the original image
+    src = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale * shift
+    src[1, :] = center + src_dir + scale * shift
+    src[2, :] = _get_3rd_point(src[0, :], src[1, :])
+
+    # get four corners of the dst rectangle in the input image
+    dst = np.zeros((3, 2), dtype=np.float32)
+    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return warp_mat
+
+
+def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict,
+                    img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Get the bbox image as the model input by affine transform.
+
+    Args:
+        input_size (dict): The input size of the model.
+        bbox_scale (dict): The bbox scale of the img.
+        bbox_center (dict): The bbox center of the img.
+        img (np.ndarray): The original image.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: img after affine transform.
+        - np.ndarray[float32]: bbox scale after affine transform.
+    """
+    w, h = input_size
+    warp_size = (int(w), int(h))
+
+    # reshape bbox to fixed aspect ratio
+    bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
+
+    # get the affine matrix
+    center = bbox_center
+    scale = bbox_scale
+    rot = 0
+    warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
+
+    # do affine transform
+    img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
+
+    return img, bbox_scale
+
+
+def get_simcc_maximum(simcc_x: np.ndarray,
+                      simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Get maximum response location and value from simcc representations.
+
+    Note:
+        instance number: N
+        num_keypoints: K
+        heatmap height: H
+        heatmap width: W
+
+    Args:
+        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
+        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
+
+    Returns:
+        tuple:
+        - locs (np.ndarray): locations of maximum heatmap responses in shape
+            (K, 2) or (N, K, 2)
+        - vals (np.ndarray): values of maximum heatmap responses in shape
+            (K,) or (N, K)
+    """
+    N, K, Wx = simcc_x.shape
+    simcc_x = simcc_x.reshape(N * K, -1)
+    simcc_y = simcc_y.reshape(N * K, -1)
+
+    # get maximum value locations
+    x_locs = np.argmax(simcc_x, axis=1)
+    y_locs = np.argmax(simcc_y, axis=1)
+    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
+    max_val_x = np.amax(simcc_x, axis=1)
+    max_val_y = np.amax(simcc_y, axis=1)
+
+    # get maximum value across x and y axis
+    mask = max_val_x > max_val_y
+    max_val_x[mask] = max_val_y[mask]
+    vals = max_val_x
+    locs[vals <= 0.] = -1
+
+    # reshape
+    locs = locs.reshape(N, K, 2)
+    vals = vals.reshape(N, K)
+
+    return locs, vals
+
+
+def decode(simcc_x: np.ndarray, simcc_y: np.ndarray,
+           simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]:
+    """Modulate simcc distribution with Gaussian.
+
+    Args:
+        simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
+        simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
+        simcc_split_ratio (int): The split ratio of simcc.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
+        - np.ndarray[float32]: scores in shape (K,) or (n, K)
+    """
+    keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
+    keypoints /= simcc_split_ratio
+
+    return keypoints, scores
+
+
+def inference_pose(session, out_bbox, oriImg, model_input_size=(288, 384), dtype=np.float32):
+    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
+    outputs = inference(session, resized_img, dtype)
+    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
+
+    return keypoints, scores

custom_nodes/ComfyUI-tbox/src/dwpose/dw_onnx/cv_ox_yolo_nas.py

@@ -0,0 +1,60 @@
+# Source: https://github.com/Hyuto/yolo-nas-onnx/tree/master/yolo-nas-py
+# Inspired from: https://github.com/Deci-AI/super-gradients/blob/3.1.1/src/super_gradients/training/processing/processing.py
+
+import numpy as np
+import cv2
+
+def preprocess(img, input_size, swap=(2, 0, 1)):
+    if len(img.shape) == 3:
+        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+    else:
+        padded_img = np.ones(input_size, dtype=np.uint8) * 114
+
+    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+    resized_img = cv2.resize(
+        img,
+        (int(img.shape[1] * r), int(img.shape[0] * r)),
+        interpolation=cv2.INTER_LINEAR,
+    ).astype(np.uint8)
+    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+    padded_img = padded_img.transpose(swap)
+    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+    return padded_img, r
+
+def inference_detector(session, oriImg, detect_classes=[0], dtype=np.uint8):
+    """
+    This function is only compatible with onnx models exported from the new API with built-in NMS
+    ```py
+    from super_gradients.conversion.conversion_enums import ExportQuantizationMode
+    from super_gradients.common.object_names import Models
+    from super_gradients.training import models
+
+    model = models.get(Models.YOLO_NAS_L, pretrained_weights="coco")
+
+    export_result = model.export(
+        "yolo_nas/yolo_nas_l_fp16.onnx",
+        quantization_mode=ExportQuantizationMode.FP16,
+        device="cuda"
+    )
+    ```
+    """
+    input_shape = (640,640)
+    img, ratio = preprocess(oriImg, input_shape)
+    input = img[None, :, :, :]
+    input = input.astype(dtype)
+    if "InferenceSession" in type(session).__name__:
+        input_name = session.get_inputs()[0].name
+        output = session.run(None, {input_name: input})
+    else:
+        outNames = session.getUnconnectedOutLayersNames()
+        session.setInput(input)
+        output = session.forward(outNames)
+    num_preds, pred_boxes, pred_scores, pred_classes = output
+    num_preds = num_preds[0,0]
+    if num_preds == 0:
+        return None
+    idxs = np.where((np.isin(pred_classes[0, :num_preds], detect_classes)) & (pred_scores[0, :num_preds] > 0.3))
+    if (len(idxs) == 0) or (idxs[0].size == 0):
+        return None
+    return pred_boxes[0, idxs].squeeze(axis=0) / ratio

custom_nodes/ComfyUI-tbox/src/dwpose/dw_torchscript/__init__.py

@@ -0,0 +1 @@
+#Dummy file ensuring this package will be recognized

custom_nodes/ComfyUI-tbox/src/dwpose/dw_torchscript/jit_det.py

@@ -0,0 +1,125 @@
+import cv2
+import numpy as np
+import torch
+
+def nms(boxes, scores, nms_thr):
+    """Single class NMS implemented in Numpy."""
+    x1 = boxes[:, 0]
+    y1 = boxes[:, 1]
+    x2 = boxes[:, 2]
+    y2 = boxes[:, 3]
+
+    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1 + 1)
+        h = np.maximum(0.0, yy2 - yy1 + 1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+
+        inds = np.where(ovr <= nms_thr)[0]
+        order = order[inds + 1]
+
+    return keep
+
+def multiclass_nms(boxes, scores, nms_thr, score_thr):
+    """Multiclass NMS implemented in Numpy. Class-aware version."""
+    final_dets = []
+    num_classes = scores.shape[1]
+    for cls_ind in range(num_classes):
+        cls_scores = scores[:, cls_ind]
+        valid_score_mask = cls_scores > score_thr
+        if valid_score_mask.sum() == 0:
+            continue
+        else:
+            valid_scores = cls_scores[valid_score_mask]
+            valid_boxes = boxes[valid_score_mask]
+            keep = nms(valid_boxes, valid_scores, nms_thr)
+            if len(keep) > 0:
+                cls_inds = np.ones((len(keep), 1)) * cls_ind
+                dets = np.concatenate(
+                    [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1
+                )
+                final_dets.append(dets)
+    if len(final_dets) == 0:
+        return None
+    return np.concatenate(final_dets, 0)
+
+def demo_postprocess(outputs, img_size, p6=False):
+    grids = []
+    expanded_strides = []
+    strides = [8, 16, 32] if not p6 else [8, 16, 32, 64]
+
+    hsizes = [img_size[0] // stride for stride in strides]
+    wsizes = [img_size[1] // stride for stride in strides]
+
+    for hsize, wsize, stride in zip(hsizes, wsizes, strides):
+        xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize))
+        grid = np.stack((xv, yv), 2).reshape(1, -1, 2)
+        grids.append(grid)
+        shape = grid.shape[:2]
+        expanded_strides.append(np.full((*shape, 1), stride))
+
+    grids = np.concatenate(grids, 1)
+    expanded_strides = np.concatenate(expanded_strides, 1)
+    outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides
+    outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides
+
+    return outputs
+
+def preprocess(img, input_size, swap=(2, 0, 1)):
+    if len(img.shape) == 3:
+        padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114
+    else:
+        padded_img = np.ones(input_size, dtype=np.uint8) * 114
+
+    r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1])
+    resized_img = cv2.resize(
+        img,
+        (int(img.shape[1] * r), int(img.shape[0] * r)),
+        interpolation=cv2.INTER_LINEAR,
+    ).astype(np.uint8)
+    padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img
+
+    padded_img = padded_img.transpose(swap)
+    padded_img = np.ascontiguousarray(padded_img, dtype=np.float32)
+    return padded_img, r
+
+def inference_detector(model, oriImg, detect_classes=[0]):
+    input_shape = (640,640)
+    img, ratio = preprocess(oriImg, input_shape)
+
+    device, dtype = next(model.parameters()).device, next(model.parameters()).dtype
+    input = img[None, :, :, :]
+    input = torch.from_numpy(input).to(device, dtype)
+
+    output = model(input).float().cpu().detach().numpy()
+    predictions = demo_postprocess(output[0], input_shape)
+
+    boxes = predictions[:, :4]
+    scores = predictions[:, 4:5] * predictions[:, 5:]
+
+    boxes_xyxy = np.ones_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
+    boxes_xyxy /= ratio
+    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
+    if dets is None:
+        return None
+    final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
+    isscore = final_scores>0.3
+    iscat = np.isin(final_cls_inds, detect_classes)
+    isbbox = [ i and j for (i, j) in zip(isscore, iscat)]
+    final_boxes = final_boxes[isbbox]
+    return final_boxes

custom_nodes/ComfyUI-tbox/src/dwpose/dw_torchscript/jit_pose.py

@@ -0,0 +1,363 @@
+from typing import List, Tuple
+
+import cv2
+import numpy as np
+import torch
+
+def preprocess(
+    img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256)
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Do preprocessing for DWPose model inference.
+
+    Args:
+        img (np.ndarray): Input image in shape.
+        input_size (tuple): Input image size in shape (w, h).
+
+    Returns:
+        tuple:
+        - resized_img (np.ndarray): Preprocessed image.
+        - center (np.ndarray): Center of image.
+        - scale (np.ndarray): Scale of image.
+    """
+    # get shape of image
+    img_shape = img.shape[:2]
+    out_img, out_center, out_scale = [], [], []
+    if len(out_bbox) == 0:
+        out_bbox = [[0, 0, img_shape[1], img_shape[0]]]
+    for i in range(len(out_bbox)):
+        x0 = out_bbox[i][0]
+        y0 = out_bbox[i][1]
+        x1 = out_bbox[i][2]
+        y1 = out_bbox[i][3]
+        bbox = np.array([x0, y0, x1, y1])
+
+        # get center and scale
+        center, scale = bbox_xyxy2cs(bbox, padding=1.25)
+
+        # do affine transformation
+        resized_img, scale = top_down_affine(input_size, scale, center, img)
+
+        # normalize image
+        mean = np.array([123.675, 116.28, 103.53])
+        std = np.array([58.395, 57.12, 57.375])
+        resized_img = (resized_img - mean) / std
+
+        out_img.append(resized_img)
+        out_center.append(center)
+        out_scale.append(scale)
+
+    return out_img, out_center, out_scale
+
+def inference(model, img, bs=5):
+    """Inference DWPose model implemented in TorchScript.
+
+    Args:
+        model : TorchScript Model.
+        img : Input image in shape.
+
+    Returns:
+        outputs : Output of DWPose model.
+    """
+    all_out = []
+    # build input
+    orig_img_count = len(img)
+    #Pad zeros to fit batch size
+    for _ in range(bs - (orig_img_count % bs)):
+        img.append(np.zeros_like(img[0]))
+    input = np.stack(img, axis=0).transpose(0, 3, 1, 2)
+    device, dtype = next(model.parameters()).device, next(model.parameters()).dtype
+    input = torch.from_numpy(input).to(device, dtype)
+
+    out1, out2 = [], []
+    for i in range(input.shape[0] // bs):
+        curr_batch_output = model(input[i*bs:(i+1)*bs])
+        out1.append(curr_batch_output[0].float())
+        out2.append(curr_batch_output[1].float())
+    out1, out2 = torch.cat(out1, dim=0)[:orig_img_count], torch.cat(out2, dim=0)[:orig_img_count]
+    out1, out2 = out1.float().cpu().detach().numpy(), out2.float().cpu().detach().numpy()
+    all_outputs = out1, out2
+
+    for batch_idx in range(len(all_outputs[0])):
+        outputs = [all_outputs[i][batch_idx:batch_idx+1,...] for i in range(len(all_outputs))]
+        all_out.append(outputs)
+    return all_out
+def postprocess(outputs: List[np.ndarray],
+                model_input_size: Tuple[int, int],
+                center: Tuple[int, int],
+                scale: Tuple[int, int],
+                simcc_split_ratio: float = 2.0
+                ) -> Tuple[np.ndarray, np.ndarray]:
+    """Postprocess for DWPose model output.
+
+    Args:
+        outputs (np.ndarray): Output of RTMPose model.
+        model_input_size (tuple): RTMPose model Input image size.
+        center (tuple): Center of bbox in shape (x, y).
+        scale (tuple): Scale of bbox in shape (w, h).
+        simcc_split_ratio (float): Split ratio of simcc.
+
+    Returns:
+        tuple:
+        - keypoints (np.ndarray): Rescaled keypoints.
+        - scores (np.ndarray): Model predict scores.
+    """
+    all_key = []
+    all_score = []
+    for i in range(len(outputs)):
+        # use simcc to decode
+        simcc_x, simcc_y = outputs[i]
+        keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio)
+
+        # rescale keypoints
+        keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2
+        all_key.append(keypoints[0])
+        all_score.append(scores[0])
+
+    return np.array(all_key), np.array(all_score)
+
+
+def bbox_xyxy2cs(bbox: np.ndarray,
+                 padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]:
+    """Transform the bbox format from (x,y,w,h) into (center, scale)
+
+    Args:
+        bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted
+            as (left, top, right, bottom)
+        padding (float): BBox padding factor that will be multilied to scale.
+            Default: 1.0
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or
+            (n, 2)
+        - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or
+            (n, 2)
+    """
+    # convert single bbox from (4, ) to (1, 4)
+    dim = bbox.ndim
+    if dim == 1:
+        bbox = bbox[None, :]
+
+    # get bbox center and scale
+    x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3])
+    center = np.hstack([x1 + x2, y1 + y2]) * 0.5
+    scale = np.hstack([x2 - x1, y2 - y1]) * padding
+
+    if dim == 1:
+        center = center[0]
+        scale = scale[0]
+
+    return center, scale
+
+
+def _fix_aspect_ratio(bbox_scale: np.ndarray,
+                      aspect_ratio: float) -> np.ndarray:
+    """Extend the scale to match the given aspect ratio.
+
+    Args:
+        scale (np.ndarray): The image scale (w, h) in shape (2, )
+        aspect_ratio (float): The ratio of ``w/h``
+
+    Returns:
+        np.ndarray: The reshaped image scale in (2, )
+    """
+    w, h = np.hsplit(bbox_scale, [1])
+    bbox_scale = np.where(w > h * aspect_ratio,
+                          np.hstack([w, w / aspect_ratio]),
+                          np.hstack([h * aspect_ratio, h]))
+    return bbox_scale
+
+
+def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray:
+    """Rotate a point by an angle.
+
+    Args:
+        pt (np.ndarray): 2D point coordinates (x, y) in shape (2, )
+        angle_rad (float): rotation angle in radian
+
+    Returns:
+        np.ndarray: Rotated point in shape (2, )
+    """
+    sn, cs = np.sin(angle_rad), np.cos(angle_rad)
+    rot_mat = np.array([[cs, -sn], [sn, cs]])
+    return rot_mat @ pt
+
+
+def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+    """To calculate the affine matrix, three pairs of points are required. This
+    function is used to get the 3rd point, given 2D points a & b.
+
+    The 3rd point is defined by rotating vector `a - b` by 90 degrees
+    anticlockwise, using b as the rotation center.
+
+    Args:
+        a (np.ndarray): The 1st point (x,y) in shape (2, )
+        b (np.ndarray): The 2nd point (x,y) in shape (2, )
+
+    Returns:
+        np.ndarray: The 3rd point.
+    """
+    direction = a - b
+    c = b + np.r_[-direction[1], direction[0]]
+    return c
+
+
+def get_warp_matrix(center: np.ndarray,
+                    scale: np.ndarray,
+                    rot: float,
+                    output_size: Tuple[int, int],
+                    shift: Tuple[float, float] = (0., 0.),
+                    inv: bool = False) -> np.ndarray:
+    """Calculate the affine transformation matrix that can warp the bbox area
+    in the input image to the output size.
+
+    Args:
+        center (np.ndarray[2, ]): Center of the bounding box (x, y).
+        scale (np.ndarray[2, ]): Scale of the bounding box
+            wrt [width, height].
+        rot (float): Rotation angle (degree).
+        output_size (np.ndarray[2, ] | list(2,)): Size of the
+            destination heatmaps.
+        shift (0-100%): Shift translation ratio wrt the width/height.
+            Default (0., 0.).
+        inv (bool): Option to inverse the affine transform direction.
+            (inv=False: src->dst or inv=True: dst->src)
+
+    Returns:
+        np.ndarray: A 2x3 transformation matrix
+    """
+    shift = np.array(shift)
+    src_w = scale[0]
+    dst_w = output_size[0]
+    dst_h = output_size[1]
+
+    # compute transformation matrix
+    rot_rad = np.deg2rad(rot)
+    src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad)
+    dst_dir = np.array([0., dst_w * -0.5])
+
+    # get four corners of the src rectangle in the original image
+    src = np.zeros((3, 2), dtype=np.float32)
+    src[0, :] = center + scale * shift
+    src[1, :] = center + src_dir + scale * shift
+    src[2, :] = _get_3rd_point(src[0, :], src[1, :])
+
+    # get four corners of the dst rectangle in the input image
+    dst = np.zeros((3, 2), dtype=np.float32)
+    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
+    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
+    dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :])
+
+    if inv:
+        warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src))
+    else:
+        warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst))
+
+    return warp_mat
+
+
+def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict,
+                    img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Get the bbox image as the model input by affine transform.
+
+    Args:
+        input_size (dict): The input size of the model.
+        bbox_scale (dict): The bbox scale of the img.
+        bbox_center (dict): The bbox center of the img.
+        img (np.ndarray): The original image.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: img after affine transform.
+        - np.ndarray[float32]: bbox scale after affine transform.
+    """
+    w, h = input_size
+    warp_size = (int(w), int(h))
+
+    # reshape bbox to fixed aspect ratio
+    bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h)
+
+    # get the affine matrix
+    center = bbox_center
+    scale = bbox_scale
+    rot = 0
+    warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h))
+
+    # do affine transform
+    img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR)
+
+    return img, bbox_scale
+
+
+def get_simcc_maximum(simcc_x: np.ndarray,
+                      simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+    """Get maximum response location and value from simcc representations.
+
+    Note:
+        instance number: N
+        num_keypoints: K
+        heatmap height: H
+        heatmap width: W
+
+    Args:
+        simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx)
+        simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy)
+
+    Returns:
+        tuple:
+        - locs (np.ndarray): locations of maximum heatmap responses in shape
+            (K, 2) or (N, K, 2)
+        - vals (np.ndarray): values of maximum heatmap responses in shape
+            (K,) or (N, K)
+    """
+    N, K, Wx = simcc_x.shape
+    simcc_x = simcc_x.reshape(N * K, -1)
+    simcc_y = simcc_y.reshape(N * K, -1)
+
+    # get maximum value locations
+    x_locs = np.argmax(simcc_x, axis=1)
+    y_locs = np.argmax(simcc_y, axis=1)
+    locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32)
+    max_val_x = np.amax(simcc_x, axis=1)
+    max_val_y = np.amax(simcc_y, axis=1)
+
+    # get maximum value across x and y axis
+    mask = max_val_x > max_val_y
+    max_val_x[mask] = max_val_y[mask]
+    vals = max_val_x
+    locs[vals <= 0.] = -1
+
+    # reshape
+    locs = locs.reshape(N, K, 2)
+    vals = vals.reshape(N, K)
+
+    return locs, vals
+
+
+def decode(simcc_x: np.ndarray, simcc_y: np.ndarray,
+           simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]:
+    """Modulate simcc distribution with Gaussian.
+
+    Args:
+        simcc_x (np.ndarray[K, Wx]): model predicted simcc in x.
+        simcc_y (np.ndarray[K, Wy]): model predicted simcc in y.
+        simcc_split_ratio (int): The split ratio of simcc.
+
+    Returns:
+        tuple: A tuple containing center and scale.
+        - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2)
+        - np.ndarray[float32]: scores in shape (K,) or (n, K)
+    """
+    keypoints, scores = get_simcc_maximum(simcc_x, simcc_y)
+    keypoints /= simcc_split_ratio
+
+    return keypoints, scores
+
+def inference_pose(model, out_bbox, oriImg, model_input_size=(288, 384)):
+    resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size)
+    #outputs = inference(session, resized_img, dtype)
+    outputs = inference(model, resized_img)
+
+    keypoints, scores = postprocess(outputs, model_input_size, center, scale)
+
+    return keypoints, scores

custom_nodes/ComfyUI-tbox/src/dwpose/face.py

@@ -0,0 +1,362 @@
+import logging
+import numpy as np
+from torchvision.transforms import ToTensor, ToPILImage
+import torch
+import torch.nn.functional as F
+import cv2
+
+from . import util
+from torch.nn import Conv2d, Module, ReLU, MaxPool2d, init
+
+
+class FaceNet(Module):
+    """Model the cascading heatmaps. """
+    def __init__(self):
+        super(FaceNet, self).__init__()
+        # cnn to make feature map
+        self.relu = ReLU()
+        self.max_pooling_2d = MaxPool2d(kernel_size=2, stride=2)
+        self.conv1_1 = Conv2d(in_channels=3, out_channels=64,
+                              kernel_size=3, stride=1, padding=1)
+        self.conv1_2 = Conv2d(
+            in_channels=64, out_channels=64, kernel_size=3, stride=1,
+            padding=1)
+        self.conv2_1 = Conv2d(
+            in_channels=64, out_channels=128, kernel_size=3, stride=1,
+            padding=1)
+        self.conv2_2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_1 = Conv2d(
+            in_channels=128, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_2 = Conv2d(
+            in_channels=256, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_3 = Conv2d(
+            in_channels=256, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv3_4 = Conv2d(
+            in_channels=256, out_channels=256, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_1 = Conv2d(
+            in_channels=256, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_2 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_3 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv4_4 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv5_1 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv5_2 = Conv2d(
+            in_channels=512, out_channels=512, kernel_size=3, stride=1,
+            padding=1)
+        self.conv5_3_CPM = Conv2d(
+            in_channels=512, out_channels=128, kernel_size=3, stride=1,
+            padding=1)
+
+        # stage1
+        self.conv6_1_CPM = Conv2d(
+            in_channels=128, out_channels=512, kernel_size=1, stride=1,
+            padding=0)
+        self.conv6_2_CPM = Conv2d(
+            in_channels=512, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage2
+        self.Mconv1_stage2 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage2 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage2 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage3
+        self.Mconv1_stage3 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage3 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage3 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage4
+        self.Mconv1_stage4 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage4 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage4 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage5
+        self.Mconv1_stage5 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage5 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage5 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        # stage6
+        self.Mconv1_stage6 = Conv2d(
+            in_channels=199, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv2_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv3_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv4_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv5_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=7, stride=1,
+            padding=3)
+        self.Mconv6_stage6 = Conv2d(
+            in_channels=128, out_channels=128, kernel_size=1, stride=1,
+            padding=0)
+        self.Mconv7_stage6 = Conv2d(
+            in_channels=128, out_channels=71, kernel_size=1, stride=1,
+            padding=0)
+
+        for m in self.modules():
+            if isinstance(m, Conv2d):
+                init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        """Return a list of heatmaps."""
+        heatmaps = []
+
+        h = self.relu(self.conv1_1(x))
+        h = self.relu(self.conv1_2(h))
+        h = self.max_pooling_2d(h)
+        h = self.relu(self.conv2_1(h))
+        h = self.relu(self.conv2_2(h))
+        h = self.max_pooling_2d(h)
+        h = self.relu(self.conv3_1(h))
+        h = self.relu(self.conv3_2(h))
+        h = self.relu(self.conv3_3(h))
+        h = self.relu(self.conv3_4(h))
+        h = self.max_pooling_2d(h)
+        h = self.relu(self.conv4_1(h))
+        h = self.relu(self.conv4_2(h))
+        h = self.relu(self.conv4_3(h))
+        h = self.relu(self.conv4_4(h))
+        h = self.relu(self.conv5_1(h))
+        h = self.relu(self.conv5_2(h))
+        h = self.relu(self.conv5_3_CPM(h))
+        feature_map = h
+
+        # stage1
+        h = self.relu(self.conv6_1_CPM(h))
+        h = self.conv6_2_CPM(h)
+        heatmaps.append(h)
+
+        # stage2
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage2(h))
+        h = self.relu(self.Mconv2_stage2(h))
+        h = self.relu(self.Mconv3_stage2(h))
+        h = self.relu(self.Mconv4_stage2(h))
+        h = self.relu(self.Mconv5_stage2(h))
+        h = self.relu(self.Mconv6_stage2(h))
+        h = self.Mconv7_stage2(h)
+        heatmaps.append(h)
+
+        # stage3
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage3(h))
+        h = self.relu(self.Mconv2_stage3(h))
+        h = self.relu(self.Mconv3_stage3(h))
+        h = self.relu(self.Mconv4_stage3(h))
+        h = self.relu(self.Mconv5_stage3(h))
+        h = self.relu(self.Mconv6_stage3(h))
+        h = self.Mconv7_stage3(h)
+        heatmaps.append(h)
+
+        # stage4
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage4(h))
+        h = self.relu(self.Mconv2_stage4(h))
+        h = self.relu(self.Mconv3_stage4(h))
+        h = self.relu(self.Mconv4_stage4(h))
+        h = self.relu(self.Mconv5_stage4(h))
+        h = self.relu(self.Mconv6_stage4(h))
+        h = self.Mconv7_stage4(h)
+        heatmaps.append(h)
+
+        # stage5
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage5(h))
+        h = self.relu(self.Mconv2_stage5(h))
+        h = self.relu(self.Mconv3_stage5(h))
+        h = self.relu(self.Mconv4_stage5(h))
+        h = self.relu(self.Mconv5_stage5(h))
+        h = self.relu(self.Mconv6_stage5(h))
+        h = self.Mconv7_stage5(h)
+        heatmaps.append(h)
+
+        # stage6
+        h = torch.cat([h, feature_map], dim=1)  # channel concat
+        h = self.relu(self.Mconv1_stage6(h))
+        h = self.relu(self.Mconv2_stage6(h))
+        h = self.relu(self.Mconv3_stage6(h))
+        h = self.relu(self.Mconv4_stage6(h))
+        h = self.relu(self.Mconv5_stage6(h))
+        h = self.relu(self.Mconv6_stage6(h))
+        h = self.Mconv7_stage6(h)
+        heatmaps.append(h)
+
+        return heatmaps
+
+
+LOG = logging.getLogger(__name__)
+TOTEN = ToTensor()
+TOPIL = ToPILImage()
+
+
+params = {
+    'gaussian_sigma': 2.5,
+    'inference_img_size': 736,  # 368, 736, 1312
+    'heatmap_peak_thresh': 0.1,
+    'crop_scale': 1.5,
+    'line_indices': [
+        [0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 6],
+        [6, 7], [7, 8], [8, 9], [9, 10], [10, 11], [11, 12], [12, 13],
+        [13, 14], [14, 15], [15, 16],
+        [17, 18], [18, 19], [19, 20], [20, 21],
+        [22, 23], [23, 24], [24, 25], [25, 26],
+        [27, 28], [28, 29], [29, 30],
+        [31, 32], [32, 33], [33, 34], [34, 35],
+        [36, 37], [37, 38], [38, 39], [39, 40], [40, 41], [41, 36],
+        [42, 43], [43, 44], [44, 45], [45, 46], [46, 47], [47, 42],
+        [48, 49], [49, 50], [50, 51], [51, 52], [52, 53], [53, 54],
+        [54, 55], [55, 56], [56, 57], [57, 58], [58, 59], [59, 48],
+        [60, 61], [61, 62], [62, 63], [63, 64], [64, 65], [65, 66],
+        [66, 67], [67, 60]
+    ],
+}
+
+
+class Face(object):
+    """
+    The OpenPose face landmark detector model.
+
+    Args:
+        inference_size: set the size of the inference image size, suggested:
+            368, 736, 1312, default 736
+        gaussian_sigma: blur the heatmaps, default 2.5
+        heatmap_peak_thresh: return landmark if over threshold, default 0.1
+
+    """
+    def __init__(self, face_model_path,
+                 inference_size=None,
+                 gaussian_sigma=None,
+                 heatmap_peak_thresh=None):
+        self.inference_size = inference_size or params["inference_img_size"]
+        self.sigma = gaussian_sigma or params['gaussian_sigma']
+        self.threshold = heatmap_peak_thresh or params["heatmap_peak_thresh"]
+        self.model = FaceNet()
+        self.model.load_state_dict(torch.load(face_model_path))
+        # if torch.cuda.is_available():
+        #     self.model = self.model.cuda()
+            # print('cuda')
+        self.model.eval()
+
+    def __call__(self, face_img):
+        H, W, C = face_img.shape
+
+        w_size = 384
+        x_data = torch.from_numpy(util.smart_resize(face_img, (w_size, w_size))).permute([2, 0, 1]) / 256.0 - 0.5
+
+        x_data = x_data.to(self.cn_device)
+
+        with torch.no_grad():
+            hs = self.model(x_data[None, ...])
+            heatmaps = F.interpolate(
+                hs[-1],
+                (H, W),
+                mode='bilinear', align_corners=True).cpu().numpy()[0]
+        return heatmaps
+
+    def compute_peaks_from_heatmaps(self, heatmaps):
+        all_peaks = []
+        for part in range(heatmaps.shape[0]):
+            map_ori = heatmaps[part].copy()
+            binary = np.ascontiguousarray(map_ori > 0.05, dtype=np.uint8)
+
+            if np.sum(binary) == 0:
+                continue
+
+            positions = np.where(binary > 0.5)
+            intensities = map_ori[positions]
+            mi = np.argmax(intensities)
+            y, x = positions[0][mi], positions[1][mi]
+            all_peaks.append([x, y])
+
+        return np.array(all_peaks)

custom_nodes/ComfyUI-tbox/src/dwpose/hand.py

@@ -0,0 +1,94 @@
+import cv2
+import json
+import numpy as np
+import math
+import time
+from scipy.ndimage.filters import gaussian_filter
+import matplotlib.pyplot as plt
+import matplotlib
+import torch
+from skimage.measure import label
+
+from .model import handpose_model
+from . import util
+
+class Hand(object):
+    def __init__(self, model_path):
+        self.model = handpose_model()
+        # if torch.cuda.is_available():
+        #     self.model = self.model.cuda()
+            # print('cuda')
+        model_dict = util.transfer(self.model, torch.load(model_path))
+        self.model.load_state_dict(model_dict)
+        self.model.eval()
+
+    def __call__(self, oriImgRaw):
+        scale_search = [0.5, 1.0, 1.5, 2.0]
+        # scale_search = [0.5]
+        boxsize = 368
+        stride = 8
+        padValue = 128
+        thre = 0.05
+        multiplier = [x * boxsize for x in scale_search]
+
+        wsize = 128
+        heatmap_avg = np.zeros((wsize, wsize, 22))
+
+        Hr, Wr, Cr = oriImgRaw.shape
+
+        oriImg = cv2.GaussianBlur(oriImgRaw, (0, 0), 0.8)
+
+        for m in range(len(multiplier)):
+            scale = multiplier[m]
+            imageToTest = util.smart_resize(oriImg, (scale, scale))
+
+            imageToTest_padded, pad = util.padRightDownCorner(imageToTest, stride, padValue)
+            im = np.transpose(np.float32(imageToTest_padded[:, :, :, np.newaxis]), (3, 2, 0, 1)) / 256 - 0.5
+            im = np.ascontiguousarray(im)
+
+            data = torch.from_numpy(im).float()
+            if torch.cuda.is_available():
+                data = data.cuda()
+
+            with torch.no_grad():
+                data = data.to(self.cn_device)
+                output = self.model(data).cpu().numpy()
+
+            # extract outputs, resize, and remove padding
+            heatmap = np.transpose(np.squeeze(output), (1, 2, 0))  # output 1 is heatmaps
+            heatmap = util.smart_resize_k(heatmap, fx=stride, fy=stride)
+            heatmap = heatmap[:imageToTest_padded.shape[0] - pad[2], :imageToTest_padded.shape[1] - pad[3], :]
+            heatmap = util.smart_resize(heatmap, (wsize, wsize))
+
+            heatmap_avg += heatmap / len(multiplier)
+
+        all_peaks = []
+        for part in range(21):
+            map_ori = heatmap_avg[:, :, part]
+            one_heatmap = gaussian_filter(map_ori, sigma=3)
+            binary = np.ascontiguousarray(one_heatmap > thre, dtype=np.uint8)
+
+            if np.sum(binary) == 0:
+                all_peaks.append([0, 0])
+                continue
+            label_img, label_numbers = label(binary, return_num=True, connectivity=binary.ndim)
+            max_index = np.argmax([np.sum(map_ori[label_img == i]) for i in range(1, label_numbers + 1)]) + 1
+            label_img[label_img != max_index] = 0
+            map_ori[label_img == 0] = 0
+
+            y, x = util.npmax(map_ori)
+            y = int(float(y) * float(Hr) / float(wsize))
+            x = int(float(x) * float(Wr) / float(wsize))
+            all_peaks.append([x, y])
+        return np.array(all_peaks)
+
+if __name__ == "__main__":
+    hand_estimation = Hand('../model/hand_pose_model.pth')
+
+    # test_image = '../images/hand.jpg'
+    test_image = '../images/hand.jpg'
+    oriImg = cv2.imread(test_image)  # B,G,R order
+    peaks = hand_estimation(oriImg)
+    canvas = util.draw_handpose(oriImg, peaks, True)
+    cv2.imshow('', canvas)
+    cv2.waitKey(0)

custom_nodes/ComfyUI-tbox/src/dwpose/model.py

@@ -0,0 +1,218 @@
+import torch
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+
+def make_layers(block, no_relu_layers):
+    layers = []
+    for layer_name, v in block.items():
+        if 'pool' in layer_name:
+            layer = nn.MaxPool2d(kernel_size=v[0], stride=v[1],
+                                    padding=v[2])
+            layers.append((layer_name, layer))
+        else:
+            conv2d = nn.Conv2d(in_channels=v[0], out_channels=v[1],
+                               kernel_size=v[2], stride=v[3],
+                               padding=v[4])
+            layers.append((layer_name, conv2d))
+            if layer_name not in no_relu_layers:
+                layers.append(('relu_'+layer_name, nn.ReLU(inplace=True)))
+
+    return nn.Sequential(OrderedDict(layers))
+
+class bodypose_model(nn.Module):
+    def __init__(self):
+        super(bodypose_model, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = ['conv5_5_CPM_L1', 'conv5_5_CPM_L2', 'Mconv7_stage2_L1',\
+                          'Mconv7_stage2_L2', 'Mconv7_stage3_L1', 'Mconv7_stage3_L2',\
+                          'Mconv7_stage4_L1', 'Mconv7_stage4_L2', 'Mconv7_stage5_L1',\
+                          'Mconv7_stage5_L2', 'Mconv7_stage6_L1', 'Mconv7_stage6_L1']
+        blocks = {}
+        block0 = OrderedDict([
+                      ('conv1_1', [3, 64, 3, 1, 1]),
+                      ('conv1_2', [64, 64, 3, 1, 1]),
+                      ('pool1_stage1', [2, 2, 0]),
+                      ('conv2_1', [64, 128, 3, 1, 1]),
+                      ('conv2_2', [128, 128, 3, 1, 1]),
+                      ('pool2_stage1', [2, 2, 0]),
+                      ('conv3_1', [128, 256, 3, 1, 1]),
+                      ('conv3_2', [256, 256, 3, 1, 1]),
+                      ('conv3_3', [256, 256, 3, 1, 1]),
+                      ('conv3_4', [256, 256, 3, 1, 1]),
+                      ('pool3_stage1', [2, 2, 0]),
+                      ('conv4_1', [256, 512, 3, 1, 1]),
+                      ('conv4_2', [512, 512, 3, 1, 1]),
+                      ('conv4_3_CPM', [512, 256, 3, 1, 1]),
+                      ('conv4_4_CPM', [256, 128, 3, 1, 1])
+                  ])
+
+
+        # Stage 1
+        block1_1 = OrderedDict([
+                        ('conv5_1_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_2_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_3_CPM_L1', [128, 128, 3, 1, 1]),
+                        ('conv5_4_CPM_L1', [128, 512, 1, 1, 0]),
+                        ('conv5_5_CPM_L1', [512, 38, 1, 1, 0])
+                    ])
+
+        block1_2 = OrderedDict([
+                        ('conv5_1_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_2_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_3_CPM_L2', [128, 128, 3, 1, 1]),
+                        ('conv5_4_CPM_L2', [128, 512, 1, 1, 0]),
+                        ('conv5_5_CPM_L2', [512, 19, 1, 1, 0])
+                    ])
+        blocks['block1_1'] = block1_1
+        blocks['block1_2'] = block1_2
+
+        self.model0 = make_layers(block0, no_relu_layers)
+
+        # Stages 2 - 6
+        for i in range(2, 7):
+            blocks['block%d_1' % i] = OrderedDict([
+                    ('Mconv1_stage%d_L1' % i, [185, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d_L1' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d_L1' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d_L1' % i, [128, 38, 1, 1, 0])
+                ])
+
+            blocks['block%d_2' % i] = OrderedDict([
+                    ('Mconv1_stage%d_L2' % i, [185, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d_L2' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d_L2' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d_L2' % i, [128, 19, 1, 1, 0])
+                ])
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_1 = blocks['block1_1']
+        self.model2_1 = blocks['block2_1']
+        self.model3_1 = blocks['block3_1']
+        self.model4_1 = blocks['block4_1']
+        self.model5_1 = blocks['block5_1']
+        self.model6_1 = blocks['block6_1']
+
+        self.model1_2 = blocks['block1_2']
+        self.model2_2 = blocks['block2_2']
+        self.model3_2 = blocks['block3_2']
+        self.model4_2 = blocks['block4_2']
+        self.model5_2 = blocks['block5_2']
+        self.model6_2 = blocks['block6_2']
+
+
+    def forward(self, x):
+
+        out1 = self.model0(x)
+
+        out1_1 = self.model1_1(out1)
+        out1_2 = self.model1_2(out1)
+        out2 = torch.cat([out1_1, out1_2, out1], 1)
+
+        out2_1 = self.model2_1(out2)
+        out2_2 = self.model2_2(out2)
+        out3 = torch.cat([out2_1, out2_2, out1], 1)
+
+        out3_1 = self.model3_1(out3)
+        out3_2 = self.model3_2(out3)
+        out4 = torch.cat([out3_1, out3_2, out1], 1)
+
+        out4_1 = self.model4_1(out4)
+        out4_2 = self.model4_2(out4)
+        out5 = torch.cat([out4_1, out4_2, out1], 1)
+
+        out5_1 = self.model5_1(out5)
+        out5_2 = self.model5_2(out5)
+        out6 = torch.cat([out5_1, out5_2, out1], 1)
+
+        out6_1 = self.model6_1(out6)
+        out6_2 = self.model6_2(out6)
+
+        return out6_1, out6_2
+
+class handpose_model(nn.Module):
+    def __init__(self):
+        super(handpose_model, self).__init__()
+
+        # these layers have no relu layer
+        no_relu_layers = ['conv6_2_CPM', 'Mconv7_stage2', 'Mconv7_stage3',\
+                          'Mconv7_stage4', 'Mconv7_stage5', 'Mconv7_stage6']
+        # stage 1
+        block1_0 = OrderedDict([
+                ('conv1_1', [3, 64, 3, 1, 1]),
+                ('conv1_2', [64, 64, 3, 1, 1]),
+                ('pool1_stage1', [2, 2, 0]),
+                ('conv2_1', [64, 128, 3, 1, 1]),
+                ('conv2_2', [128, 128, 3, 1, 1]),
+                ('pool2_stage1', [2, 2, 0]),
+                ('conv3_1', [128, 256, 3, 1, 1]),
+                ('conv3_2', [256, 256, 3, 1, 1]),
+                ('conv3_3', [256, 256, 3, 1, 1]),
+                ('conv3_4', [256, 256, 3, 1, 1]),
+                ('pool3_stage1', [2, 2, 0]),
+                ('conv4_1', [256, 512, 3, 1, 1]),
+                ('conv4_2', [512, 512, 3, 1, 1]),
+                ('conv4_3', [512, 512, 3, 1, 1]),
+                ('conv4_4', [512, 512, 3, 1, 1]),
+                ('conv5_1', [512, 512, 3, 1, 1]),
+                ('conv5_2', [512, 512, 3, 1, 1]),
+                ('conv5_3_CPM', [512, 128, 3, 1, 1])
+            ])
+
+        block1_1 = OrderedDict([
+            ('conv6_1_CPM', [128, 512, 1, 1, 0]),
+            ('conv6_2_CPM', [512, 22, 1, 1, 0])
+        ])
+
+        blocks = {}
+        blocks['block1_0'] = block1_0
+        blocks['block1_1'] = block1_1
+
+        # stage 2-6
+        for i in range(2, 7):
+            blocks['block%d' % i] = OrderedDict([
+                    ('Mconv1_stage%d' % i, [150, 128, 7, 1, 3]),
+                    ('Mconv2_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv3_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv4_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv5_stage%d' % i, [128, 128, 7, 1, 3]),
+                    ('Mconv6_stage%d' % i, [128, 128, 1, 1, 0]),
+                    ('Mconv7_stage%d' % i, [128, 22, 1, 1, 0])
+                ])
+
+        for k in blocks.keys():
+            blocks[k] = make_layers(blocks[k], no_relu_layers)
+
+        self.model1_0 = blocks['block1_0']
+        self.model1_1 = blocks['block1_1']
+        self.model2 = blocks['block2']
+        self.model3 = blocks['block3']
+        self.model4 = blocks['block4']
+        self.model5 = blocks['block5']
+        self.model6 = blocks['block6']
+
+    def forward(self, x):
+        out1_0 = self.model1_0(x)
+        out1_1 = self.model1_1(out1_0)
+        concat_stage2 = torch.cat([out1_1, out1_0], 1)
+        out_stage2 = self.model2(concat_stage2)
+        concat_stage3 = torch.cat([out_stage2, out1_0], 1)
+        out_stage3 = self.model3(concat_stage3)
+        concat_stage4 = torch.cat([out_stage3, out1_0], 1)
+        out_stage4 = self.model4(concat_stage4)
+        concat_stage5 = torch.cat([out_stage4, out1_0], 1)
+        out_stage5 = self.model5(concat_stage5)
+        concat_stage6 = torch.cat([out_stage5, out1_0], 1)
+        out_stage6 = self.model6(concat_stage6)
+        return out_stage6
+

custom_nodes/ComfyUI-tbox/src/dwpose/types.py

@@ -0,0 +1,30 @@
+from typing import NamedTuple, List, Optional
+
+class Keypoint(NamedTuple):
+    x: float
+    y: float
+    score: float = 1.0
+    id: int = -1
+
+
+class BodyResult(NamedTuple):
+    # Note: Using `Optional` instead of `|` operator as the ladder is a Python
+    # 3.10 feature.
+    # Annotator code should be Python 3.8 Compatible, as controlnet repo uses
+    # Python 3.8 environment.
+    # https://github.com/lllyasviel/ControlNet/blob/d3284fcd0972c510635a4f5abe2eeb71dc0de524/environment.yaml#L6
+    keypoints: List[Optional[Keypoint]]
+    total_score: float = 0.0
+    total_parts: int = 0
+
+
+HandResult = List[Keypoint]
+FaceResult = List[Keypoint]
+AnimalPoseResult = List[Keypoint]
+
+
+class PoseResult(NamedTuple):
+    body: BodyResult
+    left_hand: Optional[HandResult]
+    right_hand: Optional[HandResult]
+    face: Optional[FaceResult]