123 lines
3.6 KiB
Python
123 lines
3.6 KiB
Python
import os
|
|
import json
|
|
import re
|
|
import pandas as pd
|
|
import comfy.samplers # Taps into ComfyUI's native sampler/scheduler listings
|
|
import folder_paths # Taps into ComfyUI's native model folder tracking
|
|
|
|
STYLES_DIR = os.path.join(os.path.dirname(__file__), "styles")
|
|
CONFIGS_DIR = os.path.join(os.path.dirname(__file__), "configs")
|
|
GEN_PRESETS_FILE = os.path.join(CONFIGS_DIR, "generation_presets.json")
|
|
|
|
|
|
def ensure_styles_dir():
|
|
os.makedirs(STYLES_DIR, exist_ok=True)
|
|
|
|
|
|
def get_style_files() -> list[str]:
|
|
ensure_styles_dir()
|
|
files = sorted(
|
|
f
|
|
for f in os.listdir(STYLES_DIR)
|
|
if f.lower().endswith((".csv", ".xlsx", ".xls"))
|
|
)
|
|
return files if files else ["none"]
|
|
|
|
|
|
def load_styles(filename: str) -> dict:
|
|
"""Return {name: {prompt, negative_prompt, description, category}}."""
|
|
|
|
if not filename or filename == "none":
|
|
return {}
|
|
|
|
path = os.path.join(STYLES_DIR, filename)
|
|
|
|
if not os.path.exists(path):
|
|
return {}
|
|
|
|
ext = os.path.splitext(filename)[1].lower()
|
|
|
|
if ext == ".csv":
|
|
df = pd.read_csv(path, encoding="utf-8-sig")
|
|
elif ext in (".xlsx", ".xls"):
|
|
df = pd.read_excel(path)
|
|
else:
|
|
return {}
|
|
|
|
styles = {}
|
|
|
|
for _, row in df.iterrows():
|
|
name = str(row.get("name", "")).strip()
|
|
if not name:
|
|
continue
|
|
|
|
styles[name] = {
|
|
"prompt": str(row.get("prompt", "")).strip(),
|
|
"negative_prompt": str(row.get("negative_prompt", "")).strip(),
|
|
"description": str(row.get("description", "")).strip(),
|
|
"category": str(row.get("category", "")).strip(),
|
|
}
|
|
|
|
return styles
|
|
|
|
|
|
def get_style_names(csv_filename: str) -> list[str]:
|
|
names = list(load_styles(csv_filename).keys())
|
|
return names if names else ["none"]
|
|
|
|
|
|
def apply_style_text(input_text: str, style_text: str) -> str:
|
|
"""
|
|
If style_text contains {prompt}, replace it with input_text.
|
|
Otherwise, append input_text to the end (comma-separated).
|
|
"""
|
|
if not style_text:
|
|
return input_text
|
|
if re.search(r"(?i)\{prompt\}", style_text):
|
|
return re.sub(r"(?i)\{prompt\}", input_text, style_text)
|
|
if input_text:
|
|
return f"{style_text}, {input_text}"
|
|
return style_text
|
|
|
|
|
|
def ensure_configs_dir():
|
|
os.makedirs(CONFIGS_DIR, exist_ok=True)
|
|
|
|
|
|
def get_generation_presets():
|
|
ensure_configs_dir()
|
|
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": "dreamshaper 8",
|
|
"width": 512,
|
|
"height": 768,
|
|
"cfg": 8.0,
|
|
"steps": 30,
|
|
"sampler_name": "dpmpp_sde",
|
|
"scheduler": "karras",
|
|
"hires_steps": 15,
|
|
"hires_cfg": 5.5,
|
|
"hires_denoise": 0.4,
|
|
"hires_upscale_by": 2.0,
|
|
"hires_model": default_model,
|
|
"is_latent": False,
|
|
}
|
|
]
|
|
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
|