113 lines
2.9 KiB
Python
113 lines
2.9 KiB
Python
import uuid
|
|
from dataclasses import dataclass
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RatingCriterion:
|
|
key: str
|
|
label: str
|
|
description: str
|
|
|
|
|
|
RATING_CRITERIA: list[RatingCriterion] = [
|
|
RatingCriterion(
|
|
key="prompt_adherence",
|
|
label="Prompt Adherence",
|
|
description="Subject, attributes, and composition match the prompt exactly",
|
|
),
|
|
RatingCriterion(
|
|
key="aesthetics",
|
|
label="Aesthetics",
|
|
description="Visual beauty, color harmony, and overall appeal",
|
|
),
|
|
RatingCriterion(
|
|
key="technical_quality",
|
|
label="Technical Quality",
|
|
description="Sharpness, clean render, and lack of artifacts",
|
|
),
|
|
RatingCriterion(
|
|
key="lighting",
|
|
label="Lighting & Atmosphere",
|
|
description="Light quality, shadow detail, highlights, and mood",
|
|
),
|
|
RatingCriterion(
|
|
key="detail_richness",
|
|
label="Detail & Texture",
|
|
description="Fine-grained surface textures and element detail",
|
|
),
|
|
RatingCriterion(
|
|
key="coherence",
|
|
label="Coherence",
|
|
description="Correct anatomy, physics, and spatial relationships",
|
|
),
|
|
]
|
|
|
|
RATING_OPTIONS = [
|
|
"1 - Poor",
|
|
"2 - Weak",
|
|
"3 - Okay",
|
|
"4 - Good",
|
|
"5 - Excellent",
|
|
]
|
|
|
|
_RATING_VALUE_MAP = {
|
|
option: index
|
|
for index, option in enumerate(RATING_OPTIONS, start=1)
|
|
}
|
|
|
|
|
|
class PBRating:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
required = {
|
|
"notes": (
|
|
"STRING",
|
|
{
|
|
"multiline": True,
|
|
"default": "",
|
|
"tooltip": "Optional notes about why this image was rated this way.",
|
|
},
|
|
)
|
|
}
|
|
|
|
for criterion in RATING_CRITERIA:
|
|
required[criterion.key] = (
|
|
"COMBO",
|
|
{
|
|
"options": RATING_OPTIONS,
|
|
"default": RATING_OPTIONS[2],
|
|
"tooltip": criterion.description,
|
|
},
|
|
)
|
|
|
|
return {"required": required}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return str(uuid.uuid4())
|
|
|
|
RETURN_TYPES = ("RATING",)
|
|
RETURN_NAMES = ("rating",)
|
|
FUNCTION = "rate_prompt"
|
|
CATEGORY = "Prompt Bench"
|
|
|
|
def rate_prompt(self, notes: str = "", **ratings: str):
|
|
rating_id = str(uuid.uuid4())
|
|
|
|
criterion_scores: dict[str, int] = {}
|
|
|
|
for criterion in RATING_CRITERIA:
|
|
selected_option = ratings.get(criterion.key, None)
|
|
if selected_option:
|
|
criterion_scores[criterion.key] = _RATING_VALUE_MAP[selected_option]
|
|
|
|
average_score = sum(criterion_scores.values()) / len(RATING_CRITERIA)
|
|
payload = {
|
|
"rating_id": rating_id,
|
|
"notes": notes,
|
|
"criteria": criterion_scores,
|
|
"average_score": average_score,
|
|
}
|
|
|
|
return (payload,)
|