diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..28bef68 --- /dev/null +++ b/__init__.py @@ -0,0 +1 @@ +from .tag_blocker import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS diff --git a/assets/warning.png b/assets/warning.png new file mode 100644 index 0000000..988530b Binary files /dev/null and b/assets/warning.png differ diff --git a/tag_blocker.py b/tag_blocker.py new file mode 100644 index 0000000..6830c7c --- /dev/null +++ b/tag_blocker.py @@ -0,0 +1,137 @@ +import os +import torch +from PIL import Image +from torchvision import transforms + +class TagKeywordBlocker: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE",), + "tags": ("STRING", {"multiline": True, "default": ""}), + "keywords": ("STRING", {"multiline": True, "default": "penis,futanari,erection"}) + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "filter_image" + CATEGORY = "Keyword Image Blocker" + + def __init__(self): + base = os.path.dirname(__file__) + self.warning_path = os.path.join(base, "assets", "warning.png") + self.to_tensor = transforms.ToTensor() + + def _to_hwc(self, t: torch.Tensor) -> torch.Tensor: + """Convert any [1,C,H,W], [C,H,W] or [1,H,W,C], [H,W,C] → [H,W,3]""" + print(f"[DEBUG _to_hwc] before: dim={t.dim()}, shape={tuple(t.shape)}") + # drop batch dim + if t.dim() == 4 and t.shape[0] == 1: + t = t.squeeze(0) + print(f"[DEBUG _to_hwc] dropped batch → shape={tuple(t.shape)}") + + # now t.dim()==3 + if t.dim() != 3: + raise RuntimeError(f"[DEBUG _to_hwc] Unexpected tensor dim: {t.shape}") + + # channels-first [C,H,W]? + if t.shape[0] == 3: + t = t.permute(1, 2, 0) # → [H,W,C] + print(f"[DEBUG _to_hwc] permuted CHW→HWC → shape={tuple(t.shape)}") + + # or maybe it's already [H,W,C]? + elif t.shape[-1] == 3: + print(f"[DEBUG _to_hwc] already HWC → shape={tuple(t.shape)}") + else: + raise RuntimeError(f"[DEBUG _to_hwc] No channel dim found in: {t.shape}") + + return t + + def filter_image(self, image, tags, keywords): + # break inputs into list + frames = image if isinstance(image, (list, tuple)) else [image] + + # log incoming frames + print(f"[DEBUG filter_image] got {len(frames)} frame(s)") + for i, frm in enumerate(frames): + if isinstance(frm, torch.Tensor): + print(f" Frame {i}: Tensor dim={frm.dim()}, shape={tuple(frm.shape)}, " + f"dtype={frm.dtype}, device={frm.device}") + else: + print(f" Frame {i}: {type(frm)} size={frm.size if hasattr(frm, 'size') else 'N/A'}") + + # keyword check + raw = str(tags).lower() + kws = [k.strip().lower() for k in keywords.split(",") if k.strip()] + hit = any(kw in raw for kw in kws) + print(f"[DEBUG filter_image] tags='{raw}', keywords={kws}, hit={hit}") + + # pick largest tensor frame to size warning + best, area = None, -1 + for frm in frames: + if not isinstance(frm, torch.Tensor): + continue + t = frm + # detect H×W + if t.dim() == 4 and t.shape[-1] == 3: + H, W = t.shape[1], t.shape[2] + elif t.dim() == 4 and t.shape[1] == 3: + H, W = t.shape[2], t.shape[3] + elif t.dim() == 3 and t.shape[0] == 3: + H, W = t.shape[1], t.shape[2] + elif t.dim() == 3 and t.shape[-1] == 3: + H, W = t.shape[0], t.shape[1] + else: + continue + a = H * W + if a > area: + best, area = frm, a + + # if no tensor frames, bail out and just pass originals in HWC + if best is None or not hit: + out = [] + for frm in frames: + if isinstance(frm, torch.Tensor): + out.append(self._to_hwc(frm)) + else: + # fallback: convert PIL→tensor HWC + t = self.to_tensor(frm).permute(1, 2, 0) + out.append(t) + print("[DEBUG filter_image] no hit or no tensors → passing through HWC tensors") + return (out,) + + # on hit: size warning to best frame + # extract H,W from best + t0 = best + if t0.dim() == 4 and t0.shape[-1] == 3: + H, W = t0.shape[1], t0.shape[2] + elif t0.dim() == 4 and t0.shape[1] == 3: + H, W = t0.shape[2], t0.shape[3] + elif t0.dim() == 3 and t0.shape[0] == 3: + H, W = t0.shape[1], t0.shape[2] + else: + H, W = t0.shape[0], t0.shape[1] + + # load & resize warning + pil_warn = Image.open(self.warning_path).convert("RGB") + pil_warn = pil_warn.resize((W, H), resample=Image.LANCZOS) + + # to tensor CHW→HWC + warn_t = self.to_tensor(pil_warn).permute(1, 2, 0) + warn_t = warn_t.to(best.device).to(best.dtype) + print(f"[DEBUG filter_image] hit → warning tensor shape={tuple(warn_t.shape)}") + + # return single-element list + return ([warn_t],) + + + + +NODE_CLASS_MAPPINGS = { + "TagKeywordBlocker": TagKeywordBlocker +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "TagKeywordBlocker": "Keyword Image Blocker" +}