356 lines
11 KiB
Python
356 lines
11 KiB
Python
import os
|
|
import json
|
|
import comfy.samplers # Taps into ComfyUI's native sampler/scheduler listings
|
|
import folder_paths # Taps into ComfyUI's native model folder tracking
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
PROMPT_PRESETS_FILE = os.path.join(BASE_DIR, "prompt_presets.json")
|
|
GEN_PRESETS_FILE = os.path.join(BASE_DIR, "generation_presets.json")
|
|
|
|
|
|
# ==========================================
|
|
# HELPERS
|
|
# ==========================================
|
|
def get_presets():
|
|
if not os.path.exists(PROMPT_PRESETS_FILE):
|
|
default_data = [
|
|
{
|
|
"name": "sample",
|
|
"positive_prompt": "masterpiece, {prompt}",
|
|
"negative_prompt": "bad quality, {prompt}",
|
|
}
|
|
]
|
|
with open(PROMPT_PRESETS_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(default_data, f, indent=2)
|
|
|
|
presets_dict = {}
|
|
try:
|
|
with open(PROMPT_PRESETS_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
if isinstance(data, list):
|
|
for item in data:
|
|
if "name" in item:
|
|
presets_dict[item["name"]] = {
|
|
"positive": item.get("positive_prompt", ""),
|
|
"negative": item.get("negative_prompt", ""),
|
|
}
|
|
except Exception:
|
|
pass
|
|
return presets_dict
|
|
|
|
|
|
def get_generation_presets():
|
|
"""Reads the generation_presets.json file or creates a baseline configuration if missing."""
|
|
if not os.path.exists(GEN_PRESETS_FILE):
|
|
# Fetch available upscale models to set a valid default if possible
|
|
models = folder_paths.get_filename_list("upscale_models")
|
|
default_model = models[0] if models else "None"
|
|
|
|
default_data = [
|
|
{
|
|
"name": "sample",
|
|
"width": 512,
|
|
"height": 768,
|
|
"cfg": 7.0,
|
|
"steps": 25,
|
|
"sampler_name": "euler",
|
|
"scheduler": "normal",
|
|
"hires_steps": 15,
|
|
"hires_cfg": 5.5,
|
|
"hires_denoise": 0.45,
|
|
"hires_upscale_by": 2.0,
|
|
"hires_model": default_model,
|
|
"is_latent": True,
|
|
}
|
|
]
|
|
with open(GEN_PRESETS_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(default_data, f, indent=2)
|
|
|
|
presets_dict = {}
|
|
try:
|
|
with open(GEN_PRESETS_FILE, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
if isinstance(data, list):
|
|
for item in data:
|
|
if "name" in item:
|
|
presets_dict[item["name"]] = item
|
|
except Exception as e:
|
|
print(f"[Generation Preset Error] Failed to read settings: {e}")
|
|
return presets_dict
|
|
|
|
|
|
# ==========================================
|
|
# PROMPT NODES
|
|
# ==========================================
|
|
class PresetPromptLoader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
presets = get_presets()
|
|
preset_names = list(presets.keys()) if presets else ["No presets found"]
|
|
return {
|
|
"required": {
|
|
"positive": ("STRING", {"forceInput": True}),
|
|
"negative": ("STRING", {"forceInput": True}),
|
|
"preset_name": (preset_names, {"default": preset_names[0]}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("positive", "negative")
|
|
FUNCTION = "replace_preset_prompts"
|
|
CATEGORY = "Config Presets"
|
|
|
|
def replace_preset_prompts(self, positive, negative, preset_name):
|
|
presets = get_presets()
|
|
preset_positive = ""
|
|
preset_negative = ""
|
|
if preset_name in presets:
|
|
preset_positive = presets[preset_name]["positive"]
|
|
preset_negative = presets[preset_name]["negative"]
|
|
output_positive = (
|
|
preset_positive.replace("{prompt}", positive)
|
|
if "{prompt}" in preset_positive
|
|
else f"{preset_positive} {positive}"
|
|
)
|
|
output_negative = (
|
|
preset_negative.replace("{prompt}", negative)
|
|
if "{prompt}" in preset_negative
|
|
else f"{preset_negative} {negative}"
|
|
)
|
|
return (output_positive, output_negative)
|
|
|
|
|
|
class PresetPromptSaver:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"preset_name": ("STRING", {"default": "my-new-preset"}),
|
|
"positive_layout": (
|
|
"STRING",
|
|
{"multiline": True, "default": "masterpiece, {prompt}"},
|
|
),
|
|
"negative_layout": (
|
|
"STRING",
|
|
{"multiline": True, "default": "bad quality, {prompt}"},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_new_preset"
|
|
CATEGORY = "Config Presets"
|
|
OUTPUT_NODE = True
|
|
|
|
def save_new_preset(self, preset_name, positive_layout, negative_layout):
|
|
existing_data = []
|
|
if os.path.exists(PROMPT_PRESETS_FILE):
|
|
try:
|
|
with open(PROMPT_PRESETS_FILE, "r", encoding="utf-8") as f:
|
|
existing_data = json.load(f)
|
|
except Exception:
|
|
pass
|
|
existing_data = [
|
|
item for item in existing_data if item.get("name") != preset_name
|
|
]
|
|
existing_data.append(
|
|
{
|
|
"name": preset_name,
|
|
"positive_prompt": positive_layout,
|
|
"negative_prompt": negative_layout,
|
|
}
|
|
)
|
|
with open(PROMPT_PRESETS_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(existing_data, f, indent=2)
|
|
return {}
|
|
|
|
|
|
# ==========================================
|
|
# EXTENDED GENERATION SETTINGS NODES
|
|
# ==========================================
|
|
class PresetSettingsLoader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
presets = get_generation_presets()
|
|
preset_names = list(presets.keys()) if presets else ["No presets found"]
|
|
return {
|
|
"required": {
|
|
"preset_name": (preset_names, {"default": preset_names[0]}),
|
|
}
|
|
}
|
|
|
|
# Added the 6 new output sockets matching your requested data types
|
|
RETURN_TYPES = (
|
|
"INT",
|
|
"INT",
|
|
"FLOAT",
|
|
"INT",
|
|
"*",
|
|
"*",
|
|
"INT",
|
|
"FLOAT",
|
|
"FLOAT",
|
|
"FLOAT",
|
|
"*",
|
|
"BOOLEAN",
|
|
)
|
|
RETURN_NAMES = (
|
|
"width",
|
|
"height",
|
|
"cfg",
|
|
"steps",
|
|
"sampler_name",
|
|
"scheduler",
|
|
"hires_steps",
|
|
"hires_cfg",
|
|
"hires_denoise",
|
|
"hires_upscale_by",
|
|
"hires_model",
|
|
"is_latent",
|
|
)
|
|
FUNCTION = "load_settings"
|
|
CATEGORY = "Config Presets"
|
|
|
|
def load_settings(self, preset_name):
|
|
presets = get_generation_presets()
|
|
|
|
# Safe default fallback structure
|
|
p = presets.get(preset_name, {})
|
|
|
|
return (
|
|
int(p.get("width", 512)),
|
|
int(p.get("height", 512)),
|
|
float(p.get("cfg", 7.0)),
|
|
int(p.get("steps", 20)),
|
|
str(p.get("sampler_name", "euler")),
|
|
str(p.get("scheduler", "normal")),
|
|
# New Hi-Res Outputs
|
|
int(p.get("hires_steps", 10)),
|
|
float(p.get("hires_cfg", 7.0)),
|
|
float(p.get("hires_denoise", 0.5)),
|
|
float(p.get("hires_upscale_by", 2.0)),
|
|
str(p.get("hires_model", "None")),
|
|
bool(p.get("is_latent", True)),
|
|
)
|
|
|
|
|
|
class PresetSettingsSaver:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Dynamically scan ComfyUI's upscaler models folder
|
|
upscale_models = folder_paths.get_filename_list("upscale_models")
|
|
if not upscale_models:
|
|
upscale_models = ["None found"]
|
|
|
|
return {
|
|
"required": {
|
|
"preset_name": ("STRING", {"default": "my-generation-preset"}),
|
|
"width": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
|
|
"height": ("INT", {"default": 512, "min": 64, "max": 8192, "step": 8}),
|
|
"cfg": (
|
|
"FLOAT",
|
|
{"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1},
|
|
),
|
|
"steps": ("INT", {"default": 20, "min": 1, "max": 10000, "step": 1}),
|
|
"sampler_name": (
|
|
comfy.samplers.KSampler.SAMPLERS,
|
|
{"default": "euler"},
|
|
),
|
|
"scheduler": (
|
|
comfy.samplers.KSampler.SCHEDULERS,
|
|
{"default": "normal"},
|
|
),
|
|
# New Hi-Res Options
|
|
"hires_steps": (
|
|
"INT",
|
|
{"default": 10, "min": 1, "max": 10000, "step": 1},
|
|
),
|
|
"hires_cfg": (
|
|
"FLOAT",
|
|
{"default": 7.0, "min": 0.0, "max": 100.0, "step": 0.1},
|
|
),
|
|
"hires_denoise": (
|
|
"FLOAT",
|
|
{"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01},
|
|
),
|
|
"hires_upscale_by": (
|
|
"FLOAT",
|
|
{"default": 2.0, "min": 0.1, "max": 10.0, "step": 0.05},
|
|
),
|
|
"hires_model": (upscale_models, {"default": upscale_models[0]}),
|
|
"is_latent": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_settings"
|
|
CATEGORY = "Config Presets"
|
|
OUTPUT_NODE = True
|
|
|
|
def save_settings(
|
|
self,
|
|
preset_name,
|
|
width,
|
|
height,
|
|
cfg,
|
|
steps,
|
|
sampler_name,
|
|
scheduler,
|
|
hires_steps,
|
|
hires_cfg,
|
|
hires_denoise,
|
|
hires_upscale_by,
|
|
hires_model,
|
|
is_latent,
|
|
):
|
|
existing_data = []
|
|
if os.path.exists(GEN_PRESETS_FILE):
|
|
try:
|
|
with open(GEN_PRESETS_FILE, "r", encoding="utf-8") as f:
|
|
existing_data = json.load(f)
|
|
if not isinstance(existing_data, list):
|
|
existing_data = []
|
|
except Exception:
|
|
existing_data = []
|
|
|
|
# Remove duplicate configurations if overwriting an existing name
|
|
existing_data = [
|
|
item for item in existing_data if item.get("name") != preset_name
|
|
]
|
|
|
|
new_entry = {
|
|
"name": preset_name,
|
|
"width": width,
|
|
"height": height,
|
|
"cfg": cfg,
|
|
"steps": steps,
|
|
"sampler_name": sampler_name,
|
|
"scheduler": scheduler,
|
|
# Save the new properties to JSON
|
|
"hires_steps": hires_steps,
|
|
"hires_cfg": hires_cfg,
|
|
"hires_denoise": hires_denoise,
|
|
"hires_upscale_by": hires_upscale_by,
|
|
"hires_model": hires_model,
|
|
"is_latent": is_latent,
|
|
}
|
|
existing_data.append(new_entry)
|
|
|
|
with open(GEN_PRESETS_FILE, "w", encoding="utf-8") as f:
|
|
json.dump(existing_data, f, indent=2)
|
|
|
|
print(f"🎉 Saved generation preset '{preset_name}' with Hi-Res profiles!")
|
|
return {}
|