113 lines
4.1 KiB
Python
113 lines
4.1 KiB
Python
import comfy
|
|
|
|
class SetMetadataString:
|
|
"""
|
|
Set a single custom metadata field and optionally update 'parameters' field in png metadata
|
|
"""
|
|
CATEGORY = "utils"
|
|
RETURN_TYPES = ()
|
|
OUTPUT_NODE = True
|
|
FUNCTION = "process"
|
|
|
|
SET_PARAMETERS = True
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"name": ("STRING", {}),
|
|
"value": ("STRING", {"forceInput": True, "default": ""}),
|
|
},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO"
|
|
},
|
|
}
|
|
|
|
def process(self, name: str, value: str = "", update_parameters: bool = True, extra_pnginfo=None):
|
|
if extra_pnginfo is not None and name and value:
|
|
extra_pnginfo[name] = value
|
|
if self.SET_PARAMETERS and name != "parameters":
|
|
extra_pnginfo["parameters"] = self.update_parameters(extra_pnginfo.get("parameters", ""), name, value)
|
|
return (None,)
|
|
|
|
def update_parameters(self, params_string: str, name: str, value: str) -> str:
|
|
"""Modify the a1111 compatible (mostly) parameters string to include the new parameter
|
|
|
|
Parameters format:
|
|
<Prompt>
|
|
[Negative prompt: <Negative prompt>]
|
|
[Steps: <steps>, Seed: <seed>, ...]
|
|
"""
|
|
prompt, *other = params_string.split("\n") + [""] * 3
|
|
negative = None
|
|
params = None
|
|
for item in other:
|
|
if item:
|
|
if item.startswith("Negative prompt: "):
|
|
negative = item
|
|
else:
|
|
params = item
|
|
|
|
if name.lower() == "prompt":
|
|
prompt = value
|
|
elif name.lower() in ("negative prompt", "negative"):
|
|
negative = f"Negative prompt: {value}"
|
|
else:
|
|
if isinstance(value, str) and any(c in value for c in (',', '\n', '"')):
|
|
value = '"' + value.replace('"', '\\"') + '"'
|
|
if params:
|
|
params += f", {name}: {value}"
|
|
else:
|
|
params = f"{name}: {value}"
|
|
|
|
return "\n".join(filter(None, [prompt, negative, params]))
|
|
|
|
class SetMetadataAll(SetMetadataString):
|
|
"""
|
|
Set multiple metadata at once and write them to 'parameters' field in png metadata
|
|
"""
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {},
|
|
"optional": {
|
|
"prompt": ("STRING", {"forceInput": True, "default": None}),
|
|
"negative": ("STRING", {"forceInput": True, "default": None}),
|
|
"model_name": ("STRING", {"forceInput": True, "default": None}),
|
|
"steps": ("INT", {"forceInput": True, "default": None}),
|
|
"cfg": ("FLOAT", {"forceInput": True, "default": None}),
|
|
"sampler_name": (comfy.samplers.KSampler.SAMPLERS, {"forceInput": True, "default": None}),
|
|
"scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"forceInput": True, "default": None}),
|
|
"seed": ("INT", {"forceInput": True, "default": None}),
|
|
},
|
|
"hidden": {
|
|
"extra_pnginfo": "EXTRA_PNGINFO"
|
|
},
|
|
}
|
|
|
|
def process(self, **kwargs):
|
|
extra_pnginfo = kwargs.pop("extra_pnginfo", None)
|
|
if extra_pnginfo is not None:
|
|
params = extra_pnginfo.get("parameters", "")
|
|
for k, v in kwargs.items():
|
|
if v is not None:
|
|
params = self.update_parameters(params, self.convert_name(k), self.convert_value(v))
|
|
if params:
|
|
extra_pnginfo["parameters"] = params
|
|
return (None,)
|
|
|
|
def convert_name(self, name: str) -> str:
|
|
names_dict = {
|
|
"prompt": "Prompt",
|
|
"negative": "Negative prompt",
|
|
"model_name": "Model",
|
|
"steps": "Steps",
|
|
"seed": "Seed",
|
|
"cfg": "CFG scale",
|
|
"sampler_name": "Sampler",
|
|
"scheduler": "Scheduler",
|
|
}
|
|
return names_dict.get(name, name)
|
|
|
|
def convert_value(self, value):
|
|
return value |