judge.py

"""
Qwen-Image-Bench Judge Model Inference Tool

Evaluate text-to-image generated images using a fine-tuned Qwen3.6-27B judge model.
Uses ms-swift PtEngine for batch inference.

Per-row output preserves all original input fields plus:
  - judge_model_output: combined raw scores JSON across all L1 dimensions
  - <dim>_judge_output: raw judge model text for each L1 dimension

Bench-level scores (L1 / L2 / Total) are aggregated following the
compute_scores.py methodology and saved alongside the per-row output.
"""

import argparse
import json
import sys
import time
from collections import defaultdict
from contextlib import contextmanager
from pathlib import Path

import pandas as pd
from PIL import Image
from tqdm import tqdm

from checklists import (
    DIM_TO_CHECKLIST,
    SYSTEM_PROMPT,
    USER_PROMPT_TEMPLATE,
    parse_dims_by_level1,
)
from score_utils import (
    aggregate_total_score,
    compute_dimension_score,
    extract_json_from_response,
    fix_score_json,
)


DIM_OUTPUT_MAP = {
    "Quality": "quality_judge_output",
    "Aesthetics": "aesthetics_judge_output",
    "Alignment": "alignment_judge_output",
    "Real-world Fidelity": "real_world_fidelity_judge_output",
    "Creative Generation": "creative_generation_judge_output",
}


@contextmanager
def _timer(label, store=None):
    """Measure wall time of a block, log it, and optionally record it in `store`."""
    start = time.perf_counter()
    yield
    elapsed = time.perf_counter() - start
    print(f"[TIMER] {label}: {elapsed:.2f}s")
    if store is not None:
        store[label] = elapsed


def load_and_resize_image(path):
    """Load image and resize to 1024x1024 if any dimension > 1024."""
    img = Image.open(path)
    if img.mode != "RGB":
        img = img.convert("RGB")
    if max(img.size) > 1024:
        img = img.resize((1024, 1024), Image.LANCZOS)
    img.load()
    return img


def load_input_file(file_path):
    """Load CSV or JSON/JSONL input file."""
    ext = Path(file_path).suffix.lower()
    if ext == ".csv":
        return pd.read_csv(file_path)
    elif ext == ".json":
        with open(file_path, "r", encoding="utf-8") as f:
            content = f.read().strip()
        if content.startswith("["):
            return pd.DataFrame(json.loads(content))
        else:
            records = [json.loads(line) for line in content.splitlines() if line.strip()]
            return pd.DataFrame(records)
    elif ext == ".jsonl":
        records = []
        with open(file_path, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    records.append(json.loads(line))
        return pd.DataFrame(records)
    else:
        raise ValueError(f"Unsupported file format: {ext}. Use .csv, .json, or .jsonl")


def load_bench_metadata(hf_bench_repo=None, local_metadata=None):
    """Load bench metadata containing dims_en per ID."""
    if local_metadata:
        return load_input_file(local_metadata)

    if hf_bench_repo:
        from huggingface_hub import hf_hub_download

        local_file = hf_hub_download(
            repo_id=hf_bench_repo,
            filename="qwen_image_bench_hf_v0518.jsonl",
            repo_type="dataset",
        )
        records = []
        with open(local_file, "r", encoding="utf-8") as f:
            for line in f:
                if line.strip():
                    records.append(json.loads(line))
        return pd.DataFrame(records)

    default_path = Path(__file__).parent / "metadata" / "bench_metadata.json"
    if default_path.exists():
        return load_input_file(str(default_path))

    print("ERROR: No metadata source found. Provide --hf-bench-repo, --local-metadata,")
    print(f"       or place metadata at {default_path}")
    sys.exit(1)


def _parse_output_to_scores(output_text, level1_dim):
    """Parse raw judge model output → fixed score_json. Returns None on failure."""
    score_json = extract_json_from_response(output_text)
    if score_json is None:
        return None
    return fix_score_json(score_json, level1_dim)


def _run_batch_inference(judge, args, input_df, metadata_df, desc, timings=None):
    """Batch inference over all rows and dimensions."""
    tasks = []
    task_meta = []  # (row_idx, level1_dim)
    skipped_rows = set()
    image_failures = 0

    with _timer("image load + task build", timings):
        for row_idx, (_, row) in enumerate(input_df.iterrows()):
            row_id = row["ID"]
            prompt = row["prompt"]
            image_path = row["image_path"]

            meta_row = metadata_df[metadata_df["ID"] == row_id]
            if meta_row.empty:
                continue

            dims_en = meta_row.iloc[0]["dims_en"]
            dims_by_level1 = parse_dims_by_level1(dims_en)

            try:
                img = load_and_resize_image(image_path)
            except Exception as e:
                image_failures += 1
                skipped_rows.add(row_idx)
                print(f"WARNING: Failed to load image for ID={row_id}, path={image_path}: {e}")
                continue

            for level1_dim in dims_by_level1:
                if level1_dim not in DIM_TO_CHECKLIST:
                    continue
                checklist = DIM_TO_CHECKLIST[level1_dim]
                user_text = USER_PROMPT_TEMPLATE.format(
                    prompt=prompt,
                    level1_dim=level1_dim,
                    format_checklist=checklist,
                )
                tasks.append({
                    "system_prompt": SYSTEM_PROMPT,
                    "user_text": user_text,
                    "image": img,
                })
                task_meta.append((row_idx, level1_dim))

    print(f"Total inference tasks: {len(tasks)}")
    all_outputs = []
    # vLLM does its own continuous batching, so submit everything in one call and
    # let max_num_seqs bound concurrency. PtEngine uses static batches.
    if args.backend == "vllm":
        chunk_size = max(len(tasks), 1)
    else:
        chunk_size = args.batch_size

    with _timer("batch inference", timings):
        for i in tqdm(range(0, len(tasks), chunk_size), desc=desc):
            batch = tasks[i:i + chunk_size]
            outputs = judge.generate_batch(batch)
            all_outputs.extend(outputs)

    row_dim_raw_scores = {}
    row_raw_outputs = {}
    parse_failures = 0

    with _timer("parse + assemble results", timings):
        for (row_idx, level1_dim), output_text in zip(task_meta, all_outputs):
            if row_idx not in row_dim_raw_scores:
                row_dim_raw_scores[row_idx] = {}
                row_raw_outputs[row_idx] = {}
            row_raw_outputs[row_idx][level1_dim] = output_text
            score_json = _parse_output_to_scores(output_text, level1_dim)
            if score_json is None:
                parse_failures += 1
            row_dim_raw_scores[row_idx][level1_dim] = score_json

        results = []
        all_dim_raw_scores = []
        for row_idx, (_, row) in enumerate(input_df.iterrows()):
            if row_idx in skipped_rows:
                results.append(_empty_result(row))
                all_dim_raw_scores.append({})
                continue
            dim_raw_scores = row_dim_raw_scores.get(row_idx, {})
            dim_raw_outputs = row_raw_outputs.get(row_idx, {})
            results.append(_build_row_result(row, dim_raw_scores, dim_raw_outputs))
            all_dim_raw_scores.append(dim_raw_scores)

    return results, parse_failures, all_dim_raw_scores, image_failures


def _build_judge(args, timings=None):
    """Construct the inference backend selected by --backend."""
    with _timer("model load", timings):
        if args.backend == "vllm":
            from backends.vllm_backend import VllmJudge

            print(f"Loading model from: {args.model} (backend=vllm)")
            return VllmJudge(
                model_path=args.model,
                max_new_tokens=args.max_new_tokens,
                max_num_seqs=args.max_num_seqs,
                gpu_memory_utilization=args.gpu_memory_utilization,
                tensor_parallel_size=args.tensor_parallel_size,
                max_model_len=args.max_model_len,
            )

        from backends.ms_swift_backend import MsSwiftJudge

        print(f"Loading model from: {args.model} (backend=pt)")
        return MsSwiftJudge(
            model_path=args.model,
            max_batch_size=args.max_batch_size,
            max_new_tokens=args.max_new_tokens,
        )


def run_inference(args, input_df, metadata_df, timings=None):
    """Run inference using the selected backend (vLLM or ms-swift PtEngine)."""
    judge = _build_judge(args, timings=timings)
    print("Model loaded successfully.")
    return _run_batch_inference(
        judge, args, input_df, metadata_df, desc="Batch inference", timings=timings
    )


def _empty_result(row):
    """Build an empty result row for skipped entries."""
    result = dict(row)
    result["judge_model_output"] = None
    for col in DIM_OUTPUT_MAP.values():
        result[col] = None
    return result


def _build_row_result(row, dim_raw_scores, dim_raw_outputs):
    """
    Build per-row JSONL record.

    Schema (in order):
      - all original row fields (transparent pass-through)
      - judge_model_output: JSON-serialized {L1_dim: fixed_score_json} for all parsed dims
      - <dim>_judge_output: raw judge text for each L1 dim
    """
    result = dict(row)

    raw_output = {
        dim_name: score_json
        for dim_name, score_json in dim_raw_scores.items()
        if score_json is not None
    }
    result["judge_model_output"] = (
        json.dumps(raw_output, ensure_ascii=False) if raw_output else None
    )

    for dim_name, col_name in DIM_OUTPUT_MAP.items():
        result[col_name] = dim_raw_outputs.get(dim_name)

    return result


def save_output(results, input_path):
    """Save per-row results to file in same directory as input."""
    input_p = Path(input_path)
    ext = input_p.suffix.lower()
    output_name = f"{input_p.stem}_judged{ext}"
    output_path = input_p.parent / output_name

    df = pd.DataFrame(results)

    if ext == ".csv":
        df.to_csv(output_path, index=False, encoding="utf-8")
    else:
        with open(output_path, "w", encoding="utf-8") as f:
            for _, row in df.iterrows():
                f.write(json.dumps(row.to_dict(), ensure_ascii=False) + "\n")

    return str(output_path)


def _safe_mean(xs):
    return sum(xs) / len(xs) if xs else None


def compute_bench_scores(all_dim_raw_scores):
    """
    Bench-level aggregation following compute_scores.py methodology:
      per-row L3→L2→L1→Total nested averaging,
      then arithmetic mean across rows (None values skipped).
    """
    l1_accum = defaultdict(list)
    l2_accum = defaultdict(lambda: defaultdict(list))
    total_accum = []

    for row_scores in all_dim_raw_scores:
        dim_results = {}
        for l1_dim, score_json in row_scores.items():
            if score_json is None:
                continue
            dim_results[l1_dim] = compute_dimension_score(score_json)

        row_total = aggregate_total_score(dim_results)
        if row_total is not None:
            total_accum.append(row_total)

        for l1_dim, dim_data in dim_results.items():
            if dim_data["level1_score"] is not None:
                l1_accum[l1_dim].append(dim_data["level1_score"])
            for l2_name, l2_score in dim_data["level2_scores"].items():
                if l2_score is not None:
                    l2_accum[l1_dim][l2_name].append(l2_score)

    return {
        "level1": {d: _safe_mean(v) for d, v in l1_accum.items()},
        "level2": {
            d: {l2: _safe_mean(v) for l2, v in l2d.items()}
            for d, l2d in l2_accum.items()
        },
        "total": _safe_mean(total_accum),
    }


def save_bench_scores(bench, input_path):
    """Save bench-level scores as JSON + Excel beside the input file."""
    input_p = Path(input_path)
    base = input_p.parent / f"{input_p.stem}_bench_scores"
    json_path = f"{base}.json"
    xlsx_path = f"{base}.xlsx"

    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(bench, f, ensure_ascii=False, indent=2)

    with pd.ExcelWriter(xlsx_path, engine="openpyxl") as writer:
        l1_rows = [{"Dimension": d, "Score": s} for d, s in bench["level1"].items()]
        l1_rows.append({"Dimension": "Total", "Score": bench["total"]})
        pd.DataFrame(l1_rows).to_excel(writer, sheet_name="Level-1 Summary", index=False)

        for dim, l2_dict in bench["level2"].items():
            if not l2_dict:
                continue
            df = pd.DataFrame(
                [{"Sub-dimension": l2, "Score": s} for l2, s in l2_dict.items()]
            )
            df.to_excel(writer, sheet_name=dim[:31], index=False)

    return json_path, xlsx_path


def print_bench_scores(bench):
    """Pretty-print bench-level scores to terminal."""
    print("\n" + "=" * 70)
    print("BENCH-LEVEL SCORES")
    print("=" * 70)
    for dim, score in bench["level1"].items():
        s = f"{score:.2f}" if score is not None else "N/A"
        print(f"  L1 {dim:30s}: {s}")
    total = bench["total"]
    total_str = f"{total:.2f}" if total is not None else "N/A"
    print(f"  {'TOTAL':33s}: {total_str}")
    print("-" * 70)
    for dim, l2_dict in bench["level2"].items():
        if not l2_dict:
            continue
        print(f"  [{dim}]")
        for l2, s in l2_dict.items():
            v = f"{s:.2f}" if s is not None else "N/A"
            print(f"    L2 {l2:28s}: {v}")
    print("=" * 70)


def main():
    parser = argparse.ArgumentParser(
        description="Qwen-Image-Bench Judge Model Inference Tool"
    )
    parser.add_argument("--input", required=True, help="Input CSV/JSON/JSONL with ID, prompt, image_path")
    parser.add_argument("--model", required=True, help="HuggingFace model ID or local model path")
    parser.add_argument("--hf-bench-repo", default=None, help="HF dataset repo for bench metadata")
    parser.add_argument("--local-metadata", default=None, help="Local metadata file path (skip HF download)")
    parser.add_argument("--backend", choices=["pt", "vllm"], default="vllm",
                        help="Inference backend: 'vllm' (continuous batching, default) or "
                             "'pt' (ms-swift PtEngine / HF static batching).")
    parser.add_argument("--max-batch-size", type=int, default=24,
                        help="ms-swift PtEngine max_batch_size (pt backend only; default: 24)")
    parser.add_argument("--max-new-tokens", type=int, default=4096)
    # vLLM-only knobs (ignored by the pt backend).
    parser.add_argument("--max-num-seqs", type=int, default=256,
                        help="vLLM max concurrent sequences (default: 256)")
    parser.add_argument("--gpu-memory-utilization", type=float, default=0.9,
                        help="vLLM GPU memory fraction (default: 0.9)")
    parser.add_argument("--tensor-parallel-size", type=int, default=1,
                        help="vLLM tensor parallel size / GPUs (default: 1)")
    parser.add_argument("--max-model-len", type=int, default=None,
                        help="vLLM max model context length (default: model config)")

    args = parser.parse_args()
    args.batch_size = args.max_batch_size

    timings = {}
    wall_start = time.perf_counter()

    # Load input
    print(f"Loading input: {args.input}")
    with _timer("input load", timings):
        input_df = load_input_file(args.input)
    required_cols = {"ID", "prompt", "image_path"}
    missing = required_cols - set(input_df.columns)
    if missing:
        print(f"ERROR: Input file missing required columns: {missing}")
        sys.exit(1)
    print(f"Input: {len(input_df)} rows")

    # Load metadata
    print("Loading bench metadata...")
    with _timer("metadata load", timings):
        metadata_df = load_bench_metadata(
            hf_bench_repo=args.hf_bench_repo,
            local_metadata=args.local_metadata,
        )
    print(f"Metadata: {len(metadata_df)} rows")

    # Run inference
    with _timer("inference (total)", timings):
        results, parse_failures, all_dim_raw_scores, image_failures = run_inference(
            args, input_df, metadata_df, timings=timings
        )

    # Save per-row JSONL
    with _timer("save per-row output", timings):
        saved_path = save_output(results, args.input)
    print(f"\nPer-row results saved to: {saved_path}")
    if image_failures:
        print(f"Skipped (broken images): {image_failures}")
    print(f"Parse failures: {parse_failures}")

    # Compute & save bench-level scores
    with _timer("compute bench scores", timings):
        bench = compute_bench_scores(all_dim_raw_scores)
    with _timer("save bench scores", timings):
        json_path, xlsx_path = save_bench_scores(bench, args.input)
    print(f"Bench scores saved to: {json_path}")
    print(f"Bench scores saved to: {xlsx_path}")
    print_bench_scores(bench)

    total_wall = time.perf_counter() - wall_start
    print("\n" + "=" * 70)
    print("TIMER SUMMARY")
    print("=" * 70)
    for label, elapsed in timings.items():
        print(f"  [TIMER] {label:28s}: {elapsed:8.2f}s")
    print("-" * 70)
    print(f"  {'total wall time':28s}: {total_wall:8.2f}s")
    print("=" * 70)


if __name__ == "__main__":
    main()