55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
from .categories import CATE_GEN
|
|
from .utils.ask_stream import openai_ask_background
|
|
from .config.configs import config
|
|
from .utils.any_hack import any
|
|
from .utils.node_defs import FEAlwaysChangeNode
|
|
from .utils.cached_fn_result import CachedFunctionResult
|
|
from .utils.job import Job
|
|
|
|
|
|
class FEGenStringNBus(FEAlwaysChangeNode):
|
|
global_cached = CachedFunctionResult()
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"pkws": (any,),
|
|
"api": ("STRING", {"default": config["nbus_api"], "multiline": False}),
|
|
"model": ("STRING", {"default": "sd_axl3plus_v2", "multiline": False}),
|
|
"cache_charge": ("BOOLEAN", {"default": True})
|
|
},
|
|
"optional": {
|
|
"background": (any,),
|
|
}
|
|
}
|
|
|
|
INPUT_IS_LIST = True
|
|
RETURN_TYPES = ("STRING", any)
|
|
RETURN_NAMES = ('prompt', 'pkw')
|
|
FUNCTION = "query"
|
|
CATEGORY = CATE_GEN
|
|
OUTPUT_NODE = True
|
|
|
|
def query(self, pkws, api, model, cache_charge: bool, background=None):
|
|
if background is None or not background:
|
|
background = {}
|
|
if isinstance(background, list):
|
|
background = background[0]
|
|
if not api:
|
|
api = config["nbus_api"]
|
|
if isinstance(api, list):
|
|
api = api[0]
|
|
if isinstance(model, list):
|
|
model = model[0]
|
|
if isinstance(cache_charge, list):
|
|
cache_charge = cache_charge[0]
|
|
print("generating prompt...", end="")
|
|
if cache_charge:
|
|
fn = Job(target=openai_ask_background, model_name=model, background=background, pkw_list=pkws, api=api)
|
|
result, pkw, _, _ = FEGenStringNBus.global_cached.queue(model, fn)
|
|
else:
|
|
result, pkw, _, _ = openai_ask_background(model_name=model, background=background, pkw_list=pkws, api=api)
|
|
print("done")
|
|
return (result, pkw,)
|