99 lines
2.8 KiB
Python
99 lines
2.8 KiB
Python
from pathlib import Path
|
|
|
|
import folder_paths
|
|
import logging
|
|
|
|
from collections.abc import Iterable
|
|
from abc import ABC
|
|
from dynamicprompts.sampling_context import SamplingContext
|
|
from dynamicprompts.enums import SamplingMethod
|
|
from dynamicprompts.wildcards.wildcard_manager import WildcardManager
|
|
import re
|
|
|
|
logger = logging.getLogger(__name__)
|
|
wildcardManager = WildcardManager(
|
|
Path(folder_paths.get_folder_paths("wildcards")[0]))
|
|
|
|
|
|
class RandomPrompt(ABC):
|
|
def __init__(self):
|
|
self._current_prompt = None
|
|
self.lora_spec_re = re.compile("(<(?:lora|lyco):[^>]+>)")
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {
|
|
"multiline": True,
|
|
"default": ""
|
|
}),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, text):
|
|
# Force re-evaluation of the node
|
|
return float("NaN")
|
|
|
|
RETURN_TYPES = ("STRING", "STRING",)
|
|
RETURN_NAMES = ("prompt", "loras",)
|
|
|
|
FUNCTION = "process"
|
|
|
|
# OUTPUT_NODE = False
|
|
|
|
CATEGORY = "Malefish/conditioning"
|
|
|
|
@property
|
|
def context(self) -> SamplingContext:
|
|
return SamplingContext(
|
|
wildcard_manager=wildcardManager,
|
|
default_sampling_method=SamplingMethod.RANDOM,
|
|
)
|
|
|
|
def _get_next_prompt(self, prompts: Iterable[str], current_prompt: str) -> str:
|
|
"""
|
|
Get the next prompt from the prompts generator.
|
|
"""
|
|
try:
|
|
return next(prompts)
|
|
except (StopIteration, RuntimeError):
|
|
self._prompts = self.context.sample_prompts(current_prompt)
|
|
try:
|
|
return next(prompts)
|
|
except StopIteration:
|
|
logger.exception("No more prompts to generate!")
|
|
return ""
|
|
|
|
def has_prompt_changed(self, text: str) -> bool:
|
|
"""
|
|
Check if the prompt has changed.
|
|
"""
|
|
return self._current_prompt != text
|
|
|
|
def process(self, text):
|
|
if text.strip() == "":
|
|
return ("",)
|
|
|
|
if self.has_prompt_changed(text):
|
|
self._current_prompt = text
|
|
self._prompts = self.context.sample_prompts(self._current_prompt)
|
|
|
|
if self._prompts is None:
|
|
logger.exception("Something went wrong. Prompts is None!")
|
|
return ("",)
|
|
|
|
if self._current_prompt is None:
|
|
logger.exception("Something went wrong. Current prompt is None!")
|
|
return ("",)
|
|
|
|
new_prompt = self._get_next_prompt(self._prompts, self._current_prompt)
|
|
|
|
extracted_loras = self.lora_spec_re.findall(new_prompt)
|
|
filtered_text = self.lora_spec_re.sub("", new_prompt)
|
|
|
|
print(f"New prompt: {new_prompt}")
|
|
return (filtered_text, "\n".join(extracted_loras),)
|