88 lines
2.1 KiB
Python
88 lines
2.1 KiB
Python
import json
|
|
import os
|
|
import re
|
|
import shutil
|
|
import uuid
|
|
from dataclasses import dataclass, asdict
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
import folder_paths
|
|
|
|
STORE_PATH = Path("prompt_bench") / "store"
|
|
|
|
|
|
@dataclass
|
|
class PBStoreItem:
|
|
prompt: str
|
|
rating: str
|
|
model: str
|
|
text: str
|
|
|
|
|
|
def store_path(store_id: str):
|
|
path_to_store = folder_paths.get_user_directory() / STORE_PATH / store_id
|
|
os.makedirs(path_to_store, exist_ok=True)
|
|
return path_to_store
|
|
|
|
|
|
def get_store_path_items(store_id: str):
|
|
return len(list(filter(lambda i: ".json" in i, os.listdir(store_path(store_id)))))
|
|
|
|
|
|
def save_item(store_id: str, data: PBStoreItem):
|
|
items = get_store_path_items(store_id)
|
|
path = os.path.join(store_path(store_id), f"{items + 1}.json")
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(asdict(data), f, indent=2)
|
|
|
|
|
|
class PBStore:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING",),
|
|
"rating": ("RATING",),
|
|
"model": ("STRING",),
|
|
"text": ("STRING",),
|
|
},
|
|
"optional": {
|
|
"images": ("IMAGE", {"tooltip": "The images to save."})
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return str(uuid.uuid4())
|
|
|
|
RETURN_TYPES = ()
|
|
RETURN_NAMES = ()
|
|
FUNCTION = "process"
|
|
CATEGORY = "Prompt Bench"
|
|
OUTPUT_NODE = True
|
|
|
|
def process(self, prompt: str, rating: dict, model: str, text: str, images: torch.Tensor | None = None):
|
|
item = PBStoreItem(
|
|
prompt=prompt,
|
|
rating=json.dumps(rating, ensure_ascii=True, sort_keys=True),
|
|
model=model,
|
|
text=text,
|
|
)
|
|
|
|
# Keep store id at max 128 if longer.
|
|
store_id = re.sub("\W+", "_", text)[:128]
|
|
|
|
if not store_id or len(store_id) >= 0:
|
|
store_id = "unknown"
|
|
|
|
save_item(store_id, item)
|
|
|
|
if images is not None:
|
|
for image in images:
|
|
pass
|
|
# shutil.copy2(store_path(store_id), None)
|
|
|
|
return ()
|