diff --git a/src/pruna/data/__init__.py b/src/pruna/data/__init__.py index fd14a496..1f0ed5f6 100644 --- a/src/pruna/data/__init__.py +++ b/src/pruna/data/__init__.py @@ -34,7 +34,13 @@ setup_hps_dataset, setup_imgedit_dataset, setup_long_text_bench_dataset, + setup_oneig_anime_stylization_dataset, setup_oneig_dataset, + setup_oneig_general_object_dataset, + setup_oneig_knowledge_reasoning_dataset, + setup_oneig_multilingualism_dataset, + setup_oneig_portrait_dataset, + setup_oneig_text_rendering_dataset, setup_parti_prompts_dataset, ) from pruna.data.datasets.question_answering import setup_polyglot_dataset @@ -103,19 +109,33 @@ "image_classification_collate", {"img_size": 224}, ), - "DrawBench": (setup_drawbench_dataset, "prompt_collate", {}), + "DrawBench": (setup_drawbench_dataset, "prompt_with_auxiliaries_collate", {}), "PartiPrompts": ( setup_parti_prompts_dataset, "prompt_with_auxiliaries_collate", {}, ), - "GenAIBench": (setup_genai_bench_dataset, "prompt_collate", {}), + "GenAIBench": (setup_genai_bench_dataset, "prompt_with_auxiliaries_collate", {}), "GenEval": (setup_geneval_dataset, "prompt_with_auxiliaries_collate", {}), "HPS": (setup_hps_dataset, "prompt_with_auxiliaries_collate", {}), "ImgEdit": (setup_imgedit_dataset, "prompt_with_auxiliaries_collate", {}), "LongTextBench": (setup_long_text_bench_dataset, "prompt_with_auxiliaries_collate", {}), "GEditBench": (setup_gedit_dataset, "prompt_with_auxiliaries_collate", {}), "OneIG": (setup_oneig_dataset, "prompt_with_auxiliaries_collate", {}), + "OneIGAnimeStylization": ( + setup_oneig_anime_stylization_dataset, + "prompt_with_auxiliaries_collate", + {}, + ), + "OneIGGeneralObject": (setup_oneig_general_object_dataset, "prompt_with_auxiliaries_collate", {}), + "OneIGKnowledgeReasoning": ( + setup_oneig_knowledge_reasoning_dataset, + "prompt_with_auxiliaries_collate", + {}, + ), + "OneIGMultilingualism": (setup_oneig_multilingualism_dataset, "prompt_with_auxiliaries_collate", {}), + "OneIGPortrait": (setup_oneig_portrait_dataset, "prompt_with_auxiliaries_collate", {}), + "OneIGTextRendering": (setup_oneig_text_rendering_dataset, "prompt_with_auxiliaries_collate", {}), "DPG": (setup_dpg_dataset, "prompt_with_auxiliaries_collate", {}), "TinyIMDB": (setup_tiny_imdb_dataset, "text_generation_collate", {}), "VBench": (setup_vbench_dataset, "prompt_with_auxiliaries_collate", {}), diff --git a/src/pruna/data/datasets/prompt.py b/src/pruna/data/datasets/prompt.py index 7764d23b..b18d03d4 100644 --- a/src/pruna/data/datasets/prompt.py +++ b/src/pruna/data/datasets/prompt.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from pathlib import Path from typing import Literal, Tuple, get_args from datasets import Dataset, load_dataset @@ -123,21 +124,95 @@ DPGCategory = Literal["entity", "attribute", "relation", "global", "other"] -def _to_oneig_record(row: dict, questions_by_key: dict[str, dict]) -> dict: - """Convert OneIG row to unified record format.""" +def _warn_ignored_benchmark_seed(seed: int | None, *, dataset: str) -> None: + if seed is not None: + pruna_logger.warning( + "%s: `seed` is ignored for this test-only benchmark; sampling does not shuffle the test split.", + dataset, + ) + + +def _oneig_alignment_language_zh(row: dict) -> bool: + """Return True when the official Q_D file for this row should use the ``*_zh`` graphs.""" + if row.get("category", "") == "Multilingualism": + return True + lang = row.get("language") or row.get("lang") + return isinstance(lang, str) and lang.lower() in {"zh", "zh-cn", "zh_cn", "chinese", "cn"} + + +def _oneig_qd_prefix(row: dict) -> str: + """Map dataset ``category`` (+ language) to Q_D JSON stem (e.g. ``object``, ``anime_zh``).""" + row_category = row.get("category", "") + use_zh = _oneig_alignment_language_zh(row) + if row_category == "Multilingualism": + return "multilingualism_zh" + base = _CATEGORY_TO_QD.get(row_category, "") + if not base: + return "" + return f"{base}_zh" if use_zh else base + + +def _to_oneig_record( + row: dict, + questions_by_key: dict[str, dict], + reasoning_gt_en: dict[str, str], + reasoning_gt_zh: dict[str, str], + reasoning_language: str = "EN", +) -> dict: + """Convert OneIG row to unified record format. + + Parameters + ---------- + row : dict + Raw Hugging Face row (``category``, ``id``, ``class``). EN configs use ``prompt_en``; the + ``OneIG-Bench-ZH`` **Multilingualism** split uses ``prompt_cn`` instead of ``prompt_en``. + questions_by_key : dict[str, dict] + Merged Q_D index keyed as ``{qd_stem}_{prompt_id}`` (see ``_fetch_oneig_alignment``). + reasoning_gt_en : dict[str, str] + Official ``gt_answer.json`` keyed by prompt id (e.g. ``"000"``). + reasoning_gt_zh : dict[str, str] + Official ``gt_answer_zh.json`` keyed by prompt id. + reasoning_language : str, optional + Which reasoning GT to use: ``"EN"`` or ``"ZH"``. Default is ``"EN"``. + + Returns + ------- + dict + Unified record including ``questions``, ``dependencies``, and ``reasoning_gt_answer`` when + applicable (Knowledge_Reasoning only). + """ row_category = row.get("category", "") row_class = row.get("class", "None") or "None" - qd_name = _CATEGORY_TO_QD.get(row_category, "") - lookup_key = f"{qd_name}_{row.get('id', '')}" if qd_name else "" + prompt_id = str(row.get("id", "")) + qd_prefix = _oneig_qd_prefix(row) + lookup_key = f"{qd_prefix}_{prompt_id}" if qd_prefix else "" q_info = questions_by_key.get(lookup_key, {}) + text = row.get("prompt") or row.get("prompt_en") or row.get("prompt_cn") or "" + reasoning_gt_answer: str | None = None + if row_category == "Knowledge_Reasoning": + if reasoning_language.upper() == "ZH": + reasoning_gt_answer = reasoning_gt_zh.get(prompt_id) + else: + reasoning_gt_answer = reasoning_gt_en.get(prompt_id) + is_text_rendering = row_category in ("Text_Rendering", "Text Rendering") + if is_text_rendering and text: + import re as _re + + quoted = _re.findall(r'"([^"]+)"', text) + text_content: str | None = " ".join(quoted) if quoted else (row_class if row_class != "None" else None) + else: + text_content = row_class if row_class != "None" else None + questions = {k: v for k, v in q_info.get("questions", {}).items() if v is not None} + dependencies = {k: v for k, v in q_info.get("dependencies", {}).items() if v is not None} return { - "text": row.get("prompt_en", row.get("prompt", "")), - "subset": "Text_Rendering" if row_category in ("Text_Rendering", "Text Rendering") else row_category, - "text_content": row_class if row_class != "None" else None, + "text": text, + "subset": "Text_Rendering" if is_text_rendering else row_category, + "text_content": text_content, "category": row_category, "class": row_class, - "questions": q_info.get("questions", {}), - "dependencies": q_info.get("dependencies", {}), + "questions": questions, + "dependencies": dependencies, + "reasoning_gt_answer": reasoning_gt_answer, } @@ -159,7 +234,7 @@ def setup_drawbench_dataset() -> Tuple[Dataset, Dataset, Dataset]: def setup_parti_prompts_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -172,8 +247,8 @@ def setup_parti_prompts_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -188,6 +263,7 @@ def setup_parti_prompts_dataset( Tuple[Dataset, Dataset, Dataset] The Parti Prompts dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="PartiPrompts") ds = load_dataset("nateraw/parti-prompts")["train"] # type: ignore[index] if category is not None: @@ -226,7 +302,7 @@ def _generate_geneval_question(entry: dict) -> list[str]: def setup_geneval_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -239,8 +315,8 @@ def setup_geneval_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -255,6 +331,7 @@ def setup_geneval_dataset( Tuple[Dataset, Dataset, Dataset] The GenEval dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="GenEval") import json import requests @@ -286,7 +363,7 @@ def setup_geneval_dataset( def setup_hps_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -299,8 +376,8 @@ def setup_hps_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -315,6 +392,7 @@ def setup_hps_dataset( Tuple[Dataset, Dataset, Dataset] The HPD dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="HPS") import json from huggingface_hub import hf_hub_download @@ -338,7 +416,7 @@ def setup_hps_dataset( def setup_long_text_bench_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -350,8 +428,8 @@ def setup_long_text_bench_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -364,6 +442,7 @@ def setup_long_text_bench_dataset( Tuple[Dataset, Dataset, Dataset] The Long Text Bench dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="LongTextBench") ds = load_dataset("X-Omni/LongText-Bench")["train"] # type: ignore[index] ds = ds.rename_column("text", "text_content") ds = ds.rename_column("prompt", "text") @@ -389,8 +468,74 @@ def setup_genai_bench_dataset() -> Tuple[Dataset, Dataset, Dataset]: return ds.select([0]), ds.select([0]), ds +def ensure_imgedit_benchmark_images_extracted() -> Path: + """ + Download ``Benchmark.tar`` (if needed), extract it, and return the ``singleturn`` folder. + + Returns + ------- + Path + ``Benchmark/singleturn`` directory whose files are addressed by ``image_id``. + + Raises + ------ + RuntimeError + If the archive cannot be downloaded or extracted. + """ + import tarfile + + from huggingface_hub import hf_hub_download + + tar_path = Path(hf_hub_download(repo_id="sysuyy/ImgEdit", filename="Benchmark.tar", repo_type="dataset")) + extract_dir = tar_path.parent / "imgedit_singleturn" + candidate = extract_dir / "Benchmark" / "singleturn" + if not candidate.is_dir() or not any(candidate.iterdir()): + extract_dir.mkdir(parents=True, exist_ok=True) + with tarfile.open(tar_path, "r") as tar: + tar.extractall(path=extract_dir) + if not candidate.is_dir() or not any(candidate.iterdir()): + raise RuntimeError(f"ImgEdit: failed to extract Benchmark.tar to {candidate}") + return candidate + + +def load_imgedit_source_image_bytes(image_id: str, *, image_folder: Path | None = None) -> bytes: + """ + Read one ImgEdit source image as JPEG bytes (RGB). + + Parameters + ---------- + image_id : str + Path relative to the singleturn folder (from the official ``basic_edit.json`` ``id``). + image_folder : Path | None, optional + ``Benchmark/singleturn`` directory; when ``None``, calls + :func:`ensure_imgedit_benchmark_images_extracted`. + + Returns + ------- + bytes + JPEG-encoded bytes for the source image. + + Raises + ------ + FileNotFoundError + If ``image_id`` does not exist under ``image_folder``. + Exception + If PIL cannot open or convert the image. + """ + from io import BytesIO + + from PIL import Image + + folder = image_folder if image_folder is not None else ensure_imgedit_benchmark_images_extracted() + img_path = folder / image_id + pil = Image.open(img_path).convert("RGB") + buf = BytesIO() + pil.save(buf, format="JPEG") + return buf.getvalue() + + def setup_imgedit_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -403,8 +548,8 @@ def setup_imgedit_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -420,6 +565,7 @@ def setup_imgedit_dataset( Tuple[Dataset, Dataset, Dataset] The ImgEdit dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="ImgEdit") import json import requests @@ -433,6 +579,8 @@ def setup_imgedit_dataset( instructions: dict = json.loads(response_instructions.text) judge_prompts: dict = json.loads(response_judge_prompts.text) + image_folder = ensure_imgedit_benchmark_images_extracted() + categories = [category] if category is not None and not isinstance(category, list) else category records = [] for _, instruction in instructions.items(): @@ -441,14 +589,17 @@ def setup_imgedit_dataset( if categories is not None and edit_type not in categories: continue - records.append( - { - "text": instruction.get("prompt", ""), - "category": edit_type, - "image_id": instruction.get("id", ""), - "judge_prompt": judge_prompts.get(edit_type, ""), - } - ) + image_id = instruction.get("id", "") + record: dict = { + "text": instruction.get("prompt", ""), + "category": edit_type, + "image_id": image_id, + "judge_prompt": judge_prompts.get(edit_type, ""), + } + src = load_imgedit_source_image_bytes(image_id, image_folder=image_folder) + if src is not None: + record["source_image_bytes"] = src + records.append(record) ds = Dataset.from_list(records) ds = stratify_dataset(ds, sample_size=test_sample_size, fraction=fraction) @@ -466,18 +617,47 @@ def setup_imgedit_dataset( "General_Object": "object", } -_ONEIG_ALIGNMENT_BASE = "https://raw.githubusercontent.com/OneIG-Bench/OneIG-Benchmark/41b49831e79e6dde5323618c164da1c4cf0f699d/scripts/alignment/Q_D" +_ONEIG_BENCHMARK_REF = "41b49831e79e6dde5323618c164da1c4cf0f699d" +_ONEIG_RAW_BASE = f"https://raw.githubusercontent.com/OneIG-Bench/OneIG-Benchmark/{_ONEIG_BENCHMARK_REF}" +_ONEIG_ALIGNMENT_QD_URL = f"{_ONEIG_RAW_BASE}/scripts/alignment/Q_D" +_ONEIG_REASONING_GT_URL_EN = f"{_ONEIG_RAW_BASE}/scripts/reasoning/gt_answer.json" +_ONEIG_REASONING_GT_URL_ZH = f"{_ONEIG_RAW_BASE}/scripts/reasoning/gt_answer_zh.json" + +_ONEIG_QD_JSON_STEMS: tuple[str, ...] = ( + "anime", + "human", + "object", + "anime_zh", + "human_zh", + "object_zh", + "multilingualism_zh", +) def _fetch_oneig_alignment() -> dict[str, dict]: - """Fetch alignment questions from per-category Q_D files (InferBench-style).""" + """Load OneIG question/dependency graphs from the official repo (HTTP, no on-disk cache). + + Fetches every ``scripts/alignment/Q_D/*.json`` file used by upstream ``alignment_score.py`` (EN + ZH), + including ``multilingualism_zh.json``. Keys in the returned map are ``{stem}_{prompt_id}`` matching + upstream file stems (e.g. ``object_012``, ``multilingualism_zh_000``). + + Returns + ------- + dict[str, dict] + ``prompt_id``-level ``questions`` and ``dependencies`` dicts (parsed from JSON strings when needed). + + Raises + ------ + requests.HTTPError + If any asset URL is missing or the response is not successful. + """ import json import requests questions_by_key: dict[str, dict] = {} - for qd_name in ("anime", "human", "object"): - url = f"{_ONEIG_ALIGNMENT_BASE}/{qd_name}.json" + for stem in _ONEIG_QD_JSON_STEMS: + url = f"{_ONEIG_ALIGNMENT_QD_URL}/{stem}.json" resp = requests.get(url, timeout=30) resp.raise_for_status() data = json.loads(resp.text) @@ -488,16 +668,55 @@ def _fetch_oneig_alignment() -> dict[str, dict]: q = json.loads(q) if isinstance(d, str): d = json.loads(d) - questions_by_key[f"{qd_name}_{prompt_id}"] = {"questions": q, "dependencies": d} + questions_by_key[f"{stem}_{prompt_id}"] = {"questions": q, "dependencies": d} return questions_by_key +def _fetch_oneig_reasoning_gt() -> tuple[dict[str, str], dict[str, str]]: + """Load official knowledge-reasoning reference answers (HTTP, no on-disk cache). + + Mirrors ``scripts/reasoning/gt_answer.json`` and ``gt_answer_zh.json`` from the same pinned commit as Q_D. + Keys are prompt ids (``str``), values are answer strings; downstream metrics may slice filenames to the + first three characters like ``reasoning_score.py``. + + Returns + ------- + tuple[dict[str, str], dict[str, str]] + ``(en_by_id, zh_by_id)``. + + Raises + ------ + requests.HTTPError + If any asset URL is missing or the response is not successful. + """ + import json + + import requests + + def _load(url: str) -> dict[str, str]: + resp = requests.get(url, timeout=60) + resp.raise_for_status() + raw = json.loads(resp.text) + return {str(k): str(v) for k, v in raw.items()} + + return _load(_ONEIG_REASONING_GT_URL_EN), _load(_ONEIG_REASONING_GT_URL_ZH) + + +def _oneig_needs_zh_multilingualism_hub(category: OneIGCategory | list[OneIGCategory] | None) -> bool: + """Whether ``OneIG-Bench-ZH`` must be loaded for ``Multilingualism`` rows.""" + if category is None: + return True + categories = [category] if not isinstance(category, list) else category + return "Multilingualism" in categories + + def setup_oneig_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, category: OneIGCategory | list[OneIGCategory] | None = None, + reasoning_language: str = "EN", ) -> Tuple[Dataset, Dataset, Dataset]: """ Setup the OneIG benchmark dataset. @@ -506,8 +725,8 @@ def setup_oneig_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -517,16 +736,43 @@ def setup_oneig_dataset( category : OneIGCategory | list[OneIGCategory] | None Filter by dataset category (Anime_Stylization, Portrait, etc.) or class (fauvism, watercolor, etc.). If None, returns all subsets. + reasoning_language : str, optional + Which reasoning GT to use for Knowledge_Reasoning rows: ``"EN"`` or ``"ZH"``. Default is ``"EN"``. Returns ------- Tuple[Dataset, Dataset, Dataset] - The OneIG dataset (dummy train, dummy val, test). + The OneIG dataset (dummy train, dummy val, test). Rows include ``questions`` and + ``dependencies`` from official Q_D JSON (EN + ZH stems, including ``multilingualism_zh``), + plus ``reasoning_gt_answer`` for ``Knowledge_Reasoning`` (language chosen by ``reasoning_language``). + Rows cover EN categories from ``OneIG-Bench`` plus ``Multilingualism`` from ``OneIG-Bench-ZH``. + Assets are downloaded over HTTP on each call (pinned commit ``_ONEIG_BENCHMARK_REF``); there is + no local disk cache. + + Notes + ----- + Non-multilingual prompts are loaded from the Hub config ``OneIG-Bench``; **Multilingualism** rows + are taken only from ``OneIG-Bench-ZH`` (they use ``prompt_cn``). The ZH config is fetched only when + the requested ``category`` is ``None`` (full suite) or explicitly includes ``Multilingualism``. + Q_D / reasoning JSON URLs are defined next to ``_fetch_oneig_alignment`` and + ``_fetch_oneig_reasoning_gt``. """ + _warn_ignored_benchmark_seed(seed, dataset="OneIG") questions_by_key = _fetch_oneig_alignment() - - ds_raw = load_dataset("OneIG-Bench/OneIG-Bench", "OneIG-Bench")["train"] # type: ignore[index] - records = [_to_oneig_record(dict(row), questions_by_key) for row in ds_raw] + reasoning_gt_en, reasoning_gt_zh = _fetch_oneig_reasoning_gt() + + ds_en = load_dataset("OneIG-Bench/OneIG-Bench", "OneIG-Bench")["train"] # type: ignore[index] + records = [ + _to_oneig_record(dict(row), questions_by_key, reasoning_gt_en, reasoning_gt_zh, reasoning_language) + for row in ds_en + ] + if _oneig_needs_zh_multilingualism_hub(category): + ds_zh = load_dataset("OneIG-Bench/OneIG-Bench", "OneIG-Bench-ZH")["train"] # type: ignore[index] + ds_zh_ml = ds_zh.filter(lambda r: r["category"] == "Multilingualism") + records.extend( + _to_oneig_record(dict(row), questions_by_key, reasoning_gt_en, reasoning_gt_zh, reasoning_language) + for row in ds_zh_ml + ) ds = Dataset.from_list(records) if category is not None: @@ -544,8 +790,252 @@ def setup_oneig_dataset( return ds.select([0]), ds.select([0]), ds +# functools.partial is not used for these wrappers: get_literal_values_from_param would unwrap +# partial objects back to setup_oneig_dataset and expose every OneIGCategory instead of one. + + +def setup_oneig_anime_stylization_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Anime_Stylization``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Anime_Stylization subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Anime_Stylization", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_general_object_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``General_Object``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the General_Object subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="General_Object", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_knowledge_reasoning_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Knowledge_Reasoning``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Knowledge_Reasoning subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Knowledge_Reasoning", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_multilingualism_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Multilingualism``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Multilingualism subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Multilingualism", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_portrait_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Portrait``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Portrait subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Portrait", + reasoning_language=reasoning_language, + ) + + +def setup_oneig_text_rendering_dataset( + seed: int | None = None, + fraction: float = 1.0, + train_sample_size: int | None = None, + test_sample_size: int | None = None, + reasoning_language: str = "EN", +) -> Tuple[Dataset, Dataset, Dataset]: + """ + Load OneIG-Bench with ``category`` fixed to ``Text_Rendering``. + + License: Apache 2.0 + + Parameters + ---------- + seed : int | None, optional + Ignored; see :func:`setup_oneig_dataset`. + fraction : float + Fraction of the subset to use. + train_sample_size : int | None + Unused; train/val are dummy. + test_sample_size : int | None + Test sample size cap for the subset. + reasoning_language : str + Passed to :func:`setup_oneig_dataset`. + + Returns + ------- + Tuple[Dataset, Dataset, Dataset] + Dummy train, dummy val, and test split for the Text_Rendering subset. + """ + return setup_oneig_dataset( + seed=seed, + fraction=fraction, + train_sample_size=train_sample_size, + test_sample_size=test_sample_size, + category="Text_Rendering", + reasoning_language=reasoning_language, + ) + + def setup_gedit_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -558,8 +1048,8 @@ def setup_gedit_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -576,6 +1066,7 @@ def setup_gedit_dataset( Tuple[Dataset, Dataset, Dataset] The GEditBench dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="GEditBench") task_type_map = { "subject_add": "subject-add", "subject_remove": "subject-remove", @@ -595,12 +1086,18 @@ def setup_gedit_dataset( for row in ds: task_type = row.get("task_type", "") category_name = task_type_to_category.get(task_type, task_type) - records.append( - { - "text": row.get("instruction", ""), - "category": category_name, - } - ) + record: dict = { + "text": row.get("instruction", ""), + "category": category_name, + } + src = row.get("input_image_raw") + if src is not None: + from io import BytesIO + + buf = BytesIO() + src.save(buf, format="JPEG") + record["source_image_bytes"] = buf.getvalue() + records.append(record) ds = Dataset.from_list(records) ds = stratify_dataset(ds, sample_size=test_sample_size, fraction=fraction) @@ -613,7 +1110,7 @@ def setup_gedit_dataset( def setup_dpg_dataset( - seed: int, + seed: int | None = None, fraction: float = 1.0, train_sample_size: int | None = None, test_sample_size: int | None = None, @@ -626,8 +1123,8 @@ def setup_dpg_dataset( Parameters ---------- - seed : int - The seed to use. + seed : int | None, optional + Ignored; test order is deterministic. If not None, a warning is logged. fraction : float The fraction of the dataset to use. train_sample_size : int | None @@ -642,6 +1139,7 @@ def setup_dpg_dataset( Tuple[Dataset, Dataset, Dataset] The DPG dataset (dummy train, dummy val, test). """ + _warn_ignored_benchmark_seed(seed, dataset="DPG") import csv import io from collections import defaultdict diff --git a/src/pruna/evaluation/benchmarks.py b/src/pruna/evaluation/benchmarks.py index 36e50c5e..cf45deba 100644 --- a/src/pruna/evaluation/benchmarks.py +++ b/src/pruna/evaluation/benchmarks.py @@ -18,6 +18,9 @@ from pruna.data.utils import get_literal_values_from_param from pruna.evaluation.metrics import MetricRegistry +TASK_TYPE_TEXT_IMAGE = "text_to_image" +TASK_TYPE_TEXT_PLUS_IMAGE_IMAGE = "text+image_image" + @dataclass class Benchmark: @@ -31,9 +34,11 @@ class Benchmark: description : str Description of what the benchmark evaluates. metrics : list[str] - List of metric names used for evaluation. + Metric names from ``MetricRegistry`` that the ``reference`` paper + explicitly names for that benchmark. task_type : str - Type of task the benchmark evaluates (e.g., 'text_to_image'). + Type of task the benchmark evaluates (e.g., ``text_to_image``, + ``text+image_image``, ``text_to_video``). reference : str | None URL to the canonical paper (e.g., arXiv) for this benchmark. """ @@ -62,24 +67,18 @@ class BenchmarkRegistry: """ Registry for benchmarks. - Metrics per benchmark are set to those explicitly used in the reference - paper (see reference URL). All entries verified from paper evaluation - sections (ar5iv/HTML or PDF) as of verification pass: + Each entry's ``metrics`` lists ``MetricRegistry`` names with a direct counterpart + in the reference paper when available; otherwise the list is empty. Paper notes: - - Parti Prompts (2206.10789 ?5.2, ?5.4): human side-by-side only on P222. - - DrawBench (2205.11487 ?4.3): human raters only; COCO uses FID + CLIP. - - GenAI Bench (2406.13743): VQAScore only (web/PWC; ar5iv failed). + - Parti Prompts (2206.10789): human side-by-side only on P222. + - DrawBench (2205.11487): human raters only; COCO uses FID + CLIP. + - GenAI Bench (2406.13743): VQAScore only. - VBench (2311.17982): 16 dimension-specific methods; no single Pruna metric. - - COCO (2205.11487 ?4.1): FID and CLIP score for fidelity and alignment. - - ImageNet (1409.0575 ?4): top-1/top-5 classification accuracy. - - WikiText (1609.07843 ?5): perplexity on validation/test. - - GenEval (2310.11513 ?3.2): Mask2Former + CLIP color pipeline, binary score. - - HPS (2306.09341): HPS v2 scoring model (CLIP fine-tuned on HPD v2). - - ImgEdit (2505.20275 ?4.2): GPT-4o 1ÿÿÿ5 ratings and ImgEdit-Judge. - - Long Text Bench (2507.22058 ?4): Text Accuracy (OCR, Qwen2.5-VL-7B). - - GEditBench (2504.17761 ?4.2): VIEScore (SQ, PQ, O via GPT-4.1/Qwen2.5-VL). - - OneIG (2506.07977 ?4.1): per-dimension metrics (semantic alignment, ED, etc.). - - DPG (2403.05135): DSG-style graph score, mPLUG-large adjudicator. + - GenEval (2310.11513): Mask2Former + CLIP color pipeline, binary score. + - Long Text Bench (2507.22058): word accuracy (X-Omni); not wired yet. + - GEditBench (2504.17761): VIEScore via GPT-4.1/Qwen2.5-VL. + - OneIG (2506.07977): alignment, text, reasoning metrics. + - DPG (2403.05135): DSG-style graph score. """ _registry: dict[str, Benchmark] = {} @@ -88,9 +87,7 @@ class BenchmarkRegistry: def _register(cls, benchmark: Benchmark) -> None: missing = [m for m in benchmark.metrics if not MetricRegistry.has_metric(m)] if missing: - raise ValueError( - f"Benchmark '{benchmark.name}' references metrics not in MetricRegistry: {missing}." - ) + raise ValueError(f"Benchmark '{benchmark.name}' references metrics not in MetricRegistry: {missing}.") if benchmark.lookup_key not in base_datasets: available = ", ".join(base_datasets.keys()) raise ValueError( @@ -154,7 +151,7 @@ def list(cls, task_type: str | None = None) -> list[str]: "perspectives, and symbol rendering from basic to complex compositions." ), metrics=[], # Paper uses human evaluation only; pass explicit metrics if needed - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2206.10789", ), Benchmark( @@ -164,7 +161,7 @@ def list(cls, task_type: str | None = None) -> list[str]: "Enables side-by-side comparison on sample quality and image-text alignment with human raters." ), metrics=[], # Paper uses human evaluation only; pass explicit metrics if needed - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2205.11487", ), Benchmark( @@ -174,8 +171,8 @@ def list(cls, task_type: str | None = None) -> list[str]: "Covers basic skills (scene, attributes, spatial relationships) to advanced reasoning " "(counting, comparison, logic/negation) with over 24k human ratings." ), - metrics=[], # Paper uses VQAScore only; not in Pruna - task_type="text_to_image", + metrics=["vqa", "clip_score"], # VQAScore + CLIPScore both named (arXiv:2406.13743) + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2406.13743", ), Benchmark( @@ -195,8 +192,8 @@ def list(cls, task_type: str | None = None) -> list[str]: "MS-COCO for text-to-image evaluation (Imagen, 2205.11487). Paper reports " "FID for fidelity and CLIP score for image-text alignment." ), - metrics=["fid", "clip_score"], # ?4.1: FID + CLIP score - task_type="text_to_image", + metrics=["fid", "clip_score"], + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2205.11487", ), Benchmark( @@ -223,11 +220,13 @@ def list(cls, task_type: str | None = None) -> list[str]: name="GenEval", description=( "Compositional text-to-image benchmark with 6 categories: single object, two object, " - "counting, colors, position, color attributes. Evaluates fine-grained alignment " - "between prompts and generated images via VQA-style questions." + "counting, colors, position, color attributes. Uses atomic yes/no questions per prompt; " + "``Task.from_benchmark`` wires ``qa_accuracy`` with strict per-image aggregation " + "(all questions must pass) plus ``clip_score``. For holistic VQAScore-style scoring " + "use GenAI Bench with ``vqa``." ), - metrics=["qa_accuracy", "clip_score"], # strict QA + CLIP score - task_type="text_to_image", + metrics=["qa_accuracy", "clip_score"], + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2310.11513", ), Benchmark( @@ -237,7 +236,7 @@ def list(cls, task_type: str | None = None) -> list[str]: "Covers anime, concept-art, paintings, and photo styles with human preference data." ), metrics=[], # Paper uses HPS scoring model; not in Pruna - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2306.09341", ), Benchmark( @@ -246,18 +245,19 @@ def list(cls, task_type: str | None = None) -> list[str]: "Image editing benchmark with 8 edit types: replace, add, remove, adjust, extract, " "style, background, compose. Evaluates instruction-following for inpainting and editing." ), - metrics=[], # Paper uses GPT-4o/ImgEdit-Judge; not in Pruna - task_type="text_to_image", + metrics=["img_edit_score"], # Paper uses GPT-4o rubric scores + FakeShield judge. + task_type=TASK_TYPE_TEXT_PLUS_IMAGE_IMAGE, reference="https://arxiv.org/abs/2505.20275", ), Benchmark( name="Long Text Bench", description=( - "Text-to-image benchmark for long, detailed prompts. Evaluates model ability to " - "handle complex multi-clause descriptions and maintain coherence across long instructions." + "Text rendering benchmark evaluating whether T2I models correctly render specific text strings " + "provided in prompts. Uses ``text_score`` (normalized character accuracy in [0, 1]). " + "This is OCR correctness, not long-prompt semantic alignment." ), metrics=[], # Paper uses word accuracy (X-Omni); not wired to text_score yet - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2507.22058", ), Benchmark( @@ -265,45 +265,53 @@ def list(cls, task_type: str | None = None) -> list[str]: description=( "General image editing benchmark with 11 task types: background change, color alter, " "material alter, motion change, style change, subject add/remove/replace, text change, " - "tone transfer, and human retouching." + "tone transfer, and human retouching. " + "Evaluated with VIEScore in text-image-edit (TIE) mode when source image bytes are available." ), - metrics=[], # Paper uses VIEScore; not in Pruna - task_type="text_to_image", + metrics=["vie_score"], # VIEScore is explicitly named in GEditBench. + task_type=TASK_TYPE_TEXT_PLUS_IMAGE_IMAGE, reference="https://arxiv.org/abs/2504.17761", ), Benchmark( name="OneIG Anime Stylization", description="OneIG subset: anime and stylized imagery.", metrics=["oneig_alignment"], - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2506.07977", ), Benchmark( name="OneIG General Object", description="OneIG subset: everyday objects and scenes.", metrics=["oneig_alignment"], - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, + reference="https://arxiv.org/abs/2506.07977", + ), + Benchmark( + name="OneIG Knowledge Reasoning", + description="OneIG subset: knowledge- and reasoning-heavy prompts.", + metrics=["oneig_reasoning"], + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2506.07977", ), Benchmark( name="OneIG Multilingualism", description="OneIG subset: multilingual prompts (incl. Chinese splits).", metrics=["oneig_alignment"], - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2506.07977", ), Benchmark( name="OneIG Portrait", description="OneIG subset: people and portraits.", metrics=["oneig_alignment"], - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2506.07977", ), Benchmark( name="OneIG Text Rendering", description="OneIG subset: text and graphics painted into the image.", metrics=["oneig_text_score"], - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2506.07977", ), Benchmark( @@ -313,7 +321,7 @@ def list(cls, task_type: str | None = None) -> list[str]: "global, and other descriptive aspects with natural-language questions for alignment." ), metrics=[], # Paper uses custom evaluation; not in Pruna - task_type="text_to_image", + task_type=TASK_TYPE_TEXT_IMAGE, reference="https://arxiv.org/abs/2403.05135", ), ]: diff --git a/src/pruna/evaluation/metrics/__init__.py b/src/pruna/evaluation/metrics/__init__.py index 4f18b9f8..989ad548 100644 --- a/src/pruna/evaluation/metrics/__init__.py +++ b/src/pruna/evaluation/metrics/__init__.py @@ -20,6 +20,7 @@ from pruna.evaluation.metrics.metric_elapsed_time import LatencyMetric, ThroughputMetric, TotalTimeMetric from pruna.evaluation.metrics.metric_energy import CO2EmissionsMetric, EnergyConsumedMetric from pruna.evaluation.metrics.metric_evalharness import LMEvalMetric +from pruna.evaluation.metrics.metric_img_edit_score import ImageEditScoreMetric from pruna.evaluation.metrics.metric_memory import DiskMemoryMetric, InferenceMemoryMetric, TrainingMemoryMetric from pruna.evaluation.metrics.metric_model_architecture import TotalMACsMetric, TotalParamsMetric from pruna.evaluation.metrics.metric_oneig_alignment import OneIGAlignmentMetric @@ -66,6 +67,7 @@ "RapidataMetric", "TextScoreMetric", "VQAMetric", + "ImageEditScoreMetric", "VieScoreMetric", "BaseVLM", "LitellmVLM", diff --git a/src/pruna/evaluation/metrics/metric_img_edit_score.py b/src/pruna/evaluation/metrics/metric_img_edit_score.py new file mode 100644 index 00000000..4c1832af --- /dev/null +++ b/src/pruna/evaluation/metrics/metric_img_edit_score.py @@ -0,0 +1,219 @@ +# Copyright 2025 - Pruna AI GmbH. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Image Edit Score metric. + +VLM-based instruction-following score for image editing. Evaluates how well an edited image +follows the given editing instruction on a 0-10 scale. Related work: EditScore (arXiv:2509.23909), +ADIEE (ICCV 2025). + +When the ``ImgEdit`` benchmark provides a per-sample ``judge_prompt`` and +``source_image_bytes`` in the auxiliaries, the metric mirrors the ImgEdit paper +evaluation protocol: the judge_prompt rubric (three 1-5 criterion scores) is +filled with the editing instruction, both source and edited images are shown to +the VLM, and the minimum of the three criterion scores is normalised to [0, 1] by +dividing by 5 (consistent with VIEScore methodology: the weakest criterion governs). +Without these auxiliaries the metric falls back to a single-image generic 0-10 prompt. +""" + +from __future__ import annotations + +from typing import Any, Literal + +import torch + +from pruna.evaluation.metrics.registry import MetricRegistry +from pruna.evaluation.metrics.result import MetricResult +from pruna.evaluation.metrics.utils import ( + SINGLE, + metric_data_processor, +) +from pruna.evaluation.metrics.vlm_base import ( + BaseVLM, + StatefulVLMMeanScoresMetric, + auxiliary_dicts_from_gt, + prompts_from_y_x_inputs, +) +from pruna.evaluation.metrics.vlm_utils import ( + FloatOutput, + VIEScoreJsonOutput, + _process_images, + get_score_from_response, + pil_rgb_from_aux_image_bytes, + viescore_min_scores_0_10, +) + +_FALLBACK_QUESTION = ( + 'On a scale of 0 to 10, how well does this edited image follow the instruction "{prompt}"? ' + "0 = instruction not followed at all, 10 = perfectly executed. Reply with a single number." +) + +_JUDGE_JSON_SUFFIX = ( + '\n\nProvide your three criterion scores as JSON: {"score": [score1, score2, score3]} ' + "where each score is a number from 1 to 5." +) + + +@MetricRegistry.register("img_edit_score") +class ImageEditScoreMetric(StatefulVLMMeanScoresMetric): + """ + Image Edit Score metric. + + VLM-based instruction-following score for image editing. Evaluates how well an edited image + follows the given editing instruction. Higher scores indicate better editing quality. + + When auxiliaries contain ``judge_prompt`` and ``source_image_bytes`` (as provided + by the ImgEdit benchmark), the metric passes **both** the source (before) and edited + (after) images to the VLM together with the dataset-specific rubric. This matches + the ImgEdit paper's evaluation protocol. Without these fields, it falls back to a + single-image generic question. + + Related work: EditScore (arXiv:2509.23909), ADIEE (ICCV 2025). + + Parameters + ---------- + *args : Any + Additional positional arguments. + vlm : BaseVLM | None, optional + Custom VLM instance. If provided, vlm_type and model_name are ignored. + vlm_type : {"litellm", "transformers"}, optional + VLM backend. Default is "litellm". + model_name : str | None, optional + Litellm model id or HuggingFace checkpoint id. **Required** when ``vlm`` is not + provided (e.g. ``openai/gpt-4o``). + vlm_kwargs : dict, optional + Forwarded by ``get_vlm`` to ``LitellmVLM`` or ``TransformersVLM``. For local models, + set ``model_load_kwargs`` for ``from_pretrained``; for litellm, pass extra API options. + structured_output : bool, optional + Use structured generation (litellm pydantic; transformers outlines when applicable). + Default is True. + device : str | torch.device | None, optional + Device for transformers VLM. + api_key : str | None, optional + API key for litellm. + call_type : str, optional + Call type for the metric. + **kwargs : Any + Additional arguments. + + Examples + -------- + Same ``hosted`` / ``local`` pattern as :func:`~pruna.evaluation.metrics.vlm_base.get_vlm`: + + .. code-block:: python + + import torch + + from pruna.evaluation.metrics import ImageEditScoreMetric + + hosted = ImageEditScoreMetric(vlm_type="litellm", model_name="openai/gpt-4o") + local = ImageEditScoreMetric( + vlm_type="transformers", + model_name="HuggingFaceTB/SmolVLM-256M-Instruct", + device="cpu", + vlm_kwargs={"model_load_kwargs": {"torch_dtype": torch.float32}}, + ) + """ + + scores: list[float] + default_call_type: str = "y_x" + higher_is_better: bool = True + metric_name: str = "img_edit_score" + + def __init__( + self, + *args, + vlm: BaseVLM | None = None, + vlm_type: Literal["litellm", "transformers"] = "litellm", + model_name: str | None = None, + vlm_kwargs: dict | None = None, + structured_output: bool = True, + device: str | torch.device | None = None, + api_key: str | None = None, + call_type: str = SINGLE, + **kwargs: Any, + ) -> None: + super().__init__(device=device) + self.response_format = FloatOutput if structured_output else None + + self._init_vlm_scores( + vlm=vlm, + vlm_type=vlm_type, + model_name=model_name, + vlm_kwargs=vlm_kwargs, + structured_output=structured_output, + device=device, + api_key=api_key, + call_type=call_type, + ) + + def update(self, x: list[Any] | torch.Tensor, gt: torch.Tensor, outputs: torch.Tensor) -> None: + """ + Update the metric with new batch data. + + When ``gt`` auxiliaries contain ``judge_prompt`` and ``source_image_bytes``, the + metric uses the dataset rubric and a before/after two-image comparison. Otherwise + it falls back to a single-image generic question. + + Parameters + ---------- + x : List[Any] | torch.Tensor + The input data (editing instructions / prompts). + gt : torch.Tensor + Auxiliaries per sample (may contain ``judge_prompt`` and ``source_image_bytes``). + outputs : torch.Tensor + The output (edited) images. + """ + inputs = metric_data_processor(x, gt, outputs, self.call_type) + images = _process_images(inputs[0]) + prompts = prompts_from_y_x_inputs(inputs, len(images)) + aux_list = auxiliary_dicts_from_gt(gt, len(images)) + + for i, image in enumerate(images): + prompt = prompts[i] if i < len(prompts) else "" + aux_row = aux_list[i] + + judge_prompt = aux_row.get("judge_prompt", "") or "" + source_image = pil_rgb_from_aux_image_bytes(aux_row, min_bytes_in_value_scan=100) + + if judge_prompt and source_image is not None: + filled = judge_prompt.replace("", prompt).strip() + question = filled + _JUDGE_JSON_SUFFIX + try: + responses = self.vlm.generate_with_image_lists( + [[source_image, image]], [question], response_format=VIEScoreJsonOutput + ) + raw = viescore_min_scores_0_10(responses[0]) + if raw: + score = max(0.0, min(1.0, float(min(raw)) / 5.0)) + self.scores.append(score) + continue + except (NotImplementedError, AttributeError): + pass + + question = _FALLBACK_QUESTION.format(prompt=prompt) + responses = self.vlm.generate([image], [question], response_format=self.response_format) + self.scores.append(get_score_from_response(responses[0])) + + def compute(self) -> MetricResult: + """ + Compute the image edit score. + + Returns + ------- + MetricResult + The mean image edit score across all updates. + """ + return self.compute_mean_of_scores() diff --git a/tests/evaluation/test_vision_metrics.py b/tests/evaluation/test_vision_metrics.py index e7284eee..f41df366 100644 --- a/tests/evaluation/test_vision_metrics.py +++ b/tests/evaluation/test_vision_metrics.py @@ -5,6 +5,7 @@ import pytest import torch +from pruna.evaluation.metrics.metric_img_edit_score import ImageEditScoreMetric from pruna.evaluation.metrics.metric_vie_score import VieScoreMetric from pruna.evaluation.metrics.metric_vqa import VQAMetric from pruna.evaluation.metrics.vlm_base import BaseVLM @@ -42,3 +43,16 @@ def test_vie_score_uses_json_score_lists() -> None: result = metric.compute() assert abs(result.result - 0.8) < 0.01 + + +@pytest.mark.cpu +def test_img_edit_score_negative_response_clamped() -> None: + """Image edit score never goes below zero for malformed negative outputs.""" + mock_vlm = MagicMock(spec=BaseVLM) + mock_vlm.generate.return_value = ['{"score": -10.0}'] + + metric = ImageEditScoreMetric(vlm=mock_vlm, device="cpu", structured_output=True) + metric.update(["replace the boot with a mug"], torch.zeros(1), _dummy_image()) + result = metric.compute() + + assert result.result == 0.0