71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
import tomllib
|
|
import os
|
|
from typing import Any, Dict, Union, List
|
|
from tools.utils import ToolError
|
|
|
|
# Paths to config files
|
|
MODEL_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "model_presets.toml")
|
|
RES_PRESETS_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "resolution_presets.toml")
|
|
|
|
def load_toml(path: str) -> Dict[str, Any]:
|
|
try:
|
|
with open(path, "rb") as f:
|
|
return tomllib.load(f)
|
|
except Exception as e:
|
|
raise ToolError(f"Failed to load config file {path}: {str(e)}")
|
|
|
|
async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Returns information about available models or detailed guides for a specific model.
|
|
"""
|
|
model_name = args.get("model_name")
|
|
|
|
models = load_toml(MODEL_PRESETS_PATH)
|
|
res_presets = load_toml(RES_PRESETS_PATH)
|
|
|
|
if not model_name:
|
|
# Return a catalog of all models
|
|
catalog = []
|
|
for name, data in models.items():
|
|
catalog.append({
|
|
"name": name,
|
|
"description": data.get("description", "No description provided.")
|
|
})
|
|
return {
|
|
"text": "Available models:\n\n" + "\n".join([f"- {m['name']}: {m['description']}" for m in catalog]) +
|
|
"\n\nTo get a detailed prompting guide and available resolutions for a specific model, call this tool again with the 'model_name' argument."
|
|
}
|
|
|
|
if model_name not in models:
|
|
raise ToolError(f"Model '{model_name}' not found in presets. Available models: {', '.join(models.keys())}")
|
|
|
|
model_data = models[model_name]
|
|
res_set_name = model_data.get("Resolution Set")
|
|
|
|
# Get available resolution presets for this model
|
|
res_options = []
|
|
if res_set_name and res_set_name in res_presets:
|
|
res_options = list(res_presets[res_set_name].keys())
|
|
else:
|
|
res_options = ["Default (1024x1024)"]
|
|
|
|
guide = model_data.get("guide", "No detailed guide available.")
|
|
description = model_data.get("description", "")
|
|
|
|
# Filter out internal config keys to show only the human-friendly presets
|
|
presets_to_show = {k: v for k, v in model_data.items() if k not in ["description", "guide", "Resolution Set"]}
|
|
|
|
preset_text = "\n".join([f"- {k}: {v}" for k, v in presets_to_show.items()])
|
|
res_text = ", ".join(res_options)
|
|
|
|
return {
|
|
"text": (
|
|
f"Model: {model_name}\n"
|
|
f"Description: {description}\n\n"
|
|
f"--- Prompting Guide ---\n{guide}\n\n"
|
|
f"--- Active Presets ---\n{preset_text}\n\n"
|
|
f"--- Available Resolution Presets ---\n{res_text}\n\n"
|
|
f"Use 'resolution_preset' in generate_image to choose one of these."
|
|
)
|
|
}
|