166 lines
5.0 KiB
Markdown
166 lines
5.0 KiB
Markdown
|
|
# 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
|
|
```python
|
|
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()
|
|
|
|
```
|