run_inference.py

from __future__ import annotations
import argparse
import json
import re
from pathlib import Path
from typing import Any

import torch
from PIL import Image
from transformers import AutoProcessor, Qwen3VLForConditionalGeneration


HF_REPO_ID = "vcl-iisc/DynEval-Evaluator"
HF_SUBFOLDERS = {
    "2b": "DynEval-2B",
    "4b": "DynEval-4B",
}

T2IA_ELEMENT_SYSTEM = """Given an aigc prompt, extract the elements that are important for generating images.
Classify each element into a type (object, human, animal, food, activity, attribute, counting, color, material, spatial, location, shape, other).
Do NOT extract medium or framing words that only describe how the image is captured or rendered, such as "photo", "photograph", "image", "picture", "shot", "render", "3D render", "illustration", "drawing", or "wallpaper". Ignore leading phrases like "a photo of", "a picture of", or "an image of" and only extract the actual visual content that follows.
Examples are as follows, where Elements is in json format.

Prompt: A man posing for a selfie in a jacket and bow tie.
Elements: ["man (human)", "selfie (activity)", "jacket (object)", "bow tie (object)", "posing (activity)"]

Prompt: A horse and several cows feed on hay.
Elements: ["horse (animal)", "cows (animal)", "hay (object)", "feed on (activity)", "several (counting)"]

Prompt: A photo of a bench.
Elements: ["bench (object)"]

Prompt: An illustration of a red apple on a table.
Elements: ["apple (object)", "red (color)", "table (object)", "on (spatial)"]"""

T2IA_SINGLE_QUESTION_SYSTEM = """Given a prompt for image generation and one of its related elements, generate one easy Yes/No question to verify whether the element is represented in the image generated by the prompt.
The question must be answerable by inspecting the image.
Return only JSON with exactly these keys: "question" and "answer".
The answer should be "yes" unless the element explicitly describes an absence or zero count."""

EVALUATION_SYSTEM = """You are a strict visual evidence evaluator. Your job is to test for image-text alignment by looking at the image and answering the questions respectively.
CRITICAL RULES:
1. Score based ONLY on what you can SEE in the image, not what the prompt says should be there.
2. Do NOT assume elements exist just because the prompt mentions them. Look at the actual pixels.
SCORING SCALE:
  5 = YES, DEFINITELY | 4 = YES, PROBABLY | 3 = UNCERTAIN | 2 = NO, PROBABLY | 1 = NO, DEFINITELY
Return ONLY a JSON array."""

IQA_SCENE_GRAPH_PROMPT = """Here is an image generated for this prompt "{prompt}".
Now generate a scene graph by considering only the image. Consider the text prompt for your reference.
Return only valid JSON with exactly these keys: "nodes" and "edges".
"nodes" must be an array of visible objects, where each object has "id", "label", and "attributes".
"edges" must be an array of relationships, where each relationship has "source", "relation", and "target".
Do not write markdown, bullets, or prose."""

IQA_DECOMPOSITION_PROMPT = """We are doing an image quality assessment. The generated image is provided to you along with a list of object nodes JSON.
We have ALREADY done the text-to-image alignment. Therefore you must NOT ask presence, attribute, color, count, or alignment questions (for example: "Is there a bench?", "Is the bench wooden?", "Are the flowers red?", "Is the wall white?"). Such questions are forbidden.
Instead, for each object node, decompose it into its parts and generate yes/no questions that judge ONLY the rendering QUALITY of how that object is drawn in the image. Focus on perceptual defects such as: shape correctness and geometry, structural distortions or deformities, unnatural or broken proportions, blur or smearing, texture artifacts, melting/warping, incorrect or impossible 3D spatial structure, and boundary/edge coherence.
Each question must be phrased so that a well-rendered, high-quality object gets the target answer "yes".
Examples of the required style: "Are the bench slats straight and free of distortion?", "Are the flower petals rendered with clean edges and no smearing?", "Does the wall surface have a consistent texture without artifacts?".
Inspect the provided image to decide the target answer.
Return only valid JSON with exactly one key: "questions".
Each question object must contain "node_id", "question", and "target_answer".
Do not write markdown, bullets, or prose."""

IQA_FINAL_PROMPT = """Now answer these questions for this image and generate a score given the target answers on a scale of 1-5.
Use the provided question JSON directly.
Do not write paragraphs, markdown, summaries, or bullet-point analysis.
Return only valid JSON in this exact format:
{
  "questions": [
    {"question": "...", "target_answer": "yes/no", "answer": "yes/no", "score": 1, "reasoning": "..."}
  ]
}
Each "score" must be a number from 1 to 5 based on how well the image answer matches the target answer."""

def json_dumps(value: Any) -> str:
    # Serialize Python objects as readable UTF-8 JSON.
    return json.dumps(value, ensure_ascii=False, indent=2)


def strip_code_fence(text: str) -> str:
    # Remove Markdown code fences that models sometimes wrap around JSON.
    text = (text or "").strip()
    if text.startswith("```"):
        lines = text.splitlines()
        if lines and lines[0].startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip() == "```":
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    return text


def parse_first_json(raw: str) -> Any:
    # Extract the first valid JSON object or array from a raw model response.
    text = strip_code_fence(raw)
    if not text:
        return None
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    decoder = json.JSONDecoder()
    for opener in ("[", "{"):
        start = text.find(opener)
        while start >= 0:
            try:
                parsed, _ = decoder.raw_decode(text[start:])
                return parsed
            except json.JSONDecodeError:
                start = text.find(opener, start + 1)
    return None


def normalize_key(text: str) -> str:
    # Normalize text for case-insensitive duplicate detection.
    return re.sub(r"\s+", " ", str(text).strip().lower())


def dedupe_strings(items: list[str]) -> list[str]:
    # Keep string items in original order while removing normalized duplicates.
    out = []
    seen = set()
    for item in items:
        item = str(item).strip()
        key = normalize_key(item)
        if key and key not in seen:
            seen.add(key)
            out.append(item)
    return out


def dedupe_questions(items: list[dict[str, Any]]) -> list[dict[str, str]]:
    # Remove duplicate element/question pairs and normalize yes/no answers.
    out = []
    seen = set()
    for item in items:
        if not isinstance(item, dict):
            continue
        question = str(item.get("question", "")).strip()
        element = str(item.get("element", "")).strip()
        answer = str(item.get("answer", "yes")).strip().lower()
        key = (normalize_key(element), normalize_key(question))
        if question and key not in seen:
            seen.add(key)
            out.append({"element": element, "question": question, "answer": "no" if answer.startswith("no") else "yes"})
    return out


MEDIUM_ELEMENT_WORDS = {
    "photo", "photograph", "photography", "image", "picture", "pic", "shot",
    "snapshot", "render", "rendering", "3d render", "illustration", "drawing",
    "sketch", "painting", "artwork", "wallpaper", "poster", "portrait",
}


def drop_medium_elements(items: list[str]) -> list[str]:
    # Remove medium/framing words (e.g. "photo (attribute)") that are not visual
    # content to verify. This is a hard filter because the checkpoint was trained
    # to emit such words and an in-context instruction alone does not stop it.
    out = []
    for item in items:
        base = re.sub(r"\(.*?\)", "", str(item)).strip().lower()
        base = base.strip(" .'\"")
        if base in MEDIUM_ELEMENT_WORDS:
            continue
        out.append(item)
    return out


def parse_elements(raw: str) -> list[str]:
    # Parse the <T2IA> element-extraction response into a clean element list.
    parsed = parse_first_json(raw)
    if isinstance(parsed, list):
        return drop_medium_elements(dedupe_strings([str(item) for item in parsed if str(item).strip()]))
    if isinstance(parsed, dict):
        for key in ("elements", "Elements"):
            if isinstance(parsed.get(key), list):
                return drop_medium_elements(dedupe_strings([str(item) for item in parsed[key] if str(item).strip()]))
    return drop_medium_elements(dedupe_strings(re.findall(r'"([^"]+\([^()]+\))"', raw or "")))


def parse_question(raw: str, element: str) -> dict[str, str] | None:
    # Parse one generated question for a specific element.
    parsed = parse_first_json(raw)
    if isinstance(parsed, dict):
        question = str(parsed.get("question", "")).strip()
        answer = str(parsed.get("answer", "yes")).strip().lower()
        if question:
            return {"element": element, "question": question, "answer": "no" if answer.startswith("no") else "yes"}
    match = re.search(r"(?:Q:|Question:)\s*(.+?)(?:\n|$)", raw or "", re.IGNORECASE)
    if match:
        return {"element": element, "question": match.group(1).strip(), "answer": "yes"}
    return None


def normalize_score(value: Any) -> int:
    # Convert a model score to an integer in [1, 5], or 0 if invalid.
    try:
        score = int(round(float(value)))
    except (TypeError, ValueError):
        match = re.search(r"\b([1-5])\b", str(value))
        score = int(match.group(1)) if match else 0
    return score if 1 <= score <= 5 else 0


def normalize_iqa_score(value: Any) -> float:
    # Convert an IQA score to a number in [1, 5], or 0 if invalid.
    try:
        score = float(value)
    except (TypeError, ValueError):
        match = re.search(r"\b([1-5](?:\.\d+)?)\b", str(value))
        score = float(match.group(1)) if match else 0.0
    return score if 1.0 <= score <= 5.0 else 0.0


def parse_answers(raw: str, questions: list[dict[str, Any]]) -> list[dict[str, Any]]:
    # Parse <EVALUATION> scores and align them with the generated questions.
    parsed = parse_first_json(raw)
    if isinstance(parsed, dict):
        parsed = parsed.get("answers", parsed.get("scores", parsed))

    out = []
    if isinstance(parsed, list):
        for idx, question in enumerate(questions):
            item = parsed[idx] if idx < len(parsed) else {}
            if isinstance(item, dict):
                score = normalize_score(item.get("score", item.get("answer", 0)))
                reasoning = str(item.get("reasoning", "")).strip()
            else:
                score = normalize_score(item)
                reasoning = ""
            out.append({"question": question["question"], "reasoning": reasoning, "score": score})
        return out

    fallback_scores = [int(match) for match in re.findall(r"\b([1-5])\b", raw or "")]
    for idx, question in enumerate(questions):
        out.append({"question": question["question"], "reasoning": "", "score": fallback_scores[idx] if idx < len(fallback_scores) else 0})
    return out


def extract_json_array_after_key(text: str, key: str) -> list[Any] | None:
    """Return the JSON array following `"key":`, matched by balanced brackets.

    Recovers arrays even when the model wraps the object in the wrong bracket
    (e.g. emits `["nodes": [...], "edges": [...]]` instead of `{...}`) or appends
    a degenerate tail of stray `]` characters after the valid JSON.
    """
    match = re.search(rf'"{re.escape(key)}"\s*:\s*\[', text)
    if not match:
        return None
    start = match.end() - 1
    depth = 0
    in_string = False
    escape = False
    for idx in range(start, len(text)):
        char = text[idx]
        if in_string:
            if escape:
                escape = False
            elif char == "\\":
                escape = True
            elif char == '"':
                in_string = False
            continue
        if char == '"':
            in_string = True
        elif char == "[":
            depth += 1
        elif char == "]":
            depth -= 1
            if depth == 0:
                try:
                    value = json.loads(text[start : idx + 1])
                except json.JSONDecodeError:
                    return None
                return value if isinstance(value, list) else None
    return None


def extract_json_objects(raw: str) -> list[dict[str, Any]]:
    # Recover every complete top-level JSON object from text, even when the model
    # truncates the array before its closing bracket (e.g. it hit max_new_tokens).
    text = strip_code_fence(raw)
    decoder = json.JSONDecoder()
    objects: list[dict[str, Any]] = []
    idx = 0
    length = len(text)
    while idx < length:
        if text[idx] == "{":
            try:
                obj, end = decoder.raw_decode(text, idx)
            except json.JSONDecodeError:
                idx += 1
                continue
            if isinstance(obj, dict):
                objects.append(obj)
            idx = end
            continue
        idx += 1
    return objects


def parse_scene_graph(raw: str) -> dict[str, list[dict[str, Any]]]:
    # Parse the IQA scene graph into explicit node and edge lists.
    parsed = parse_first_json(raw)
    nodes: list[dict[str, Any]] = []
    edges: list[dict[str, Any]] = []

    if not isinstance(parsed, dict):
        text = strip_code_fence(raw)
        raw_nodes = extract_json_array_after_key(text, "nodes") or extract_json_array_after_key(text, "Nodes")
        raw_edges = extract_json_array_after_key(text, "edges") or extract_json_array_after_key(text, "Edges")
        if raw_nodes is not None or raw_edges is not None:
            parsed = {"nodes": raw_nodes or [], "edges": raw_edges or []}

    if isinstance(parsed, dict):
        raw_nodes = parsed.get("nodes", parsed.get("Nodes", []))
        raw_edges = parsed.get("edges", parsed.get("Edges", []))
        if isinstance(raw_nodes, list):
            for idx, item in enumerate(raw_nodes, start=1):
                if isinstance(item, dict):
                    node_id = str(item.get("id", item.get("node_id", f"node_{idx}"))).strip() or f"node_{idx}"
                    label = str(item.get("label", item.get("name", item.get("object", "")))).strip()
                    attributes = item.get("attributes", [])
                else:
                    node_id = f"node_{idx}"
                    label = str(item).strip()
                    attributes = []
                if label:
                    nodes.append({"id": node_id, "label": label, "attributes": attributes if isinstance(attributes, list) else [str(attributes)]})
        if isinstance(raw_edges, list):
            for item in raw_edges:
                if not isinstance(item, dict):
                    continue
                source = str(item.get("source", "")).strip()
                relation = str(item.get("relation", item.get("predicate", ""))).strip()
                target = str(item.get("target", "")).strip()
                if source and relation and target:
                    edges.append({"source": source, "relation": relation, "target": target})

    return {"nodes": nodes, "edges": edges}


def parse_iqa_questions(raw: str) -> list[dict[str, str]]:
    # Parse IQA question-generation output into questions with target answers.
    parsed = parse_first_json(raw)
    if isinstance(parsed, dict):
        parsed = parsed.get("questions", parsed.get("Questions", []))

    questions: list[dict[str, str]] = []
    seen = set()
    if isinstance(parsed, list):
        for item in parsed:
            if isinstance(item, dict):
                question = str(item.get("question", item.get("Question", ""))).strip()
                node_id = str(item.get("node_id", item.get("node", ""))).strip()
                target_answer = str(item.get("target_answer", item.get("target", "yes"))).strip().lower()
            else:
                question = str(item).strip()
                node_id = ""
                target_answer = "yes"
            key = normalize_key(question)
            if question and key not in seen:
                seen.add(key)
                questions.append({"node_id": node_id, "question": question, "target_answer": "no" if target_answer.startswith("no") else "yes"})

    if not questions:
        for match in re.findall(r"(?:^|\n)\s*(?:[-*]|\d+[.)])\s*(.+?\?)", raw or ""):
            question = match.strip()
            key = normalize_key(question)
            if key not in seen:
                seen.add(key)
                questions.append({"node_id": "", "question": question, "target_answer": "yes"})
    return questions


def parse_iqa_result(raw: str) -> dict[str, Any]:
    # Parse IQA questions with target answers, model answers, and 1-5 scores.
    parsed = parse_first_json(raw)
    questions: list[dict[str, Any]] = []
    seen = set()

    raw_items: list[Any] = []
    if isinstance(parsed, dict):
        maybe = parsed.get("questions", parsed.get("Questions"))
        raw_items = maybe if isinstance(maybe, list) else []
    elif isinstance(parsed, list):
        raw_items = parsed

    # Recover complete objects when the model truncates the JSON array (e.g. it
    # ran out of max_new_tokens and never closed the final object/bracket).
    if not raw_items:
        raw_items = extract_json_objects(raw)

    for item in raw_items:
        if isinstance(item, dict):
            question = str(item.get("question", item.get("Question", ""))).strip()
            score = normalize_iqa_score(item.get("score", item.get("Score", 0)))
            target_answer = str(item.get("target_answer", item.get("target", ""))).strip().lower()
            answer = str(item.get("answer", item.get("Answer", ""))).strip().lower()
            reasoning = str(item.get("reasoning", item.get("Reasoning", ""))).strip()
        else:
            question = str(item).strip()
            score = 0.0
            target_answer = ""
            answer = ""
            reasoning = ""
        if question:
            key = normalize_key(question)
            if key not in seen:
                seen.add(key)
                questions.append(
                    {
                        "question": question,
                        "target_answer": target_answer,
                        "answer": answer,
                        "score": score,
                        "reasoning": reasoning,
                    }
                )

    if not questions:
        for match in re.findall(r"(?:^|\n)\s*(?:[-*]|\d+[.)])\s*(.+?\?)", raw or ""):
            question = match.strip()
            key = normalize_key(question)
            if key not in seen:
                seen.add(key)
                questions.append({"question": question, "target_answer": "", "answer": "", "score": 0.0, "reasoning": ""})

    valid_scores = [item["score"] for item in questions if normalize_iqa_score(item.get("score")) > 0]
    score = sum(valid_scores) / len(valid_scores) if valid_scores else 0.0
    return {"questions": questions, "score": score}


def resolve_task_token(tokenizer, base: str) -> str:
    """Support both local-training tokens and HF-card pipe tokens.

    Local experimental checkpoints used <T2IA>/<IQA>/<EVALUATION>.
    The public HF card documents <|T2IA|>/<|IQA|>/<|EVALUATION|>.
    We choose whichever is registered as a single tokenizer token.
    """
    candidates = [f"<|{base}|>", f"<{base}>"]
    for token in candidates:
        token_id = tokenizer.convert_tokens_to_ids(token)
        encoded = tokenizer.encode(token, add_special_tokens=False)
        if len(encoded) == 1 and encoded[0] == token_id:
            return token
    return candidates[0]


def build_element_messages(token: str, prompt: str) -> list[dict[str, Any]]:
    # Build the chat messages for extracting visual elements from a prompt.
    return [
        {"role": "system", "content": [{"type": "text", "text": T2IA_ELEMENT_SYSTEM}]},
        {"role": "user", "content": [{"type": "text", "text": f"{token}\nPrompt: {prompt}\nElements:"}]},
    ]


def build_question_messages(token: str, prompt: str, element: str) -> list[dict[str, Any]]:
    # Build the chat messages for generating one yes/no question per element.
    user_text = f'{token}\nDescription: {prompt}\nElement: {element}\nReturn JSON:\n{{"question": "...", "answer": "yes/no"}}'
    return [
        {"role": "system", "content": [{"type": "text", "text": T2IA_SINGLE_QUESTION_SYSTEM}]},
        {"role": "user", "content": [{"type": "text", "text": user_text}]},
    ]


def build_evaluation_messages(token: str, questions: list[dict[str, Any]]) -> list[dict[str, Any]]:
    # Build multimodal chat messages for scoring generated questions on an image.
    questions_text = "\n".join(f'{idx + 1}. Question: {item["question"]}' for idx, item in enumerate(questions))
    user_text = f"{token}\nQuestions to score:\n{questions_text}"
    return [
        {"role": "system", "content": [{"type": "text", "text": EVALUATION_SYSTEM}]},
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": user_text}]},
    ]


def build_iqa_scene_graph_messages(token: str, prompt: str) -> list[dict[str, Any]]:
    # Start the in-context IQA conversation with image-conditioned scene graph generation.
    user_text = f"{token}\n{IQA_SCENE_GRAPH_PROMPT.format(prompt=prompt)}"
    return [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": user_text}]},
    ]


def build_iqa_question_messages(token: str, prompt: str, scene_graph: dict[str, Any]) -> list[dict[str, Any]]:
    # Build a fresh image chat that passes only the node ids and labels.
    # Attributes are deliberately omitted so the model does not just re-verify
    # them, which would produce alignment/attribute questions instead of quality ones.
    raw_nodes = scene_graph.get("nodes", []) if isinstance(scene_graph, dict) else []
    nodes = [
        {"node_id": str(node.get("id", "")), "label": str(node.get("label", ""))}
        for node in raw_nodes
        if isinstance(node, dict) and str(node.get("label", "")).strip()
    ]
    user_text = (
        f"{token}\nPrompt: {prompt}\n\n"
        f"Object nodes JSON:\n{json_dumps(nodes)}\n\n"
        f"{IQA_DECOMPOSITION_PROMPT}"
    )
    return [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": user_text}]},
    ]


def build_iqa_final_messages(token: str, prompt: str, scene_graph: dict[str, Any], questions: list[dict[str, str]]) -> list[dict[str, Any]]:
    # Build a fresh image chat that passes generated IQA questions for scoring.
    user_text = (
        f"{token}\nPrompt: {prompt}\n\n"
        f"Scene graph JSON:\n{json_dumps(scene_graph)}\n\n"
        f"IQA question JSON:\n{json_dumps({'questions': questions})}\n\n"
        f"{IQA_FINAL_PROMPT}"
    )
    return [
        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": user_text}]},
    ]


class DynEvalRunner:
    # Loads the DynEval checkpoint and runs the full T2IA-to-evaluation pipeline.

    def __init__(self, args: argparse.Namespace):
        # Load processor/model from a local checkpoint or the Hugging Face repo.
        load_kwargs: dict[str, Any] = {
            "torch_dtype": getattr(torch, args.dtype) if args.dtype != "auto" else "auto",
            "device_map": args.device_map,
            "trust_remote_code": True,
        }
        processor_kwargs: dict[str, Any] = {"trust_remote_code": True}

        if args.checkpoint:
            model_ref = args.checkpoint
            subfolder = None
        else:
            model_ref = args.repo_id
            subfolder = HF_SUBFOLDERS[args.variant]

        if subfolder:
            load_kwargs["subfolder"] = subfolder
            processor_kwargs["subfolder"] = subfolder

        self.processor = AutoProcessor.from_pretrained(model_ref, **processor_kwargs)
        if self.processor.tokenizer.pad_token_id is None:
            self.processor.tokenizer.pad_token = self.processor.tokenizer.eos_token
        self.processor.tokenizer.padding_side = "left"

        self.model = Qwen3VLForConditionalGeneration.from_pretrained(model_ref, **load_kwargs).eval()
        self.t2ia_token = resolve_task_token(self.processor.tokenizer, "T2IA")
        self.iqa_token = resolve_task_token(self.processor.tokenizer, "IQA")
        self.evaluation_token = resolve_task_token(self.processor.tokenizer, "EVALUATION")

    @property
    def device(self):
        # Return the device where the loaded model currently lives.
        return self.model.device

    def generate_text(self, messages: list[dict[str, Any]], max_new_tokens: int) -> str:
        # Run text-only generation for element extraction or question generation.
        inputs = self.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        )
        inputs = {key: value.to(self.device) if torch.is_tensor(value) else value for key, value in inputs.items()}
        with torch.inference_mode():
            generated = self.model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
        input_width = inputs["input_ids"].shape[1]
        return self.processor.decode(generated[0, input_width:], skip_special_tokens=True, clean_up_tokenization_spaces=False)

    def generate_image(self, messages: list[dict[str, Any]], image: Image.Image, max_new_tokens: int) -> str:
        # Run image-conditioned generation for the evaluation/scoring step.
        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self.processor(text=[text], images=[image], return_tensors="pt", padding=True)
        inputs = {key: value.to(self.device) if torch.is_tensor(value) else value for key, value in inputs.items()}
        with torch.inference_mode():
            generated = self.model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
        input_width = inputs["input_ids"].shape[1]
        return self.processor.decode(generated[0, input_width:], skip_special_tokens=True, clean_up_tokenization_spaces=False)

    def run_t2ia(self, prompt: str, image: Image.Image, args: argparse.Namespace) -> dict[str, Any]:
        # Run the T2IA alignment flow: elements, questions, then 1-5 visual scores.
        raw_elements = self.generate_text(build_element_messages(self.t2ia_token, prompt), args.max_new_tokens_elements)
        elements = parse_elements(raw_elements)

        questions = []
        raw_questions = []
        for element in elements:
            raw = self.generate_text(build_question_messages(self.t2ia_token, prompt, element), args.max_new_tokens_questions)
            raw_questions.append({"element": element, "raw": raw})
            question = parse_question(raw, element)
            if question:
                questions.append(question)
        questions = dedupe_questions(questions)

        raw_answers = self.generate_image(build_evaluation_messages(self.evaluation_token, questions), image, args.max_new_tokens_answers)
        answers = parse_answers(raw_answers, questions)

        return {
            "elements": elements,
            "questions": questions,
            "answers": answers,
            "raw_elements": raw_elements,
            "raw_questions": raw_questions,
            "raw_answers": raw_answers,
        }

    def run_iqa(self, prompt: str, image: Image.Image, args: argparse.Namespace) -> dict[str, Any]:
        # Run the IQA flow: scene graph, object quality analysis, scored questions.
        raw_scene_graph = self.generate_image(
            build_iqa_scene_graph_messages(self.iqa_token, prompt),
            image,
            args.max_new_tokens_iqa_scene_graph,
        )
        scene_graph = parse_scene_graph(raw_scene_graph)
        raw_iqa_questions = self.generate_image(
            build_iqa_question_messages(self.iqa_token, prompt, scene_graph),
            image,
            args.max_new_tokens_iqa_questions,
        )
        iqa_questions = parse_iqa_questions(raw_iqa_questions)
        raw_final = self.generate_image(
            build_iqa_final_messages(self.iqa_token, prompt, scene_graph, iqa_questions),
            image,
            args.max_new_tokens_iqa_final,
        )
        parsed_final = parse_iqa_result(raw_final)

        return {
            "scene_graph": scene_graph,
            "iqa_questions": iqa_questions,
            "questions": parsed_final["questions"],
            "score": parsed_final["score"],
            "raw_scene_graph": raw_scene_graph,
            "raw_iqa_questions": raw_iqa_questions,
            "raw_final": raw_final,
        }

    def run(self, prompt: str, image_path: Path, args: argparse.Namespace) -> dict[str, Any]:
        # Run the requested score pipeline for one prompt/image pair.
        image = Image.open(image_path).convert("RGB")

        result = {
            "prompt": prompt,
            "image_path": str(image_path),
        }
        t2ia = None
        iqa = None
        if args.score_type in ("t2ia", "both"):
            t2ia = self.run_t2ia(prompt, image, args)
            result["elements"] = t2ia["elements"]
            result["questions"] = t2ia["questions"]
            result["answers"] = t2ia["answers"]
        if args.score_type in ("iqa", "both"):
            iqa = self.run_iqa(prompt, image, args)
            result["iqa"] = {
                "scene_graph": iqa["scene_graph"],
                "iqa_questions": iqa["iqa_questions"],
                "questions": iqa["questions"],
                "score": iqa["score"],
            }
        if args.hide_elements:
            result.pop("elements", None)
        if args.include_raw:
            if t2ia is not None:
                result["raw_elements"] = t2ia["raw_elements"]
                result["raw_questions"] = t2ia["raw_questions"]
                result["raw_answers"] = t2ia["raw_answers"]
            if iqa is not None:
                result["raw_iqa_scene_graph"] = iqa["raw_scene_graph"]
                result["raw_iqa_questions"] = iqa["raw_iqa_questions"]
                result["raw_iqa_final"] = iqa["raw_final"]
        return result


def format_result(result: dict[str, Any]) -> str:
    # Format the pipeline result for readable terminal output.
    lines = ["=" * 88, "DynEval Evaluator Result", "=" * 88]
    lines.append(f"Prompt: {result['prompt']}")
    lines.append(f"Image: {result['image_path']}")

    elements = result.get("elements", [])
    if isinstance(elements, list):
        lines.extend(["", f"Elements ({len(elements)})"])
        for idx, element in enumerate(elements, start=1):
            lines.append(f"  {idx}. {element}")

    if "answers" in result:
        questions = result.get("questions", [])
        lines.extend(["", f"T2IA Questions ({len(questions)})"])
        for idx, item in enumerate(questions, start=1):
            lines.append(f"  {idx}. {item.get('question', '')}")
            if item.get("answer"):
                lines.append(f"     ground-truth answer: {item['answer']}")

        answers = result.get("answers", [])
        lines.extend(["", f"T2IA Evaluation Scores ({len(answers)})"])
        valid_scores = []
        for idx, item in enumerate(answers, start=1):
            score = normalize_score(item.get("score"))
            if 1 <= score <= 5:
                valid_scores.append(score)
            lines.append(f"  {idx}. score={score} | {item.get('question', '')}")
            if item.get("reasoning"):
                lines.append(f"     reasoning: {item['reasoning']}")
        if valid_scores:
            lines.append(f"  T2IA score: {sum(valid_scores) / len(valid_scores):.3f} / 5")

    iqa = result.get("iqa")
    if isinstance(iqa, dict):
        scene_graph = iqa.get("scene_graph", {})
        nodes = scene_graph.get("nodes", []) if isinstance(scene_graph, dict) else []
        edges = scene_graph.get("edges", []) if isinstance(scene_graph, dict) else []
        lines.extend(["", f"IQA Scene Graph Nodes ({len(nodes)})"])
        for idx, node in enumerate(nodes, start=1):
            if isinstance(node, dict):
                attrs = node.get("attributes", [])
                attr_text = f" | attributes: {attrs}" if attrs else ""
                lines.append(f"  {idx}. {node.get('id', '')}: {node.get('label', '')}{attr_text}")
        lines.extend(["", f"IQA Scene Graph Edges ({len(edges)})"])
        for idx, edge in enumerate(edges, start=1):
            if isinstance(edge, dict):
                lines.append(f"  {idx}. {edge.get('source', '')} --{edge.get('relation', '')}--> {edge.get('target', '')}")

        generated_iqa_questions = iqa.get("iqa_questions", [])
        lines.extend(["", f"Generated IQA Questions ({len(generated_iqa_questions) if isinstance(generated_iqa_questions, list) else 0})"])
        if isinstance(generated_iqa_questions, list):
            for idx, item in enumerate(generated_iqa_questions, start=1):
                lines.append(f"  {idx}. target={item.get('target_answer', '')} | {item.get('question', '')}")

        iqa_questions = iqa.get("questions", [])
        lines.extend(["", f"Scored IQA Questions ({len(iqa_questions) if isinstance(iqa_questions, list) else 0})"])
        if isinstance(iqa_questions, list):
            for idx, item in enumerate(iqa_questions, start=1):
                if isinstance(item, dict):
                    lines.append(f"  {idx}. score={normalize_iqa_score(item.get('score')):.3f} / 5 | {item.get('question', '')}")
                    if item.get("target_answer") or item.get("answer"):
                        lines.append(f"     target: {item.get('target_answer', '')} | answer: {item.get('answer', '')}")
                    if item.get("reasoning"):
                        lines.append(f"     reasoning: {item['reasoning']}")
                else:
                    lines.append(f"  {idx}. {item}")
        lines.extend(["", f"IQA score: {normalize_iqa_score(iqa.get('score')):.3f} / 5"])
    lines.append("=" * 88)
    return "\n".join(lines)


def parse_args() -> argparse.Namespace:
    # Parse command-line options for local or Hugging Face inference.
    parser = argparse.ArgumentParser(description="Run DynEval Evaluator inference from local or Hugging Face weights.")
    parser.add_argument("--checkpoint", type=str, default=None, help="Local checkpoint path. If set, this overrides --variant/--model-size and --repo-id.")
    parser.add_argument("--repo-id", type=str, default=HF_REPO_ID, help="Hugging Face repo id used when --checkpoint is not set.")
    parser.add_argument(
        "--variant",
        "--model-size",
        dest="variant",
        choices=["2b", "4b"],
        default="4b",
        help="HF model variant to load from the repo when --checkpoint is not set.",
    )
    parser.add_argument("--score-type", choices=["t2ia", "iqa", "both"], default="both", help="Which score to compute. Default computes both T2IA and IQA.")
    parser.add_argument("--prompt", type=str, required=True, help="Text-to-image prompt for the image.")
    parser.add_argument("--image", type=Path, required=True, help="Image path to evaluate.")
    parser.add_argument("--output-file", type=Path, default=None, help="Optional JSON output path.")
    parser.add_argument("--dtype", choices=["bfloat16", "float16", "float32", "auto"], default="bfloat16")
    parser.add_argument("--device-map", type=str, default="auto")
    parser.add_argument("--max-new-tokens-elements", type=int, default=256)
    parser.add_argument("--max-new-tokens-questions", type=int, default=256)
    parser.add_argument("--max-new-tokens-answers", type=int, default=768)
    parser.add_argument("--max-new-tokens-iqa-scene-graph", type=int, default=1024)
    parser.add_argument("--max-new-tokens-iqa-questions", "--max-new-tokens-iqa-decomposition", dest="max_new_tokens_iqa_questions", type=int, default=1024)
    parser.add_argument("--max-new-tokens-iqa-final", type=int, default=2048)
    parser.add_argument("--hide-elements", action="store_true", help="Hide intermediate elements from terminal and saved JSON.")
    parser.add_argument("--include-raw", action="store_true", help="Include raw model responses in saved JSON.")
    return parser.parse_args()


def main() -> None:
    # CLI entry point: validate inputs, run inference, print and optionally save JSON.
    args = parse_args()
    if not args.image.exists():
        raise FileNotFoundError(f"Image not found: {args.image}")

    runner = DynEvalRunner(args)
    result = runner.run(args.prompt, args.image, args)

    print(format_result(result))
    if args.output_file:
        args.output_file.parent.mkdir(parents=True, exist_ok=True)
        args.output_file.write_text(json_dumps(result) + "\n", encoding="utf-8")
        print(f"\nSaved JSON: {args.output_file}")


if __name__ == "__main__":
    main()