56 lines
1.7 KiB
Python
56 lines
1.7 KiB
Python
import logging
|
|
|
|
from comfy_api.latest import io
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class SplitPrompt(io.ComfyNode):
|
|
"""
|
|
A simple node to split a prompt into a batch of prompt.
|
|
Inspired from [ComfyUI-Inspire-Pack](https://github.com/ltdrdata/ComfyUI-Inspire-Pack)
|
|
"""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="SplitPrompt",
|
|
display_name="Split prompt",
|
|
description="Split a prompt in a list of prompts",
|
|
category="pecha_kucha",
|
|
inputs=[
|
|
io.String.Input(
|
|
"prompt",
|
|
default="First Prompt \nSecond Prompt",
|
|
multiline=True,
|
|
tooltip="A multiline prompt that will be split to batch generate.",
|
|
),
|
|
io.String.Input(
|
|
"split_on",
|
|
default="\n",
|
|
multiline=False,
|
|
tooltip="Will split the input prompt on this character.",
|
|
),
|
|
io.Boolean.Input(
|
|
"trim",
|
|
default=True,
|
|
tooltip="Will trim prompts and remove empty ones.",
|
|
),
|
|
],
|
|
is_output_node=False,
|
|
outputs=[io.String.Output(display_name="split prompt")],
|
|
)
|
|
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
@classmethod
|
|
def fingerprint_inputs(cls, prompt, split_on="\n", trim=True):
|
|
return f"trim:{trim}_split:{split_on}_prompt:{prompt}"
|
|
|
|
@classmethod
|
|
def execute(cls, prompt, split_on="\n", trim=True):
|
|
prompts = prompt.split(split_on)
|
|
if trim:
|
|
prompts = [prompt.strip() for prompt in prompts]
|
|
return (prompts,)
|