52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
import comfy.utils # noqa
|
|
from .categories import CATE_GEN
|
|
from .utils.ask_bc_docker import bc_docker_ask_stream, PromptItem
|
|
from .utils.job import Job
|
|
|
|
|
|
class FEBatchGenStringBCDocker:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"user_prompts": ("BCD_GEN_CTX",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ('prompt',)
|
|
FUNCTION = "query"
|
|
INPUT_IS_LIST = (True,)
|
|
OUTPUT_IS_LIST = (True,)
|
|
CATEGORY = CATE_GEN
|
|
OUTPUT_NODE = True
|
|
|
|
def get_result(self, kwds):
|
|
result = ""
|
|
if kwds.get("prompt_list", None) is not None:
|
|
kwds["messages"] = kwds['prompt_list']['prompts']
|
|
else:
|
|
kwds["messages"] = [PromptItem("user", kwds['user_prompt'])]
|
|
if "prompt_list" in kwds:
|
|
del kwds['prompt_list']
|
|
if "user_prompt" in kwds:
|
|
del kwds['user_prompt']
|
|
for _ in bc_docker_ask_stream(**kwds):
|
|
result += _
|
|
return result
|
|
|
|
def query(self, user_prompts):
|
|
steps = len(user_prompts)
|
|
tasks = []
|
|
for user_prompt in user_prompts:
|
|
tasks.append(Job(target=self.get_result, args=(user_prompt,)))
|
|
for task in tasks:
|
|
task.start()
|
|
result = []
|
|
pbar = comfy.utils.ProgressBar(steps)
|
|
for idx, task in enumerate(tasks):
|
|
result.append(task.get_result(nJoin=True))
|
|
pbar.update(1)
|
|
return (result,)
|