CavinHuang-comfyui-nodes-docs/docs/DeepDanbooruCaption.md

Documentation

• Class name: DeepDanbooruCaption
• Category: Art Venture/Utils
• Output node: False

DeepDanbooruCaption节点旨在使用DeepDanbooru模型为图像生成说明文字。它处理图像以识别和描述其内容，使用标签，可选择性地为每个说明文字添加前缀和后缀，并可配置为过滤标签、调整阈值和按字母顺序排序标签。

Input types

Required

• image
  • 需要生成说明文字的图像。这是生成说明文字的主要输入。
  • Comfy dtype: IMAGE
  • Python dtype: torch.Tensor
• threshold
  • 在说明文字中包含标签的置信度阈值。置信度低于此阈值的标签将被排除。
  • Comfy dtype: FLOAT
  • Python dtype: float
• sort_alpha
  • 决定说明文字中的标签是否应按字母顺序排序。
  • Comfy dtype: BOOLEAN
  • Python dtype: bool
• use_spaces
  • 控制说明文字中标签之间是否应使用空格。
  • Comfy dtype: BOOLEAN
  • Python dtype: bool
• escape
  • 指示是否应对标签中的特殊字符进行转义。
  • Comfy dtype: BOOLEAN
  • Python dtype: bool
• filter_tags
  • 要从说明文字中排除的标签列表。
  • Comfy dtype: STRING
  • Python dtype: List[str]

Optional

• device_mode
  • 指定运行DeepDanbooru模型的设备模式（如CPU、GPU、AUTO）。
  • Comfy dtype: COMBO[STRING]
  • Python dtype: str
• prefix
  • 添加到每个说明文字之前的前缀。
  • Comfy dtype: STRING
  • Python dtype: str
• suffix
  • 添加到每个说明文字之后的后缀。
  • Comfy dtype: STRING
  • Python dtype: str
• enabled
  • 启用或禁用说明文字生成功能。
  • Comfy dtype: BOOLEAN
  • Python dtype: bool

Output types

• caption
  • 为输入图像生成的说明文字，可能会被前缀和后缀修改。
  • Comfy dtype: STRING
  • Python dtype: List[str]

Usage tips

• Infra type: GPU
• Common nodes: unknown

Source code

class DeepDanbooruCaption:
    def __init__(self):
        pass

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "threshold": ("FLOAT", {"default": 0.5, "min": 0, "max": 1, "step": 0.01}),
                "sort_alpha": ("BOOLEAN", {"default": True}),
                "use_spaces": ("BOOLEAN", {"default": True}),
                "escape": ("BOOLEAN", {"default": True}),
                "filter_tags": ("STRING", {"default": "blacklist", "multiline": True}),
            },
            "optional": {
                "device_mode": (["AUTO", "Prefer GPU", "CPU"],),
                "prefix": ("STRING", {"default": ""}),
                "suffix": ("STRING", {"default": ""}),
                "enabled": ("BOOLEAN", {"default": True}),
            },
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("caption",)
    OUTPUT_IS_LIST = (True,)
    FUNCTION = "caption"
    CATEGORY = "Art Venture/Utils"

    def caption(
        self,
        image,
        threshold,
        sort_alpha,
        use_spaces,
        escape,
        filter_tags,
        device_mode="AUTO",
        prefix="",
        suffix="",
        enabled=True,
    ):
        if not enabled:
            return ([join_caption("", prefix, suffix)],)

        model = load_danbooru(device_mode)

        try:
            captions = []

            for img in image:
                img = tensor2pil(img)
                img = resize_image(img.convert("RGB"), 512, 512, resize_mode=2)
                arr = np.expand_dims(np.array(img, dtype=np.float32), 0) / 255

                with torch.no_grad():
                    x = torch.from_numpy(arr).to(gpu)
                    y = model(x)[0].detach().cpu().numpy()

                probability_dict = {}

                for tag, probability in zip(model.tags, y):
                    if probability < threshold:
                        continue

                    if tag.startswith("rating:"):
                        continue

                    probability_dict[tag] = probability

                if sort_alpha:
                    tags = sorted(probability_dict)
                else:
                    tags = [tag for tag, _ in sorted(probability_dict.items(), key=lambda x: -x[1])]

                res = []
                filtertags = {x.strip().replace(" ", "_") for x in filter_tags.split(",")}

                for tag in [x for x in tags if x not in filtertags]:
                    probability = probability_dict[tag]
                    tag_outformat = tag
                    if use_spaces:
                        tag_outformat = tag_outformat.replace("_", " ")
                    if escape:
                        tag_outformat = re.sub(re_special, r"\\\1", tag_outformat)

                    res.append(tag_outformat)

                caption = ", ".join(res)
                caption = join_caption(caption, prefix, suffix)
                captions.append(caption)

            return (captions,)
        except:
            raise
        finally:
            unload_danbooru()