137 lines
4.1 KiB
Python
137 lines
4.1 KiB
Python
import re
|
|
|
|
class CleanCaptionNode:
|
|
"""
|
|
A ComfyUI node that cleans captions for LoRA training by removing
|
|
common "meta" descriptions (e.g., "The image is a digital illustration...").
|
|
"""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"multiline": True, "default": ""}),
|
|
"trigger_word": ("STRING", {"default": "cayesvgrt"}),
|
|
},
|
|
"optional": {
|
|
"enabled": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("cleaned_text",)
|
|
FUNCTION = "process_caption"
|
|
CATEGORY = "LoRA/Text"
|
|
|
|
def process_caption(self, text, trigger_word, enabled):
|
|
if not enabled:
|
|
return (text,)
|
|
|
|
# --- CONFIGURATION ---
|
|
PREFIXES_TO_REMOVE = [
|
|
"The image is a digital illustration of",
|
|
"The image is a digital painting of",
|
|
"The image is an illustration of",
|
|
"The image is a 3D render of",
|
|
"The image shows",
|
|
"The image is a",
|
|
"The image is",
|
|
]
|
|
|
|
PHRASES_TO_REMOVE = [
|
|
"digital illustration",
|
|
"digital painting",
|
|
"vector art",
|
|
"of the image",
|
|
"in the image",
|
|
"of the illustration",
|
|
"in the illustration",
|
|
"to the image",
|
|
"to the scene",
|
|
|
|
]
|
|
|
|
SENTENCE_BAD_KEYWORDS = [
|
|
"overall style",
|
|
"cartoon-like",
|
|
"digital art",
|
|
"overall mood",
|
|
"overall atmosphere",
|
|
"overall color scheme",
|
|
"overall design",
|
|
"whimsical feel",
|
|
"serene feel",
|
|
"peaceful and serene",
|
|
"vibrant and lively",
|
|
"bright and cheerful",
|
|
"illustration is",
|
|
"image is",
|
|
"scene is",
|
|
"illustration of ",
|
|
"render of ",
|
|
]
|
|
# ---------------------
|
|
|
|
# 1. Remove Prefixes
|
|
sorted_prefixes = sorted(PREFIXES_TO_REMOVE, key=len, reverse=True)
|
|
for prefix in sorted_prefixes:
|
|
if text.lower().startswith(prefix.lower()):
|
|
text = text[len(prefix):].strip()
|
|
break
|
|
|
|
# 2. Sentence filtering
|
|
sentences = re.split(r'(?<=[.!?])\s+', text)
|
|
cleaned_sentences = []
|
|
|
|
for sent in sentences:
|
|
# Check if sentence contains bad keywords
|
|
if any(kw in sent.lower() for kw in SENTENCE_BAD_KEYWORDS):
|
|
continue
|
|
|
|
# Remove specific phrases within sentence
|
|
for phrase in PHRASES_TO_REMOVE:
|
|
pattern = re.compile(re.escape(phrase), re.IGNORECASE)
|
|
sent = pattern.sub("", sent)
|
|
|
|
# Clean up spaces
|
|
sent = re.sub(r'\s+', ' ', sent).strip()
|
|
|
|
# Fix punctuation issues like "center ." -> "center."
|
|
sent = re.sub(r'\s+([.!?])', r'\1', sent)
|
|
|
|
# Remove punctuation-only sentences or empty
|
|
if not any(c.isalnum() for c in sent):
|
|
continue
|
|
|
|
cleaned_sentences.append(sent)
|
|
|
|
cleaned_text = " ".join(cleaned_sentences)
|
|
|
|
# Post-cleaning cleanup
|
|
if cleaned_text.lower().startswith("of "):
|
|
cleaned_text = cleaned_text[3:].strip()
|
|
|
|
# Capitalize
|
|
if cleaned_text:
|
|
cleaned_text = cleaned_text[0].upper() + cleaned_text[1:]
|
|
|
|
# Add Trigger Word if provided and not already present
|
|
if trigger_word and trigger_word.strip():
|
|
tw = trigger_word.strip()
|
|
if not cleaned_text.startswith(tw):
|
|
cleaned_text = f"{tw}, {cleaned_text}"
|
|
|
|
return (cleaned_text,)
|
|
|
|
# Mappings for ComfyUI to recognize the node
|
|
NODE_CLASS_MAPPINGS = {
|
|
"CleanCaptionNode": CleanCaptionNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"CleanCaptionNode": "Clean Caption (LoRA)"
|
|
}
|