Spaces:
Runtime error
Runtime error
# filename: prompt_builder.py | |
import os | |
from collections import defaultdict | |
class PromptBuilder: | |
def __init__(self, options_dir="prompt_options"): | |
self.options_dir = options_dir | |
self.options = self._load_options() | |
def _load_options(self): | |
""" | |
Dynamically load all .txt files in the options directory into a dictionary. | |
Group fields by prefix (e.g., 'character.gender' or 'background.sky'). | |
""" | |
options = defaultdict(dict) | |
if not os.path.exists(self.options_dir): | |
return options | |
for filename in os.listdir(self.options_dir): | |
if filename.endswith(".txt"): | |
path = os.path.join(self.options_dir, filename) | |
key = filename.replace(".txt", "") | |
if '.' in key: | |
group, field = key.split('.', 1) | |
with open(path, "r", encoding="utf-8") as f: | |
options[group][field] = [line.strip() for line in f if line.strip()] | |
return options | |
def get_choices(self, group, field): | |
"""Retrieve available choices for a given group/field.""" | |
return self.options.get(group, {}).get(field, []) | |
def build_prompt(self, base_prompt="", custom_tags=None, **field_values): | |
""" | |
Build prompt using: | |
- base_prompt: user base | |
- custom_tags: freeform additional prompt tags | |
- field_values: key-value pairs like gender="female", lighting="soft", etc. | |
""" | |
parts = [base_prompt] | |
for key, value in field_values.items(): | |
if not value: | |
continue | |
if key in ("styles", "lighting", "mood"): | |
# general tags with descriptors | |
if key == "styles": | |
parts.append(f"in {value} style") | |
elif key == "lighting": | |
parts.append(f"with {value} lighting") | |
elif key == "mood": | |
parts.append(f"evoking a {value} mood") | |
else: | |
# other tags are descriptive or noun-like | |
parts.append(value) | |
if custom_tags: | |
parts.append(custom_tags) | |
return ", ".join(filter(None, parts)) | |
builder = PromptBuilder() | |
builder.get_choices("character", "gender") # β ['male', 'female', 'nonbinary'] | |
prompt = builder.build_prompt( | |
base_prompt="a portrait", | |
gender="female", body="slim", clothing="gothic dress", | |
land_type="desert", sky="stormy", styles="anime" | |
) | |
print(prompt) |