File size: 2,777 Bytes
70ceb8b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """
Custom node for splitting WAN 2.2 Story prompts
Handles various separators automatically
"""
class StorySplitNode:
"""
Splits WAN 2.2 Story output into 4 separate prompts
Handles different separators: \n, \n\n, \n\n\n
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"text": ("STRING", {"multiline": True, "default": ""}),
},
"optional": {
"auto_split": ("BOOLEAN", {"default": True}),
"custom_delimiter": ("STRING", {"multiline": False, "default": ""}),
}
}
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING")
RETURN_NAMES = ("prompt_1", "prompt_2", "prompt_3", "prompt_4", "debug_info")
FUNCTION = "split_story"
CATEGORY = "🔷 QwenVL-Mod/Utils"
def split_story(self, text, auto_split=True, custom_delimiter=""):
"""
Split story text into 4 prompts
"""
if not text:
return ("", "", "", "", "Empty input")
text = text.strip()
if auto_split and not custom_delimiter:
# Try different separators automatically
if '\n\n\n' in text:
prompts = text.split('\n\n\n')
elif '\n\n' in text:
prompts = text.split('\n\n')
elif '\n' in text:
prompts = text.split('\n')
else:
# No clear separator, split by paragraphs
prompts = [p for p in text.split('\n') if p.strip()]
else:
# Use custom delimiter
delimiter = custom_delimiter if custom_delimiter else '\n\n'
prompts = text.split(delimiter)
# Clean up prompts
clean_prompts = []
for prompt in prompts:
clean_prompt = prompt.strip()
if clean_prompt and not clean_prompt.startswith("Prompt") and "content describing" not in clean_prompt:
clean_prompts.append(clean_prompt)
# Ensure we have exactly 4 prompts
while len(clean_prompts) < 4:
clean_prompts.append("")
result_prompts = clean_prompts[:4]
# Debug info
debug = f"Found {len(prompts)} parts, cleaned to {len(clean_prompts)} prompts"
return (
result_prompts[0],
result_prompts[1],
result_prompts[2],
result_prompts[3],
debug
)
# Node mapping for ComfyUI
NODE_CLASS_MAPPINGS = {
"StorySplitNode": StorySplitNode
}
NODE_DISPLAY_NAME_MAPPINGS = {
"StorySplitNode": "Story Split Node"
}
|