PortraitMT / evaluation /inference_aesthetic.py

Upload folder using huggingface_hub

1737ebf verified 6 months ago

18 kB

"""
Inference script for the aesthetic scoring model (score-token approach).
For each test image, runs 13 criterion queries + 1 total-score query,
maps score tokens to numeric values, then classifies into levels:
    A — Poor   (0 <= score < 5)
    B — Medium (5 <= score < 7)
    C — Good   (7 <= score <= 10)
Also answers the per-image MCQ by feeding the question+options directly.
Usage:
    CUDA_VISIBLE_DEVICES=0,1,2,3 python evaluation/inference_aesthetic.py
"""

import os
import sys
import json
import math
import re
import glob
import time
import torch
import multiprocessing as mp
from pathlib import Path
from tqdm import tqdm
from multiprocessing import Process
from PIL import Image

Image.MAX_IMAGE_PIXELS = None
mp.set_start_method("spawn", force=True)

project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from qwenvl.train.score_tokens import ScoreTokenRegistry
from qwenvl.data.aesthetic_dataset import CRITERIA_QUESTIONS, CRITERIA_ORDER, TOTAL_SCORE_QUESTION

# ── Config ────────────────────────────────────────────────────────────────────
CHECKPOINT = ""

IMAGE_DIR = "zijielou/PortraitCraft/main"

TEST_JSON = "PortraitCraft/track_1_test.json"

OUTPUT_JSON = "PortraitCraft/track_1_test_results_2.json"

WORKERS_PER_GPU = 1
MAX_NEW_TOKENS = 1024

# ── Score token registries (must match training) ─────────────────────────────
CRIT_BINS = 20
TOTAL_BINS = 101

crit_registry = ScoreTokenRegistry(n_bins=CRIT_BINS, score_min=1.0, score_max=10.0, prefix="score")
total_registry = ScoreTokenRegistry(n_bins=TOTAL_BINS, score_min=0.0, score_max=100.0, prefix="total")

# ── Level classification ─────────────────────────────────────────────────────
def score_to_level(score: float) -> str:
    if score < 5.0:
        return "A"
    elif score < 7.0:
        return "B"
    else:
        return "C"


# ── Image path resolver (test json has flat names, images live in subdirs) ───
_image_index = None

def build_image_index(image_dir: str) -> dict:
    index = {}
    for subdir in sorted(os.listdir(image_dir)):
        subdir_path = os.path.join(image_dir, subdir)
        if not os.path.isdir(subdir_path):
            continue
        for fname in os.listdir(subdir_path):
            index[fname] = os.path.join(subdir_path, fname)
    return index


def resolve_image_path(image_name: str, image_dir: str) -> str:
    global _image_index
    full = os.path.join(image_dir, image_name)
    if os.path.exists(full):
        return full
    if _image_index is None:
        _image_index = build_image_index(image_dir)
    return _image_index.get(image_name, full)


# ── Model wrapper ─────────────────────────────────────────────────────────────
class AestheticInferenceServer:
    def __init__(self, gpu_id: int):
        from transformers import Qwen3VLForConditionalGeneration, AutoProcessor

        self.device = torch.device(f"cuda:{gpu_id}")
        print(f"Loading model on GPU {gpu_id} from {CHECKPOINT}...")

        self.model = Qwen3VLForConditionalGeneration.from_pretrained(
            CHECKPOINT,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
        ).to(self.device).eval()

        self.processor = AutoProcessor.from_pretrained(CHECKPOINT)

        self.crit_token_ids = crit_registry.get_token_ids(self.processor.tokenizer)
        self.crit_values = crit_registry.get_value_tensor(device=self.device)
        self.total_token_ids = total_registry.get_token_ids(self.processor.tokenizer)
        self.total_values = total_registry.get_value_tensor(device=self.device)

        print(f"Model loaded on GPU {gpu_id}")

    def _generate(self, image_path: str, text: str, max_new_tokens: int = None) -> tuple:
        """Returns (generated_text, generated_ids, scores_tuple)."""
        if max_new_tokens is None:
            max_new_tokens = MAX_NEW_TOKENS

        img = Image.open(image_path).convert("RGB")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": text},
            ],
        }]

        inputs = self.processor.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt",
        ).to(self.device)

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                return_dict_in_generate=True,
                output_scores=True,
            )

        generated_ids = outputs.sequences[0][inputs["input_ids"].shape[1]:]
        generated_text = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=False)
        scores_tuple = outputs.scores

        return generated_text, generated_ids, scores_tuple

    def predict_criterion(self, image_path: str, criterion: str) -> dict:
        question = CRITERIA_QUESTIONS[criterion]
        gen_text, gen_ids, scores_tuple = self._generate(image_path, question)

        score_token_id_set = set(self.crit_token_ids)
        expected_score = None
        argmax_score = None
        argmax_token = None

        gen_ids_list = gen_ids.tolist()
        for step_idx, tid in enumerate(gen_ids_list):
            if tid in score_token_id_set:
                token_str = self.processor.tokenizer.decode([tid])
                argmax_token = token_str.strip()
                argmax_score = crit_registry.token_to_score(argmax_token)
                logits_at_step = scores_tuple[step_idx]
                expected_score = crit_registry.expected_score(
                    logits_at_step[0], self.crit_token_ids, self.crit_values
                ).item()
                break

        final_score = expected_score if expected_score is not None else argmax_score
        if final_score is None:
            print(f"    WARNING: no score token found for [{criterion}], raw output: {gen_text[:200]}", flush=True)
            final_score = 5.0

        return {
            "score": round(final_score, 2),
            "level": score_to_level(final_score),
            "argmax_token": argmax_token,
            "argmax_score": round(argmax_score, 2) if argmax_score is not None else None,
            "expected_score": round(expected_score, 2) if expected_score is not None else None,
            "raw_text": gen_text,
        }

    def predict_total(self, image_path: str) -> dict:
        gen_text, gen_ids, scores_tuple = self._generate(image_path, TOTAL_SCORE_QUESTION)

        total_token_id_set = set(self.total_token_ids)
        expected_score = None
        argmax_score = None
        argmax_token = None

        gen_ids_list = gen_ids.tolist()
        for step_idx, tid in enumerate(gen_ids_list):
            if tid in total_token_id_set:
                token_str = self.processor.tokenizer.decode([tid])
                argmax_token = token_str.strip()
                argmax_score = total_registry.token_to_score(argmax_token)
                logits_at_step = scores_tuple[step_idx]
                expected_score = total_registry.expected_score(
                    logits_at_step[0], self.total_token_ids, self.total_values
                ).item()
                break

        final_score = expected_score if expected_score is not None else argmax_score
        if final_score is None:
            print(f"    WARNING: no total token found, raw output: {gen_text[:200]}", flush=True)
            final_score = 50.0

        return {
            "total_score": round(final_score, 1),
            "argmax_token": argmax_token,
            "argmax_score": round(argmax_score, 1) if argmax_score is not None else None,
            "expected_score": round(expected_score, 1) if expected_score is not None else None,
            "raw_text": gen_text,
        }

    def answer_mcq(self, image_path: str, question: str, options: dict) -> dict:
        options_text = "\n".join([f"{k}. {v}" for k, v in options.items()])
        prompt = (
            f"{question}\n\n"
            f"Options:\n{options_text}\n\n"
            f"Answer with a single letter (A, B, C, or D)."
        )

        img = Image.open(image_path).convert("RGB")
        messages = [{
            "role": "user",
            "content": [
                {"type": "image", "image": img},
                {"type": "text", "text": prompt},
            ],
        }]
        inputs = self.processor.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True,
            return_dict=True, return_tensors="pt",
        ).to(self.device)

        think_ids = self.processor.tokenizer.encode("<think>\n", add_special_tokens=False)
        think_tensor = torch.tensor([think_ids], device=self.device)
        inputs["input_ids"] = torch.cat([inputs["input_ids"], think_tensor], dim=1)
        if "attention_mask" in inputs:
            inputs["attention_mask"] = torch.cat(
                [inputs["attention_mask"], torch.ones_like(think_tensor)], dim=1
            )

        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=2048,
                do_sample=False,
                return_dict_in_generate=True,
                output_scores=True,
            )

        generated_ids = outputs.sequences[0][inputs["input_ids"].shape[1]:]
        gen_text = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=False)

        text_to_search = gen_text.split("</think>")[-1] if "</think>" in gen_text else gen_text
        match = re.search(r'\b([A-D])\b', text_to_search)
        answer = match.group(1) if match else "A"
        return {"answer": answer, "raw_text": gen_text}


# ── Worker ────────────────────────────────────────────────────────────────────
def worker_run(worker_id: int, gpu_id: int, data_chunk: list):
    server = AestheticInferenceServer(gpu_id)

    out_path = OUTPUT_JSON + f".part{worker_id}.json"

    results = []
    if os.path.exists(out_path):
        try:
            with open(out_path, "r", encoding="utf-8") as f:
                results = json.load(f)
            print(f"Worker {worker_id} resuming from {len(results)} existing results")
        except Exception:
            results = []

    done_images = {r["image_path"] for r in results}

    total_to_process = len([it for it in data_chunk if it["image_path"] not in done_images])
    processed_count = 0
    t_start_all = time.time()

    for item_idx, item in enumerate(data_chunk):
        if item["image_path"] in done_images:
            continue

        image_path = resolve_image_path(item["image_path"], IMAGE_DIR)
        if not os.path.exists(image_path):
            print(f"[W{worker_id}] IMAGE NOT FOUND: {image_path}", flush=True)
            continue

        t_img_start = time.time()

        try:
            criteria_results = {}
            print(f"\n[W{worker_id}] ({'='*60})", flush=True)
            print(f"[W{worker_id}] Image {processed_count+1}/{total_to_process}: {item['image_path']}", flush=True)
            print(f"[W{worker_id}] {'-'*60}", flush=True)

            for crit_name in CRITERIA_ORDER:
                t_crit = time.time()
                pred = server.predict_criterion(image_path, crit_name)
                criteria_results[crit_name] = {
                    "level": pred["level"],
                    "score": pred["score"],
                    "argmax_token": pred.get("argmax_token"),
                    "argmax_score": pred.get("argmax_score"),
                    "expected_score": pred.get("expected_score"),
                    "raw_text": pred.get("raw_text", ""),
                }
                elapsed = time.time() - t_crit
                argmax_info = f"argmax={pred.get('argmax_token','?')}({pred.get('argmax_score','?')})" if pred.get('argmax_token') else "NO_TOKEN"
                print(f"[W{worker_id}]   {crit_name:45s} => E[score]={pred['score']:5.2f}  level={pred['level']:6s}  {argmax_info}  ({elapsed:.1f}s)", flush=True)
                print(f"[W{worker_id}]     raw_text: {pred.get('raw_text', '')[:300]}", flush=True)

            total_pred = server.predict_total(image_path)
            t_argmax = f"argmax={total_pred.get('argmax_token','?')}({total_pred.get('argmax_score','?')})" if total_pred.get('argmax_token') else "NO_TOKEN"
            print(f"[W{worker_id}]   {'Total Score':45s} => E[score]={total_pred['total_score']:5.1f}/100  {t_argmax}", flush=True)
            print(f"[W{worker_id}]     raw_text: {total_pred.get('raw_text', '')[:300]}", flush=True)

            mcq_result = server.answer_mcq(
                image_path,
                item.get("question", ""),
                item.get("options", {}),
            )
            mcq_answer = mcq_result["answer"]
            mcq_raw = mcq_result["raw_text"]
            print(f"[W{worker_id}]   {'MCQ Answer':45s} => {mcq_answer}", flush=True)
            print(f"[W{worker_id}]     raw_text: {mcq_raw[:500]}", flush=True)

            t_img_elapsed = time.time() - t_img_start
            processed_count += 1
            avg_time = (time.time() - t_start_all) / processed_count
            eta = avg_time * (total_to_process - processed_count)
            print(f"[W{worker_id}] {'-'*60}", flush=True)
            print(f"[W{worker_id}] Done in {t_img_elapsed:.1f}s | Progress: {processed_count}/{total_to_process} | ETA: {eta/60:.1f}min", flush=True)

            level_summary = {v["level"]: 0 for v in criteria_results.values()}
            for v in criteria_results.values():
                level_summary[v["level"]] += 1
            print(f"[W{worker_id}] Level distribution: {dict(level_summary)} | total={total_pred['total_score']} | mcq={mcq_answer}", flush=True)

            record = {
                "image_path": item["image_path"],
                "criteria": criteria_results,
                "total_score": total_pred["total_score"],
                "total_score_detail": {
                    "argmax_token": total_pred.get("argmax_token"),
                    "argmax_score": total_pred.get("argmax_score"),
                    "expected_score": total_pred.get("expected_score"),
                    "raw_text": total_pred.get("raw_text", ""),
                },
                "question": item.get("question", ""),
                "options": item.get("options", {}),
                "answer": mcq_answer,
                "mcq_raw_text": mcq_raw,
            }

            results.append(record)

            tmp_path = out_path + ".tmp"
            with open(tmp_path, "w", encoding="utf-8") as f:
                json.dump(results, f, ensure_ascii=False, indent=2)
            os.replace(tmp_path, out_path)
            print(f"[W{worker_id}] Saved to {out_path} ({len(results)} total records)", flush=True)

        except Exception as e:
            print(f"[W{worker_id}] ERROR on {item['image_path']}: {e}", flush=True)
            import traceback
            traceback.print_exc()
            continue

    total_elapsed = time.time() - t_start_all
    print(f"\n[W{worker_id}] ALL DONE | {len(results)} records | {total_elapsed/60:.1f}min total", flush=True)


# ── Merge ─────────────────────────────────────────────────────────────────────
def merge_results():
    files = glob.glob(OUTPUT_JSON + ".part*.json")
    unique = {}
    for f in files:
        with open(f, "r", encoding="utf-8") as fp:
            for x in json.load(fp):
                unique[x["image_path"]] = x
    all_results = list(unique.values())
    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    print(f"Merged {len(all_results)} results to {OUTPUT_JSON}")


# ── Main ──────────────────────────────────────────────────────────────────────
def main():
    with open(TEST_JSON, "r", encoding="utf-8") as f:
        data = json.load(f)

    done_set = set()
    for part_file in glob.glob(OUTPUT_JSON + ".part*.json"):
        try:
            with open(part_file, "r") as fp:
                for x in json.load(fp):
                    done_set.add(x["image_path"])
        except Exception:
            continue

    remaining = [item for item in data if item["image_path"] not in done_set]
    print(f"Total: {len(data)}, already done: {len(done_set)}, remaining: {len(remaining)}")

    if len(remaining) == 0:
        print("All done, merging...")
        merge_results()
        return

    NUM_GPUS = torch.cuda.device_count()
    TOTAL_WORKERS = NUM_GPUS * WORKERS_PER_GPU
    print(f"GPUs: {NUM_GPUS}, workers: {TOTAL_WORKERS}")

    chunk_size = math.ceil(len(remaining) / TOTAL_WORKERS)
    chunks = [remaining[i:i + chunk_size] for i in range(0, len(remaining), chunk_size)]

    processes = []
    for i, chunk in enumerate(chunks):
        gpu_id = i // WORKERS_PER_GPU
        p = Process(target=worker_run, args=(i, gpu_id, chunk))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    merge_results()


if __name__ == "__main__":
    main()

	"""
	Inference script for the aesthetic scoring model (score-token approach).

	For each test image, runs 13 criterion queries + 1 total-score query,
	maps score tokens to numeric values, then classifies into levels:
	A — Poor (0 <= score < 5)
	B — Medium (5 <= score < 7)
	C — Good (7 <= score <= 10)

	Also answers the per-image MCQ by feeding the question+options directly.

	Usage:
	CUDA_VISIBLE_DEVICES=0,1,2,3 python evaluation/inference_aesthetic.py
	"""

	import os
	import sys
	import json
	import math
	import re
	import glob
	import time
	import torch
	import multiprocessing as mp
	from pathlib import Path
	from tqdm import tqdm
	from multiprocessing import Process
	from PIL import Image

	Image.MAX_IMAGE_PIXELS = None
	mp.set_start_method("spawn", force=True)

	project_root = Path(__file__).parent.parent
	sys.path.insert(0, str(project_root))

	from qwenvl.train.score_tokens import ScoreTokenRegistry
	from qwenvl.data.aesthetic_dataset import CRITERIA_QUESTIONS, CRITERIA_ORDER, TOTAL_SCORE_QUESTION

	# ── Config ────────────────────────────────────────────────────────────────────
	CHECKPOINT = ""

	IMAGE_DIR = "zijielou/PortraitCraft/main"

	TEST_JSON = "PortraitCraft/track_1_test.json"

	OUTPUT_JSON = "PortraitCraft/track_1_test_results_2.json"

	WORKERS_PER_GPU = 1
	MAX_NEW_TOKENS = 1024

	# ── Score token registries (must match training) ─────────────────────────────
	CRIT_BINS = 20
	TOTAL_BINS = 101

	crit_registry = ScoreTokenRegistry(n_bins=CRIT_BINS, score_min=1.0, score_max=10.0, prefix="score")
	total_registry = ScoreTokenRegistry(n_bins=TOTAL_BINS, score_min=0.0, score_max=100.0, prefix="total")

	# ── Level classification ─────────────────────────────────────────────────────
	def score_to_level(score: float) -> str:
	if score < 5.0:
	return "A"
	elif score < 7.0:
	return "B"
	else:
	return "C"


	# ── Image path resolver (test json has flat names, images live in subdirs) ───
	_image_index = None

	def build_image_index(image_dir: str) -> dict:
	index = {}
	for subdir in sorted(os.listdir(image_dir)):
	subdir_path = os.path.join(image_dir, subdir)
	if not os.path.isdir(subdir_path):
	continue
	for fname in os.listdir(subdir_path):
	index[fname] = os.path.join(subdir_path, fname)
	return index


	def resolve_image_path(image_name: str, image_dir: str) -> str:
	global _image_index
	full = os.path.join(image_dir, image_name)
	if os.path.exists(full):
	return full
	if _image_index is None:
	_image_index = build_image_index(image_dir)
	return _image_index.get(image_name, full)


	# ── Model wrapper ─────────────────────────────────────────────────────────────
	class AestheticInferenceServer:
	def __init__(self, gpu_id: int):
	from transformers import Qwen3VLForConditionalGeneration, AutoProcessor

	self.device = torch.device(f"cuda:{gpu_id}")
	print(f"Loading model on GPU {gpu_id} from {CHECKPOINT}...")

	self.model = Qwen3VLForConditionalGeneration.from_pretrained(
	CHECKPOINT,
	torch_dtype=torch.bfloat16,
	attn_implementation="flash_attention_2",
	).to(self.device).eval()

	self.processor = AutoProcessor.from_pretrained(CHECKPOINT)

	self.crit_token_ids = crit_registry.get_token_ids(self.processor.tokenizer)
	self.crit_values = crit_registry.get_value_tensor(device=self.device)
	self.total_token_ids = total_registry.get_token_ids(self.processor.tokenizer)
	self.total_values = total_registry.get_value_tensor(device=self.device)

	print(f"Model loaded on GPU {gpu_id}")

	def _generate(self, image_path: str, text: str, max_new_tokens: int = None) -> tuple:
	"""Returns (generated_text, generated_ids, scores_tuple)."""
	if max_new_tokens is None:
	max_new_tokens = MAX_NEW_TOKENS

	img = Image.open(image_path).convert("RGB")
	messages = [{
	"role": "user",
	"content": [
	{"type": "image", "image": img},
	{"type": "text", "text": text},
	],
	}]

	inputs = self.processor.apply_chat_template(
	messages, tokenize=True, add_generation_prompt=True,
	return_dict=True, return_tensors="pt",
	).to(self.device)

	with torch.no_grad():
	outputs = self.model.generate(
	**inputs,
	max_new_tokens=max_new_tokens,
	do_sample=False,
	return_dict_in_generate=True,
	output_scores=True,
	)

	generated_ids = outputs.sequences[0][inputs["input_ids"].shape[1]:]
	generated_text = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=False)
	scores_tuple = outputs.scores

	return generated_text, generated_ids, scores_tuple

	def predict_criterion(self, image_path: str, criterion: str) -> dict:
	question = CRITERIA_QUESTIONS[criterion]
	gen_text, gen_ids, scores_tuple = self._generate(image_path, question)

	score_token_id_set = set(self.crit_token_ids)
	expected_score = None
	argmax_score = None
	argmax_token = None

	gen_ids_list = gen_ids.tolist()
	for step_idx, tid in enumerate(gen_ids_list):
	if tid in score_token_id_set:
	token_str = self.processor.tokenizer.decode([tid])
	argmax_token = token_str.strip()
	argmax_score = crit_registry.token_to_score(argmax_token)
	logits_at_step = scores_tuple[step_idx]
	expected_score = crit_registry.expected_score(
	logits_at_step[0], self.crit_token_ids, self.crit_values
	).item()
	break

	final_score = expected_score if expected_score is not None else argmax_score
	if final_score is None:
	print(f" WARNING: no score token found for [{criterion}], raw output: {gen_text[:200]}", flush=True)
	final_score = 5.0

	return {
	"score": round(final_score, 2),
	"level": score_to_level(final_score),
	"argmax_token": argmax_token,
	"argmax_score": round(argmax_score, 2) if argmax_score is not None else None,
	"expected_score": round(expected_score, 2) if expected_score is not None else None,
	"raw_text": gen_text,
	}

	def predict_total(self, image_path: str) -> dict:
	gen_text, gen_ids, scores_tuple = self._generate(image_path, TOTAL_SCORE_QUESTION)

	total_token_id_set = set(self.total_token_ids)
	expected_score = None
	argmax_score = None
	argmax_token = None

	gen_ids_list = gen_ids.tolist()
	for step_idx, tid in enumerate(gen_ids_list):
	if tid in total_token_id_set:
	token_str = self.processor.tokenizer.decode([tid])
	argmax_token = token_str.strip()
	argmax_score = total_registry.token_to_score(argmax_token)
	logits_at_step = scores_tuple[step_idx]
	expected_score = total_registry.expected_score(
	logits_at_step[0], self.total_token_ids, self.total_values
	).item()
	break

	final_score = expected_score if expected_score is not None else argmax_score
	if final_score is None:
	print(f" WARNING: no total token found, raw output: {gen_text[:200]}", flush=True)
	final_score = 50.0

	return {
	"total_score": round(final_score, 1),
	"argmax_token": argmax_token,
	"argmax_score": round(argmax_score, 1) if argmax_score is not None else None,
	"expected_score": round(expected_score, 1) if expected_score is not None else None,
	"raw_text": gen_text,
	}

	def answer_mcq(self, image_path: str, question: str, options: dict) -> dict:
	options_text = "\n".join([f"{k}. {v}" for k, v in options.items()])
	prompt = (
	f"{question}\n\n"
	f"Options:\n{options_text}\n\n"
	f"Answer with a single letter (A, B, C, or D)."
	)

	img = Image.open(image_path).convert("RGB")
	messages = [{
	"role": "user",
	"content": [
	{"type": "image", "image": img},
	{"type": "text", "text": prompt},
	],
	}]
	inputs = self.processor.apply_chat_template(
	messages, tokenize=True, add_generation_prompt=True,
	return_dict=True, return_tensors="pt",
	).to(self.device)

	think_ids = self.processor.tokenizer.encode("<think>\n", add_special_tokens=False)
	think_tensor = torch.tensor([think_ids], device=self.device)
	inputs["input_ids"] = torch.cat([inputs["input_ids"], think_tensor], dim=1)
	if "attention_mask" in inputs:
	inputs["attention_mask"] = torch.cat(
	[inputs["attention_mask"], torch.ones_like(think_tensor)], dim=1
	)

	with torch.no_grad():
	outputs = self.model.generate(
	**inputs,
	max_new_tokens=2048,
	do_sample=False,
	return_dict_in_generate=True,
	output_scores=True,
	)

	generated_ids = outputs.sequences[0][inputs["input_ids"].shape[1]:]
	gen_text = self.processor.tokenizer.decode(generated_ids, skip_special_tokens=False)

	text_to_search = gen_text.split("</think>")[-1] if "</think>" in gen_text else gen_text
	match = re.search(r'\b([A-D])\b', text_to_search)
	answer = match.group(1) if match else "A"
	return {"answer": answer, "raw_text": gen_text}


	# ── Worker ────────────────────────────────────────────────────────────────────
	def worker_run(worker_id: int, gpu_id: int, data_chunk: list):
	server = AestheticInferenceServer(gpu_id)

	out_path = OUTPUT_JSON + f".part{worker_id}.json"

	results = []
	if os.path.exists(out_path):
	try:
	with open(out_path, "r", encoding="utf-8") as f:
	results = json.load(f)
	print(f"Worker {worker_id} resuming from {len(results)} existing results")
	except Exception:
	results = []

	done_images = {r["image_path"] for r in results}

	total_to_process = len([it for it in data_chunk if it["image_path"] not in done_images])
	processed_count = 0
	t_start_all = time.time()

	for item_idx, item in enumerate(data_chunk):
	if item["image_path"] in done_images:
	continue

	image_path = resolve_image_path(item["image_path"], IMAGE_DIR)
	if not os.path.exists(image_path):
	print(f"[W{worker_id}] IMAGE NOT FOUND: {image_path}", flush=True)
	continue

	t_img_start = time.time()

	try:
	criteria_results = {}
	print(f"\n[W{worker_id}] ({'='*60})", flush=True)
	print(f"[W{worker_id}] Image {processed_count+1}/{total_to_process}: {item['image_path']}", flush=True)
	print(f"[W{worker_id}] {'-'*60}", flush=True)

	for crit_name in CRITERIA_ORDER:
	t_crit = time.time()
	pred = server.predict_criterion(image_path, crit_name)
	criteria_results[crit_name] = {
	"level": pred["level"],
	"score": pred["score"],
	"argmax_token": pred.get("argmax_token"),
	"argmax_score": pred.get("argmax_score"),
	"expected_score": pred.get("expected_score"),
	"raw_text": pred.get("raw_text", ""),
	}
	elapsed = time.time() - t_crit
	argmax_info = f"argmax={pred.get('argmax_token','?')}({pred.get('argmax_score','?')})" if pred.get('argmax_token') else "NO_TOKEN"
	print(f"[W{worker_id}] {crit_name:45s} => E[score]={pred['score']:5.2f} level={pred['level']:6s} {argmax_info} ({elapsed:.1f}s)", flush=True)
	print(f"[W{worker_id}] raw_text: {pred.get('raw_text', '')[:300]}", flush=True)

	total_pred = server.predict_total(image_path)
	t_argmax = f"argmax={total_pred.get('argmax_token','?')}({total_pred.get('argmax_score','?')})" if total_pred.get('argmax_token') else "NO_TOKEN"
	print(f"[W{worker_id}] {'Total Score':45s} => E[score]={total_pred['total_score']:5.1f}/100 {t_argmax}", flush=True)
	print(f"[W{worker_id}] raw_text: {total_pred.get('raw_text', '')[:300]}", flush=True)

	mcq_result = server.answer_mcq(
	image_path,
	item.get("question", ""),
	item.get("options", {}),
	)
	mcq_answer = mcq_result["answer"]
	mcq_raw = mcq_result["raw_text"]
	print(f"[W{worker_id}] {'MCQ Answer':45s} => {mcq_answer}", flush=True)
	print(f"[W{worker_id}] raw_text: {mcq_raw[:500]}", flush=True)

	t_img_elapsed = time.time() - t_img_start
	processed_count += 1
	avg_time = (time.time() - t_start_all) / processed_count
	eta = avg_time * (total_to_process - processed_count)
	print(f"[W{worker_id}] {'-'*60}", flush=True)
	print(f"[W{worker_id}] Done in {t_img_elapsed:.1f}s \| Progress: {processed_count}/{total_to_process} \| ETA: {eta/60:.1f}min", flush=True)

	level_summary = {v["level"]: 0 for v in criteria_results.values()}
	for v in criteria_results.values():
	level_summary[v["level"]] += 1
	print(f"[W{worker_id}] Level distribution: {dict(level_summary)} \| total={total_pred['total_score']} \| mcq={mcq_answer}", flush=True)

	record = {
	"image_path": item["image_path"],
	"criteria": criteria_results,
	"total_score": total_pred["total_score"],
	"total_score_detail": {
	"argmax_token": total_pred.get("argmax_token"),
	"argmax_score": total_pred.get("argmax_score"),
	"expected_score": total_pred.get("expected_score"),
	"raw_text": total_pred.get("raw_text", ""),
	},
	"question": item.get("question", ""),
	"options": item.get("options", {}),
	"answer": mcq_answer,
	"mcq_raw_text": mcq_raw,
	}

	results.append(record)

	tmp_path = out_path + ".tmp"
	with open(tmp_path, "w", encoding="utf-8") as f:
	json.dump(results, f, ensure_ascii=False, indent=2)
	os.replace(tmp_path, out_path)
	print(f"[W{worker_id}] Saved to {out_path} ({len(results)} total records)", flush=True)

	except Exception as e:
	print(f"[W{worker_id}] ERROR on {item['image_path']}: {e}", flush=True)
	import traceback
	traceback.print_exc()
	continue

	total_elapsed = time.time() - t_start_all
	print(f"\n[W{worker_id}] ALL DONE \| {len(results)} records \| {total_elapsed/60:.1f}min total", flush=True)


	# ── Merge ─────────────────────────────────────────────────────────────────────
	def merge_results():
	files = glob.glob(OUTPUT_JSON + ".part*.json")
	unique = {}
	for f in files:
	with open(f, "r", encoding="utf-8") as fp:
	for x in json.load(fp):
	unique[x["image_path"]] = x
	all_results = list(unique.values())
	with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
	json.dump(all_results, f, ensure_ascii=False, indent=2)
	print(f"Merged {len(all_results)} results to {OUTPUT_JSON}")


	# ── Main ──────────────────────────────────────────────────────────────────────
	def main():
	with open(TEST_JSON, "r", encoding="utf-8") as f:
	data = json.load(f)

	done_set = set()
	for part_file in glob.glob(OUTPUT_JSON + ".part*.json"):
	try:
	with open(part_file, "r") as fp:
	for x in json.load(fp):
	done_set.add(x["image_path"])
	except Exception:
	continue

	remaining = [item for item in data if item["image_path"] not in done_set]
	print(f"Total: {len(data)}, already done: {len(done_set)}, remaining: {len(remaining)}")

	if len(remaining) == 0:
	print("All done, merging...")
	merge_results()
	return

	NUM_GPUS = torch.cuda.device_count()
	TOTAL_WORKERS = NUM_GPUS * WORKERS_PER_GPU
	print(f"GPUs: {NUM_GPUS}, workers: {TOTAL_WORKERS}")

	chunk_size = math.ceil(len(remaining) / TOTAL_WORKERS)
	chunks = [remaining[i:i + chunk_size] for i in range(0, len(remaining), chunk_size)]

	processes = []
	for i, chunk in enumerate(chunks):
	gpu_id = i // WORKERS_PER_GPU
	p = Process(target=worker_run, args=(i, gpu_id, chunk))
	p.start()
	processes.append(p)

	for p in processes:
	p.join()

	merge_results()


	if __name__ == "__main__":
	main()