Upload folder using huggingface_hub

evaluation/inference_aesthetic.py

@@ -0,0 +1,436 @@
+"""
+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()

qwenvl/data/__init__.py

@@ -0,0 +1,77 @@
+import re
+
+# Define placeholders for dataset paths
+CAMBRIAN_737K = {
+    "annotation_path": "PATH_TO_CAMBRIAN_737K_ANNOTATION",
+    "data_path": "",
+}
+
+CAMBRIAN_737K_PACK = {
+    "annotation_path": f"PATH_TO_CAMBRIAN_737K_ANNOTATION_PACKED",
+    "data_path": f"",
+}
+
+MP_DOC = {
+    "annotation_path": "PATH_TO_MP_DOC_ANNOTATION",
+    "data_path": "PATH_TO_MP_DOC_DATA",
+}
+
+CLEVR_MC = {
+    "annotation_path": "PATH_TO_CLEVR_MC_ANNOTATION",
+    "data_path": "PATH_TO_CLEVR_MC_DATA",
+}
+
+VIDEOCHATGPT = {
+    "annotation_path": "PATH_TO_VIDEOCHATGPT_ANNOTATION",
+    "data_path": "PATH_TO_VIDEOCHATGPT_DATA",
+}
+
+SINGLEIMAGES = {
+    "annotation_path": "./demo/single_images_train_convert.json",
+    "data_path": "./demo/images",
+}
+
+AESTHETIC_TRACK1 = {
+    # image_path in json uses subdir format: images_00/xxx.jpg
+    "annotation_path": "PortraitCraft/track_1_train_subdir.json",
+    "data_path": "zijielou/PortraitCraft/main",
+    "dataset_type": "aesthetic",
+}
+
+data_dict = {
+    "cambrian_737k": CAMBRIAN_737K,
+    "cambrian_737k_pack": CAMBRIAN_737K_PACK,
+    "mp_doc": MP_DOC,
+    "clevr_mc": CLEVR_MC,
+    "videochatgpt": VIDEOCHATGPT,
+    "singleimages": SINGLEIMAGES,
+    "aesthetic_track1": AESTHETIC_TRACK1,
+}
+
+
+def parse_sampling_rate(dataset_name):
+    match = re.search(r"%(\d+)$", dataset_name)
+    if match:
+        return int(match.group(1)) / 100.0
+    return 1.0
+
+
+def data_list(dataset_names):
+    config_list = []
+    for dataset_name in dataset_names:
+        sampling_rate = parse_sampling_rate(dataset_name)
+        dataset_name = re.sub(r"%(\d+)$", "", dataset_name)
+        if dataset_name in data_dict.keys():
+            config = data_dict[dataset_name].copy()
+            config["sampling_rate"] = sampling_rate
+            config_list.append(config)
+        else:
+            raise ValueError(f"do not find {dataset_name}")
+    return config_list
+
+
+if __name__ == "__main__":
+    dataset_names = ["cambrian_737k"]
+    configs = data_list(dataset_names)
+    for config in configs:
+        print(config)

qwenvl/data/aesthetic_dataset.py

@@ -0,0 +1,419 @@
+"""
+Dataset for image aesthetic scoring using generic integer-index score tokens.
+
+Each image produces N+1 independent training samples:
+  - one sample per criterion  (N=13)
+  - one sample for total score
+
+Every sample is a single-turn VQA pair:
+    user:      <image> + criterion-specific question
+    assistant: <think>reason</think>\nThe score is <scoreK>
+
+The <scoreK> token is what we use for the logits-weighted prediction.
+The <think>...</think> block is labelled for CE loss so the model learns to reason.
+"""
+
+import json
+import time
+from pathlib import Path
+from typing import Dict, Any, List, Tuple
+
+import torch
+from torch.utils.data import Dataset
+
+from .rope2d import get_rope_index_25, get_rope_index_2, get_rope_index_3
+from ..train.score_tokens import ScoreTokenRegistry
+
+IGNORE_INDEX = -100
+
+# ── Per-criterion questions ───────────────────────────────────────────────────
+# Each question is carefully designed to probe the specific visual dimension.
+# Default instruction uses the original <score1>..<score20> tokens.
+# When existing vocab tokens are used, build_criteria_questions() overrides this.
+CRIT_SCORE_INSTRUCTION = (
+    "Rate this aspect using a score token from <score1> (lowest quality) to <score20> (highest quality)."
+)
+
+_CRITERIA_QUESTION_BODIES: Dict[str, str] = {
+    "Color Harmony": (
+        "Evaluate the color harmony of this image. "
+        "Consider how well the colors work together in terms of palette coherence, "
+        "contrast balance, and emotional consistency. "
+    ),
+    "Visual Style Consistency": (
+        "Assess the visual style consistency of this image. "
+        "Does the overall aesthetic — including lighting, color grading, subject treatment, "
+        "and compositional choices — feel unified and purposeful throughout? "
+    ),
+    "Sharpness": (
+        "Evaluate the sharpness of this image. "
+        "Focus on whether the intended subject is in crisp focus, "
+        "the level of fine detail preserved, and whether any blur is intentional or distracting. "
+    ),
+    "Light and Shadow Modeling": (
+        "Assess the light and shadow modeling in this image. "
+        "Consider the quality of tonal transitions, whether shadows flatter the subject, "
+        "and how well the lighting reveals three-dimensional form. "
+    ),
+    "Creativity and Originality": (
+        "Evaluate the creativity and originality of this image. "
+        "Does it demonstrate a distinctive perspective, novel visual language, "
+        "or an unexpected approach that sets it apart from typical photography? "
+    ),
+    "Exposure Control": (
+        "Assess the exposure control of this image. "
+        "Consider whether highlights are blown out, shadows are crushed, "
+        "and whether the overall tonal range is balanced and intentional. "
+    ),
+    "Application of Classical Composition Principles": (
+        "Evaluate how well classical composition principles are applied in this image. "
+        "Consider elements such as rule of thirds, leading lines, framing, "
+        "symmetry, and the intentionality of the subject placement within the frame. "
+    ),
+    "Depth of Field and Layering": (
+        "Assess the depth of field and visual layering in this image. "
+        "Consider whether foreground, midground, and background are purposefully separated, "
+        "and whether the depth of field effectively isolates the subject. "
+    ),
+    "Visual Center Stability": (
+        "Evaluate the visual center stability of this image. "
+        "Consider whether there is a clear focal point that anchors the viewer's gaze, "
+        "and whether competing elements undermine or reinforce that center. "
+    ),
+    "Visual Flow Guidance": (
+        "Assess the visual flow guidance in this image. "
+        "Does the composition lead the viewer's eye naturally through the frame "
+        "using lines, shapes, contrast, or subject placement? "
+    ),
+    "Structural Support Stability": (
+        "Evaluate the structural support and stability of this image's composition. "
+        "Consider whether the arrangement of elements creates a balanced, grounded feel "
+        "or whether the structure feels unresolved or tension-inducing. "
+    ),
+    "Appropriateness of Negative Space": (
+        "Assess the use of negative space in this image. "
+        "Consider whether empty areas enhance the subject's presence, "
+        "create breathing room, and serve a compositional or emotional purpose. "
+    ),
+    "Subject Integrity": (
+        "Evaluate the subject integrity of this image. "
+        "Are the main subjects fully visible, clearly defined, and free from "
+        "awkward cropping, occlusion, or distracting intrusions that compromise their presence? "
+    ),
+}
+
+
+def build_criteria_questions(crit_registry) -> Dict[str, str]:
+    """Build criteria questions with the actual token names from the registry."""
+    first_token = crit_registry.tokens[0]
+    last_token = crit_registry.tokens[-1]
+    instruction = (
+        f"Rate this aspect using a score token from {first_token} (lowest quality) "
+        f"to {last_token} (highest quality)."
+    )
+    return {k: v + instruction for k, v in _CRITERIA_QUESTION_BODIES.items()}
+
+
+def build_total_score_question(total_registry) -> str:
+    """Build total score question with the actual token names from the registry."""
+    first_token = total_registry.tokens[0]
+    last_token = total_registry.tokens[-1]
+    return (
+        "Give an overall aesthetic quality score for this image. "
+        "Consider all aspects: composition, exposure, sharpness, color, creativity, "
+        "and subject presentation. "
+        f"Rate using a score token from {first_token} (lowest, 0) to {last_token} (highest, 100)."
+    )
+
+
+# Default questions (backward compatible with original <score1>..<score20> tokens)
+CRITERIA_QUESTIONS: Dict[str, str] = {
+    k: v + CRIT_SCORE_INSTRUCTION for k, v in _CRITERIA_QUESTION_BODIES.items()
+}
+
+TOTAL_SCORE_QUESTION = (
+    "Give an overall aesthetic quality score for this image. "
+    "Consider all aspects: composition, exposure, sharpness, color, creativity, "
+    "and subject presentation. "
+    "Rate using a score token from <total1> (lowest, 0) to <total101> (highest, 100)."
+)
+
+# Ordered list used to keep sample ordering deterministic
+CRITERIA_ORDER: List[str] = list(CRITERIA_QUESTIONS.keys())
+
+
+# ── Text builders ─────────────────────────────────────────────────────────────
+
+def build_crit_assistant_text(
+    reason: str,
+    score: float,
+    crit_registry: ScoreTokenRegistry,
+) -> str:
+    token = crit_registry.score_to_token(score)
+    return f"<think>\n{reason.strip()}\n</think>\nThe score is {token}"
+    # return  f"The score is {token}"
+
+def build_total_assistant_text(
+    total_score: float,
+    total_registry: ScoreTokenRegistry,
+) -> str:
+    token = total_registry.score_to_token(total_score)
+    return f"The score is {token}"
+
+
+
+def build_mcq_messages(
+    image_path: str,
+    question: str,
+    options: dict,
+    answer: str,
+) -> List[Dict]:
+    options_text = "\n".join(f"{k}. {v}" for k, v in sorted(options.items()))
+    user_text = f"{question}\n{options_text}"
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image_path},
+                {"type": "text", "text": user_text},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": answer},
+            ],
+        },
+    ]
+
+
+def build_crit_messages(
+    image_path: str,
+    criterion: str,
+    reason: str,
+    score: float,
+    crit_registry,
+    criteria_questions: Dict[str, str] = None,
+) -> List[Dict]:
+    questions = criteria_questions if criteria_questions is not None else CRITERIA_QUESTIONS
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image_path},
+                {"type": "text", "text": questions[criterion]},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": build_crit_assistant_text(reason, score, crit_registry)},
+            ],
+        },
+    ]
+
+
+def build_total_messages(
+    image_path: str,
+    total_score: float,
+    total_registry,
+    total_score_question: str = None,
+) -> List[Dict]:
+    question = total_score_question if total_score_question is not None else TOTAL_SCORE_QUESTION
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image_path},
+                {"type": "text", "text": question},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": [
+                {"type": "text", "text": build_total_assistant_text(total_score, total_registry)},
+            ],
+        },
+    ]
+
+
+# ── Tokenisation ──────────────────────────────────────────────────────────────
+
+def _tokenise_messages(messages: List[Dict], processor) -> Dict:
+    full_result = processor.apply_chat_template(
+        messages, tokenize=True, return_dict=True, return_tensors="pt"
+    )
+    input_ids = full_result["input_ids"]
+    if isinstance(input_ids, list):
+        input_ids = torch.tensor(input_ids).unsqueeze(0)
+
+    labels = torch.full_like(input_ids, IGNORE_INDEX)
+
+    # Label the assistant turn only (after token 77091 = "assistant")
+    ids_flat = input_ids[0].tolist()
+    L = len(ids_flat)
+    pos = 0
+    while pos < L:
+        if ids_flat[pos] == 77091:
+            ans_start = pos + 2
+            ans_end = ans_start
+            while ans_end < L and ids_flat[ans_end] != 151645:
+                ans_end += 1
+            if ans_end < L:
+                labels[0, ans_start : ans_end + 2] = input_ids[0, ans_start : ans_end + 2]
+                pos = ans_end
+        pos += 1
+
+    full_result["labels"] = labels
+    full_result["input_ids"] = input_ids
+    return full_result
+
+
+# ── Dataset ───────────────────────────────────────────────────────────────────
+
+class AestheticScoringDataset(Dataset):
+    """
+    Each image in track_1_train.json expands into len(criteria) + 1 samples.
+
+    Args:
+        processor:       HF processor with score tokens already added.
+        annotation_path: Path to track_1_train.json.
+        image_dir:       Root directory for image files.
+        data_args:       DataArguments (needs model_type).
+        crit_registry:   ScoreTokenRegistry for criteria scores (1–10).
+        total_registry:  ScoreTokenRegistry for total score (0–100).
+        include_total:   Whether to include total-score samples (default True).
+    """
+
+    def __init__(
+        self,
+        processor,
+        annotation_path: str,
+        image_dir: str,
+        data_args,
+        crit_registry,
+        total_registry,
+        include_total: bool = True,
+    ):
+        super().__init__()
+        with open(annotation_path) as f:
+            raw = json.load(f)
+
+        self.image_dir = Path(image_dir)
+        self.processor = processor
+        self.crit_registry = crit_registry
+        self.total_registry = total_registry
+
+        # Build questions dynamically based on registry token names
+        self.criteria_questions = build_criteria_questions(crit_registry)
+        self.total_score_question = build_total_score_question(total_registry)
+
+        model_type = getattr(data_args, "model_type", "qwen3vl")
+        if model_type == "qwen3vl":
+            self.get_rope_index = get_rope_index_3
+        elif model_type == "qwen2.5vl":
+            self.get_rope_index = get_rope_index_25
+        elif model_type == "qwen2vl":
+            self.get_rope_index = get_rope_index_2
+        else:
+            raise ValueError(f"Unsupported model_type: {model_type}")
+
+        self.merge_size = getattr(processor.image_processor, "merge_size", 2)
+
+        # Expand: each item → (image_path, type_key, score, text_field, question, options)
+        # type_key: criterion name | "__total__" | "__mcq__"
+        self.samples: List[Tuple] = []
+        for item in raw:
+            img = item["image_path"]
+            if "question" in item:
+                # MCQ entry: single sample
+                self.samples.append((
+                    img, "__mcq__", 0.0,
+                    item["answer"], item["question"], item["options"],
+                ))
+            else:
+                # Scoring entry: 13 criteria + 1 total
+                criteria = item.get("criteria", {})
+                for key in CRITERIA_ORDER:
+                    if key in criteria:
+                        entry = criteria[key]
+                        self.samples.append((
+                            img,
+                            key,
+                            float(entry["score"]),
+                            entry.get("reason", ""),
+                            None, None,
+                        ))
+                if include_total:
+                    self.samples.append((
+                        img,
+                        "__total__",
+                        float(item["total_score"]),
+                        "",
+                        None, None,
+                    ))
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, i: int) -> Dict[str, torch.Tensor]:
+        for attempt in range(3):
+            try:
+                return self._get_item(i)
+            except Exception as e:
+                img_name = self.samples[i][0]
+                print(f"[Attempt {attempt}] Error loading sample {i} (image={img_name}): {e}")
+                # skip to a random different sample instead of retrying the same broken image
+                import random
+                i = random.randint(0, len(self.samples) - 1)
+                time.sleep(0.1)
+        return self._get_item(random.randint(0, len(self.samples) - 1))
+
+    def _get_item(self, i: int) -> Dict[str, torch.Tensor]:
+        img_name, criterion, score, text_field, question, options = self.samples[i]
+
+        if criterion == "__mcq__":
+            image_path = str((self.image_dir / img_name).resolve())
+            messages = build_mcq_messages(image_path, question, options, text_field)
+        elif criterion == "__total__":
+            image_path = str((self.image_dir / img_name).resolve())
+            messages = build_total_messages(
+                image_path, score, self.total_registry,
+                total_score_question=self.total_score_question,
+            )
+        else:
+            image_path = str((self.image_dir / img_name).resolve())
+            messages = build_crit_messages(
+                image_path, criterion, text_field, score, self.crit_registry,
+                criteria_questions=self.criteria_questions,
+            )
+
+        data_dict = _tokenise_messages(messages, self.processor)
+
+        seq_len = data_dict["input_ids"][0].size(0)
+        # print(f"seq_length:{seq_len}")
+        grid_thw = data_dict.get("image_grid_thw")
+        if grid_thw is not None and not isinstance(grid_thw, (list, tuple)):
+            grid_thw = [grid_thw]
+
+        position_ids, _ = self.get_rope_index(
+            self.merge_size,
+            data_dict["input_ids"],
+            image_grid_thw=torch.cat(grid_thw, dim=0) if grid_thw else None,
+            video_grid_thw=None,
+            second_per_grid_ts=None,
+        )
+
+        data_dict["position_ids"] = position_ids
+        data_dict["attention_mask"] = [seq_len]
+        # sample_type: 0=total, 1=crit, 2=mcq (mcq skips regression loss)
+        if criterion == "__mcq__":
+            data_dict["sample_type"] = 2
+        elif criterion == "__total__":
+            data_dict["sample_type"] = 0
+        else:
+            data_dict["sample_type"] = 1
+        data_dict["gt_score"] = score
+        return data_dict
+
+

qwenvl/data/data_processor.py

@@ -0,0 +1,693 @@
+import json
+import random
+import logging
+import re
+import time
+import itertools
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Sequence, List, Tuple, Any
+from collections.abc import Sequence
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+import transformers
+
+from . import data_list
+from .rope2d import get_rope_index_25, get_rope_index_2, get_rope_index_3
+
+IGNORE_INDEX = -100
+IMAGE_TOKEN_INDEX = 151655
+VIDEO_TOKEN_INDEX = 151656
+DEFAULT_IMAGE_TOKEN = "<image>"
+DEFAULT_VIDEO_TOKEN = "<video>"
+
+local_rank = None
+
+
+def rank0_print(*args):
+    if local_rank == 0:
+        print(*args)
+
+
+def read_jsonl(path):
+    with open(path, "r") as f:
+        return [json.loads(line) for line in f]
+
+
+def _make_abs_paths(base: Path, files: str) -> str:
+    return f"{(base / files).resolve()}"
+
+
+def update_processor_pixels(processor, data_args):
+    logger = logging.getLogger(__name__)
+
+    # --- Image Processor ---
+    ip = processor.image_processor
+    rank0_print("=== BEFORE IMAGE PROCESSOR PARAMETERS ===")
+    rank0_print(f"Image min_pixels: {getattr(ip, 'min_pixels', 'N/A')}")
+    rank0_print(f"Image max_pixels: {getattr(ip, 'max_pixels', 'N/A')}")
+    rank0_print(f"ip.size: {ip.size}")
+    rank0_print(f"Image size (shortest_edge): {ip.size.get('shortest_edge', 'N/A')}")
+    rank0_print(f"Image size (longest_edge):  {ip.size.get('longest_edge', 'N/A')}")
+
+    if hasattr(ip, "min_pixels") and hasattr(ip, "max_pixels"):
+        ip.min_pixels = data_args.min_pixels
+        ip.max_pixels = data_args.max_pixels
+        rank0_print(f"✅ Updated image_processor min_pixels to {data_args.min_pixels}")
+        rank0_print(f"✅ Updated image_processor max_pixels to {data_args.max_pixels}")
+
+    if hasattr(ip, "size") and isinstance(ip.size, dict):
+        ip.size["shortest_edge"] = data_args.min_pixels
+        ip.size["longest_edge"] = data_args.max_pixels
+        rank0_print(
+            f"✅ Updated image_processor size['shortest_edge'] to {data_args.min_pixels}"
+        )
+        rank0_print(
+            f"✅ Updated image_processor size['longest_edge'] to {data_args.max_pixels}"
+        )
+
+    rank0_print("=== AFTER IMAGE PROCESSOR PARAMETERS ===")
+    rank0_print(f"Image min_pixels: {getattr(ip, 'min_pixels', 'N/A')}")
+    rank0_print(f"Image max_pixels: {getattr(ip, 'max_pixels', 'N/A')}")
+    rank0_print(f"Image size (shortest_edge): {ip.size.get('shortest_edge', 'N/A')}")
+    rank0_print(f"Image size (longest_edge):  {ip.size.get('longest_edge', 'N/A')}")
+
+    # --- Video Processor ---
+    if hasattr(processor, "video_processor") and processor.video_processor is not None:
+        vp = processor.video_processor
+        rank0_print("\n=== BEFORE VIDEO PROCESSOR PARAMETERS ===")
+        rank0_print(f"Video min_pixels: {getattr(vp, 'min_pixels', 'N/A')}")
+        rank0_print(f"Video max_pixels: {getattr(vp, 'max_pixels', 'N/A')}")
+        rank0_print(f"Video min_frames: {getattr(vp, 'min_frames', 'N/A')}")
+        rank0_print(f"Video max_frames: {getattr(vp, 'max_frames', 'N/A')}")
+        rank0_print(f"Video fps: {getattr(vp, 'fps', 'N/A')}")
+        rank0_print(
+            f"Video size (shortest_edge): {vp.size.get('shortest_edge', 'N/A')}"
+        )
+        rank0_print(f"Video size (longest_edge):  {vp.size.get('longest_edge', 'N/A')}")
+
+        if hasattr(vp, "min_pixels") and hasattr(vp, "max_pixels"):
+            vp.min_pixels = data_args.video_min_pixels
+            vp.max_pixels = data_args.video_max_pixels
+            rank0_print(
+                f"✅ Updated Qwen2-VL video_processor min_pixels to {data_args.video_min_pixels}"
+            )
+            rank0_print(
+                f"✅ Updated Qwen2-VL video_processor max_pixels to {data_args.video_max_pixels}"
+            )
+
+        if hasattr(vp, "min_frames") and hasattr(vp, "max_frames"):
+            vp.min_frames = data_args.video_min_frames
+            vp.max_frames = data_args.video_max_frames
+            rank0_print(
+                f"✅ Updated video_processor min_frames to {data_args.video_min_frames}"
+            )
+            rank0_print(
+                f"✅ Updated video_processor max_frames to {data_args.video_max_frames}"
+            )
+
+        if hasattr(vp, "fps"):
+            vp.fps = data_args.video_fps
+            rank0_print(f"✅ Updated video_processor fps to {data_args.video_fps}")
+
+        if hasattr(vp, "size") and isinstance(vp.size, dict):
+            vp.size["shortest_edge"] = data_args.video_min_pixels
+            vp.size["longest_edge"] = data_args.video_max_pixels
+            rank0_print(
+                f"✅ Updated Video size (shortest_edge): {vp.size.get('shortest_edge', 'N/A')}"
+            )
+            rank0_print(
+                f"✅ Updated Video size (longest_edge):  {vp.size.get('longest_edge', 'N/A')}"
+            )
+
+        rank0_print("=== AFTER VIDEO PROCESSOR PARAMETERS ===")
+        rank0_print(f"Video min_pixels: {getattr(vp, 'min_pixels', 'N/A')}")
+        rank0_print(f"Video max_pixels: {getattr(vp, 'max_pixels', 'N/A')}")
+        rank0_print(f"Video min_frames: {getattr(vp, 'min_frames', 'N/A')}")
+        rank0_print(f"Video max_frames: {getattr(vp, 'max_frames', 'N/A')}")
+        rank0_print(f"Video fps: {getattr(vp, 'fps', 'N/A')}")
+        rank0_print(
+            f"Video size (shortest_edge): {vp.size.get('shortest_edge', 'N/A')}"
+        )
+        rank0_print(f"Video size (longest_edge):  {vp.size.get('longest_edge', 'N/A')}")
+
+    return processor
+
+
+def _build_messages(item: Dict[str, Any], base_path: Path) -> List[Dict[str, Any]]:
+    # Extract and normalize images and videos
+    images = item.get("image") or []
+    if isinstance(images, str):
+        images = [images]
+
+    videos = item.get("video") or []
+    if isinstance(videos, str):
+        videos = [videos]
+
+    # Build media pools with absolute paths
+    image_pool = [
+        {"type": "image", "image": _make_abs_paths(base_path, img)} for img in images
+    ]
+    video_pool = [
+        {"type": "video", "video": _make_abs_paths(base_path, vid)} for vid in videos
+    ]
+
+    messages = []
+    for turn in item["conversations"]:
+        role = "user" if turn["from"] == "human" else "assistant"
+        text: str = turn["value"]
+
+        if role == "user":
+            content = []
+            # Split text by <image> or <video> placeholders while keeping delimiters
+            text_parts = re.split(r"(<image>|<video>)", text)
+
+            for seg in text_parts:
+                if seg == "<image>":
+                    if not image_pool:
+                        raise ValueError(
+                            "Number of <image> placeholders exceeds the number of provided images"
+                        )
+                    content.append(image_pool.pop(0))
+                elif seg == "<video>":
+                    if not video_pool:
+                        raise ValueError(
+                            "Number of <video> placeholders exceeds the number of provided videos"
+                        )
+                    content.append(video_pool.pop(0))
+                elif seg.strip():
+                    content.append({"type": "text", "text": seg.strip()})
+
+            messages.append({"role": role, "content": content})
+        else:
+            # Assistant messages contain only text
+            messages.append({"role": role, "content": [{"type": "text", "text": text}]})
+
+    # Check for unused media files
+    if image_pool:
+        raise ValueError(
+            f"{len(image_pool)} image(s) remain unused (not consumed by placeholders)"
+        )
+    if video_pool:
+        raise ValueError(
+            f"{len(video_pool)} video(s) remain unused (not consumed by placeholders)"
+        )
+
+    return messages
+
+
+def preprocess_qwen_visual(
+    sources,
+    processor,
+) -> Dict:
+    if len(sources) != 1:
+        raise ValueError(f"Expected 1 source, got {len(sources)}")
+
+    source = sources[0]
+    base_path = Path(source.get("data_path", ""))
+    messages = _build_messages(source, base_path)
+
+    full_result = processor.apply_chat_template(
+        messages, tokenize=True, return_dict=True, return_tensors="pt"
+    )
+
+    input_ids = full_result["input_ids"]
+    if isinstance(input_ids, list):
+        input_ids = torch.tensor(input_ids).unsqueeze(0)
+
+    labels = torch.full_like(input_ids, IGNORE_INDEX)
+
+    input_ids_flat = input_ids[0].tolist()
+    L = len(input_ids_flat)
+    pos = 0
+    while pos < L:
+        if input_ids_flat[pos] == 77091:
+            ans_start = pos + 2
+            ans_end = ans_start
+            while ans_end < L and input_ids_flat[ans_end] != 151645:
+                ans_end += 1
+            if ans_end < L:
+                labels[0, ans_start : ans_end + 2] = input_ids[
+                    0, ans_start : ans_end + 2
+                ]
+                pos = ans_end
+        pos += 1
+
+    full_result["labels"] = labels
+    full_result["input_ids"] = input_ids
+    return full_result
+
+
+class LazySupervisedDataset(Dataset):
+    """Dataset for supervised fine-tuning."""
+
+    def __init__(self, processor, data_args):
+        super(LazySupervisedDataset, self).__init__()
+
+        dataset = data_args.dataset_use.split(",")
+        dataset_list = data_list(dataset)
+        rank0_print(f"Loading datasets: {dataset_list}")
+        self.video_max_total_pixels = getattr(
+            data_args, "video_max_total_pixels", 1664 * 28 * 28
+        )
+        self.video_min_total_pixels = getattr(
+            data_args, "video_min_total_pixels", 256 * 28 * 28
+        )
+        self.model_type = data_args.model_type
+        if data_args.model_type == "qwen3vl":
+            self.get_rope_index = get_rope_index_3
+        elif data_args.model_type == "qwen2.5vl":
+            self.get_rope_index = get_rope_index_25
+        elif data_args.model_type == "qwen2vl":
+            self.get_rope_index = get_rope_index_2
+        else:
+            raise ValueError(f"model_type: {data_args.model_type} not supported")
+
+        list_data_dict = []
+
+        for data in dataset_list:
+            file_format = data["annotation_path"].split(".")[-1]
+            if file_format == "jsonl":
+                annotations = read_jsonl(data["annotation_path"])
+            else:
+                annotations = json.load(open(data["annotation_path"], "r"))
+            sampling_rate = data.get("sampling_rate", 1.0)
+            if sampling_rate < 1.0:
+                annotations = random.sample(
+                    annotations, int(len(annotations) * sampling_rate)
+                )
+                rank0_print(f"sampling {len(annotations)} examples from dataset {data}")
+            else:
+                rank0_print(f"dataset name: {data}")
+            for ann in annotations:
+                if isinstance(ann, list):
+                    for sub_ann in ann:
+                        sub_ann["data_path"] = data["data_path"]
+                else:
+                    ann["data_path"] = data["data_path"]
+            list_data_dict += annotations
+
+        rank0_print(f"Total training samples: {len(list_data_dict)}")
+
+
+        rank0_print("Formatting inputs...Skip in lazy mode")
+        processor = update_processor_pixels(processor, data_args)
+        self.processor = processor
+        self.tokenizer = processor.tokenizer
+        self.data_args = data_args
+        self.merge_size = getattr(processor.image_processor, "merge_size", 2)
+        self.list_data_dict = list_data_dict
+
+        if data_args.data_packing:
+            self.item_fn = self._get_packed_item
+        else:
+            self.item_fn = self._get_item
+
+    def __len__(self):
+        return len(self.list_data_dict)
+
+    @property
+    def lengths(self):
+        length_list = []
+        for sample in self.list_data_dict:
+            img_tokens = 128 if "image" in sample else 0
+            length_list.append(
+                sum(len(conv["value"].split()) for conv in sample["conversations"])
+                + img_tokens
+            )
+        return length_list
+
+    @property
+    def modality_lengths(self):
+        length_list = []
+        for sample in self.list_data_dict:
+            cur_len = sum(
+                len(conv["value"].split()) for conv in sample["conversations"]
+            )
+            cur_len = (
+                cur_len if ("image" in sample) or ("video" in sample) else -cur_len
+            )
+            length_list.append(cur_len)
+        return length_list
+
+    @property
+    def pre_calculated_length(self):
+        if "num_tokens" in self.list_data_dict[0]:
+            length_list = [sample["num_tokens"] for sample in self.list_data_dict]
+            return np.array(length_list)
+        else:
+            print("No pre-calculated length available.")
+            return np.array([1] * len(self.list_data_dict))
+
+    def __getitem__(self, i) -> Dict[str, torch.Tensor]:
+        num_base_retries = 3
+        num_final_retries = 30
+
+        # try the current sample first
+        for attempt_idx in range(num_base_retries):
+            try:
+                sources = self.list_data_dict[i]
+                if isinstance(sources, dict):
+                    sources = [sources]
+                sample = self.item_fn(sources)
+                return sample
+            except Exception as e:
+                # sleep 1s in case it is a cloud disk issue
+                print(f"[Try #{attempt_idx}] Failed to fetch sample {i}. Exception:", e)
+                time.sleep(1)
+
+        # try other samples, in case it is file corruption issue
+        for attempt_idx in range(num_base_retries):
+            try:
+                next_index = min(i + 1, len(self.list_data_dict) - 1)
+                sources = self.list_data_dict[next_index]
+                if isinstance(sources, dict):
+                    sources = [sources]
+
+                sample = self.item_fn(sources)
+                return sample
+            except Exception as e:
+                # no need to sleep
+                print(
+                    f"[Try other #{attempt_idx}] Failed to fetch sample {next_index}. Exception:",
+                    e,
+                )
+                pass
+
+        try:
+            sources = self.list_data_dict[i]
+            if isinstance(sources, dict):
+                sources = [sources]
+            sample = self.item_fn(sources)
+            return sample
+        except Exception as e:
+            raise e
+
+    def _get_item(self, sources) -> Dict[str, torch.Tensor]:
+        data_dict = preprocess_qwen_visual(
+            sources,
+            self.processor,
+        )
+
+        seq_len = data_dict["input_ids"][0].size(0)
+
+        if "image_grid_thw" in data_dict:
+            grid_thw = data_dict.get("image_grid_thw")
+            if not isinstance(grid_thw, Sequence):
+                grid_thw = [grid_thw]
+        else:
+            grid_thw = None
+
+        if "video_grid_thw" in data_dict:
+            video_grid_thw = data_dict.get("video_grid_thw")
+            if not isinstance(video_grid_thw, Sequence):
+                video_grid_thw = [video_grid_thw]
+            second_per_grid_ts = [
+                self.processor.video_processor.temporal_patch_size
+                / self.processor.video_processor.fps
+            ] * len(video_grid_thw)
+        else:
+            video_grid_thw = None
+            second_per_grid_ts = None
+
+        position_ids, _ = self.get_rope_index(
+            self.merge_size,
+            data_dict["input_ids"],
+            image_grid_thw=torch.cat(grid_thw, dim=0) if grid_thw else None,
+            video_grid_thw=(
+                torch.cat(video_grid_thw, dim=0) if video_grid_thw else None
+            ),
+            second_per_grid_ts=second_per_grid_ts if second_per_grid_ts else None,
+        )
+
+        data_dict["position_ids"] = position_ids
+        data_dict["attention_mask"] = [seq_len]
+
+        text = self.processor.tokenizer.decode(
+            data_dict["input_ids"][0], skip_special_tokens=False
+        )
+
+        labels = data_dict["labels"][0]
+        labels = [
+            tid if tid != -100 else self.processor.tokenizer.pad_token_id
+            for tid in labels
+        ]
+        label = self.processor.tokenizer.decode(labels, skip_special_tokens=False)
+
+        return data_dict
+
+    def _get_packed_item(self, sources) -> Dict[str, torch.Tensor]:
+
+        if isinstance(sources, dict):
+            if isinstance(source, dict):
+                sources = [sources]
+            assert len(sources) == 1, "Don't know why it is wrapped to a list"  # FIXME
+            return self._get_item(sources)
+
+        if isinstance(sources, list):
+            data_list = []
+            new_data_dict = {}
+            for source in sources:
+                if isinstance(source, dict):
+                    source = [source]
+                assert (
+                    len(source) == 1
+                ), f"Don't know why it is wrapped to a list.\n {source}"  # FIXME
+                data_list.append(self._get_item(source))
+
+            input_ids = torch.cat([d["input_ids"] for d in data_list], dim=1)
+            labels = torch.cat([d["labels"] for d in data_list], dim=1)
+            position_ids = torch.cat([d["position_ids"] for d in data_list], dim=2)
+            attention_mask = [
+                d["attention_mask"][0] for d in data_list if "attention_mask" in d
+            ]
+            new_data_dict = {
+                "input_ids": input_ids,
+                "labels": labels,
+                "position_ids": position_ids,
+                "attention_mask": attention_mask if attention_mask else None,
+            }
+
+            if any("pixel_values" in d for d in data_list):
+                new_data_dict.update(
+                    {
+                        "pixel_values": torch.cat(
+                            [
+                                d["pixel_values"]
+                                for d in data_list
+                                if "pixel_values" in d
+                            ],
+                            dim=0,
+                        ),
+                        "image_grid_thw": torch.cat(
+                            [
+                                d["image_grid_thw"]
+                                for d in data_list
+                                if "image_grid_thw" in d
+                            ],
+                            dim=0,
+                        ),
+                    }
+                )
+
+            if any("pixel_values_videos" in d for d in data_list):
+                new_data_dict.update(
+                    {
+                        "pixel_values_videos": torch.cat(
+                            [
+                                d["pixel_values_videos"]
+                                for d in data_list
+                                if "pixel_values_videos" in d
+                            ],
+                            dim=0,
+                        ),
+                        "video_grid_thw": torch.cat(
+                            [
+                                d["video_grid_thw"]
+                                for d in data_list
+                                if "video_grid_thw" in d
+                            ],
+                            dim=0,
+                        ),
+                    }
+                )
+            return new_data_dict
+
+
+def pad_and_cat(tensor_list):
+    max_length = max(tensor.shape[2] for tensor in tensor_list)
+
+    padded_tensors = []
+    for tensor in tensor_list:
+        pad_length = max_length - tensor.shape[2]
+        padded_tensor = torch.nn.functional.pad(tensor, (0, pad_length), "constant", 1)
+        padded_tensors.append(padded_tensor)
+
+    stacked_tensor = torch.cat(padded_tensors, dim=1)
+
+    return stacked_tensor
+
+
+@dataclass
+class DataCollatorForSupervisedDataset(object):
+    """Collate examples for supervised fine-tuning."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels, position_ids = tuple(
+            [instance[key] for instance in instances]
+            for key in ("input_ids", "labels", "position_ids")
+        )
+        input_ids = [ids.squeeze(0) for ids in input_ids]
+        labels = [ids.squeeze(0) for ids in labels]
+        input_ids = torch.nn.utils.rnn.pad_sequence(
+            input_ids, batch_first=True, padding_value=self.tokenizer.pad_token_id
+        )
+        labels = torch.nn.utils.rnn.pad_sequence(
+            labels, batch_first=True, padding_value=IGNORE_INDEX
+        )
+        position_ids = pad_and_cat(position_ids)
+        input_ids = input_ids[:, : self.tokenizer.model_max_length]
+        labels = labels[:, : self.tokenizer.model_max_length]
+        position_ids = position_ids[:, :, : self.tokenizer.model_max_length]
+        batch = dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=input_ids.ne(self.tokenizer.pad_token_id),
+        )
+        images = list(
+            instance["pixel_values"]
+            for instance in instances
+            if "pixel_values" in instance
+        )
+        videos = list(
+            instance["pixel_values_videos"]
+            for instance in instances
+            if "pixel_values_videos" in instance
+        )
+        if len(images) != 0:
+            concat_images = torch.cat([image for image in images], dim=0)
+            grid_thw = [
+                instance["image_grid_thw"]
+                for instance in instances
+                if "image_grid_thw" in instance
+            ]
+            grid_thw = torch.cat(grid_thw, dim=0)
+        else:
+            concat_images = None
+            grid_thw = None
+
+        if len(videos) != 0:
+            concat_videos = torch.cat([video for video in videos], dim=0)
+            video_grid_thw = [
+                instance["video_grid_thw"]
+                for instance in instances
+                if "video_grid_thw" in instance
+            ]
+            video_grid_thw = torch.cat(video_grid_thw, dim=0)
+        else:
+            concat_videos = None
+            video_grid_thw = None
+
+        batch["pixel_values"] = concat_images
+        batch["image_grid_thw"] = grid_thw
+        batch["pixel_values_videos"] = concat_videos
+        batch["video_grid_thw"] = video_grid_thw
+        batch["position_ids"] = position_ids
+        return batch
+
+
+@dataclass
+class FlattenedDataCollatorForSupervisedDataset(DataCollatorForSupervisedDataset):
+    """Collate examples into packed sequence with multi-modal support."""
+
+    tokenizer: transformers.PreTrainedTokenizer
+
+    def __call__(self, instances: Sequence[Dict]) -> Dict[str, torch.Tensor]:
+        input_ids, labels, position_ids, attention_mask = tuple(
+            [instance[key] for instance in instances]
+            for key in ("input_ids", "labels", "position_ids", "attention_mask")
+        )
+        attention_mask = list(
+            itertools.chain(
+                *(
+                    instance["attention_mask"]
+                    for instance in instances
+                    if "attention_mask" in instance
+                )
+            )
+        )
+        seq_lens = torch.tensor([0] + attention_mask, dtype=torch.int32)
+        cumsum_seq_lens = torch.cumsum(seq_lens, dim=0, dtype=torch.int32)
+        input_ids = torch.cat(input_ids, dim=1)
+        labels = torch.cat(labels, dim=1)
+        position_ids = torch.cat(position_ids, dim=2)
+
+        batch = dict(
+            input_ids=input_ids,
+            labels=labels,
+            attention_mask=cumsum_seq_lens,
+            position_ids=position_ids,
+        )
+        images = list(
+            instance["pixel_values"]
+            for instance in instances
+            if "pixel_values" in instance
+        )
+        videos = list(
+            instance["pixel_values_videos"]
+            for instance in instances
+            if "pixel_values_videos" in instance
+        )
+        if len(images) != 0:
+            concat_images = torch.cat([image for image in images], dim=0)
+            grid_thw = [
+                instance["image_grid_thw"]
+                for instance in instances
+                if "image_grid_thw" in instance
+            ]
+            grid_thw = torch.cat(grid_thw, dim=0)
+        else:
+            concat_images = None
+            grid_thw = None
+
+        if len(videos) != 0:
+            concat_videos = torch.cat([video for video in videos], dim=0)
+            video_grid_thw = [
+                instance["video_grid_thw"]
+                for instance in instances
+                if "video_grid_thw" in instance
+            ]
+            video_grid_thw = torch.cat(video_grid_thw, dim=0)
+        else:
+            concat_videos = None
+            video_grid_thw = None
+
+        batch["pixel_values"] = concat_images
+        batch["image_grid_thw"] = grid_thw
+        batch["pixel_values_videos"] = concat_videos
+        batch["video_grid_thw"] = video_grid_thw
+
+        return batch
+
+
+def make_supervised_data_module(processor, data_args) -> Dict:
+    """Make dataset and collator for supervised fine-tuning."""
+    train_dataset = LazySupervisedDataset(processor, data_args=data_args)
+    if data_args.data_flatten or data_args.data_packing:
+        data_collator = FlattenedDataCollatorForSupervisedDataset(processor.tokenizer)
+        return dict(
+            train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator
+        )
+    data_collator = DataCollatorForSupervisedDataset(processor.tokenizer)
+    return dict(
+        train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator
+    )
+
+
+if __name__ == "__main__":
+    pass

qwenvl/train/argument.py

@@ -0,0 +1,70 @@
+import transformers
+from dataclasses import dataclass, field
+from typing import Dict, Optional, Sequence, List
+
+
+@dataclass
+class ModelArguments:
+    model_name_or_path: Optional[str] = field(default="Qwen/Qwen2.5-VL-3B-Instruct")
+    tune_mm_llm: bool = field(default=False)
+    tune_mm_mlp: bool = field(default=False)
+    tune_mm_vision: bool = field(default=False)
+
+@dataclass
+class DataArguments:
+    dataset_use: str = field(default="")
+    data_flatten: bool = field(default=False)
+    data_packing: bool = field(default=False)
+    base_interval: int = field(default=2)
+    max_pixels: int = field(default=28 * 28 * 576)
+    min_pixels: int = field(default=28 * 28 * 16)
+    video_max_frames: Optional[int] = field(default=8)
+    video_min_frames: Optional[int] = field(default=4)
+    video_max_pixels: int = field(default=1024 * 28 * 28)
+    video_min_pixels: int = field(default=256 * 28 * 28)
+    video_fps: float = 2
+    # Aesthetic scoring
+    use_score_tokens: bool = field(default=False,
+        metadata={"help": "Enable custom score tokens for aesthetic scoring tasks."})
+    use_existing_vocab_tokens: bool = field(default=False,
+        metadata={"help": "Use existing vocab tokens (aa, ab, ...) as score bins instead of adding new tokens."})
+    use_placeholder: bool = field(default=True,
+        metadata={"help": "Use placeholder tokens for score regression. If False, use direct score bin tokens."})
+    crit_score_bins: int = field(default=20,
+        metadata={"help": "Number of bins for criteria scores (1-10). Tokens: <score1>..<scoreN>."})
+    total_score_bins: int = field(default=101,
+        metadata={"help": "Number of bins for total score (0-100). Tokens: <total1>..<totalN>."})
+
+
+@dataclass
+class TrainingArguments(transformers.TrainingArguments):
+    cache_dir: Optional[str] = field(default=None)
+    optim: str = field(default="adamw_torch")
+    model_max_length: int = field(
+        default=512,
+        metadata={
+            "help": "Maximum sequence length. Sequences will be right padded (and possibly truncated)."
+        },
+    )
+    mm_projector_lr: Optional[float] = None
+    vision_tower_lr: Optional[float] = None
+
+    ## Lora config
+    lora_enable: bool = field(default=False)
+    lora_r: int = field(default=64)
+    lora_alpha: int = field(default=128)
+    lora_dropout: float = field(default=0.0)
+
+    ## Local logging
+    sample_every_steps: int = field(default=500,
+        metadata={"help": "Run model generation samples every N steps and save to output_dir/model_output_samples.txt."})
+
+    ## Score loss config
+    score_ce_weight: float = field(default=1.0,
+        metadata={"help": "Weight for CE loss in aesthetic scoring."})
+    crit_reg_weight: float = field(default=0.5,
+        metadata={"help": "Weight for criteria (1-10) regression loss."})
+    total_reg_weight: float = field(default=0.5,
+        metadata={"help": "Weight for total-score (0-100) regression loss."})
+    score_regression_loss: str = field(default="huber",
+        metadata={"help": "Regression loss type: 'mse' or 'huber'."})

qwenvl/train/logging_callback.py

@@ -0,0 +1,141 @@
+"""
+Local training logger callback.
+
+Writes two files to output_dir:
+  - train_metrics.jsonl : one JSON line per logging step with loss / grad_norm / lr / step
+  - model_output_samples.txt : periodic model generation samples showing <think> + score token
+"""
+
+import csv
+import json
+import os
+import time
+from pathlib import Path
+from typing import Optional
+
+import torch
+from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
+
+
+class LocalLoggingCallback(TrainerCallback):
+    """
+    Args:
+        processor:       AutoProcessor (for generation + decoding).
+        sample_dataset:  A small list of dicts with keys image_path / question / gt_score,
+                         used for periodic generation samples.  Pass None to skip.
+        sample_every_steps: How often (in steps) to run generation samples.
+        output_dir:      Where to write log files (defaults to training_args.output_dir).
+    """
+
+    def __init__(
+        self,
+        processor,
+        sample_dataset=None,
+        sample_every_steps: int = 500,
+        output_dir: Optional[str] = None,
+    ):
+        self.processor = processor
+        self.sample_dataset = sample_dataset or []
+        self.sample_every_steps = sample_every_steps
+        self._output_dir = output_dir
+        self._metrics_path = None
+        self._samples_path = None
+        self._csv_writer = None
+        self._csv_file = None
+
+    # ── setup ─────────────────────────────────────────────────────────────────
+
+    def on_train_begin(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs):
+        out = Path(self._output_dir or args.output_dir)
+        out.mkdir(parents=True, exist_ok=True)
+        self._metrics_path = out / "train_metrics.jsonl"
+        self._samples_path = out / "model_output_samples.txt"
+        if args.local_rank in (-1, 0):
+            with open(self._samples_path, "a") as f:
+                f.write(f"\n{'='*60}\nTraining started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n{'='*60}\n")
+
+    # ── metrics logging ───────────────────────────────────────────────────────
+
+    def on_log(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, logs=None, **kwargs):
+        if args.local_rank not in (-1, 0):
+            return
+        if logs is None:
+            return
+
+        record = {
+            "step":      state.global_step,
+            "epoch":     round(state.epoch or 0, 4),
+            "loss":      logs.get("loss"),
+            "grad_norm": logs.get("grad_norm"),
+            "lr":        logs.get("learning_rate"),
+            "wall_time": time.time(),
+        }
+        record = {k: v for k, v in record.items() if v is not None}
+
+        with open(self._metrics_path, "a") as f:
+            f.write(json.dumps(record) + "\n")
+
+    # ── generation samples ────────────────────────────────────────────────────
+
+    def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, model=None, **kwargs):
+        if args.local_rank not in (-1, 0):
+            return
+        if not self.sample_dataset:
+            return
+        if state.global_step % self.sample_every_steps != 0:
+            return
+        self._run_samples(model, state.global_step, args)
+
+    def on_train_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, model=None, **kwargs):
+        if args.local_rank not in (-1, 0):
+            return
+        if self.sample_dataset and model is not None:
+            self._run_samples(model, state.global_step, args, label="[FINAL]")
+
+    def _run_samples(self, model, step: int, args: TrainingArguments, label: str = ""):
+        model.eval()
+        lines = [f"\n--- Step {step} {label} {time.strftime('%H:%M:%S')} ---\n"]
+
+        with torch.no_grad():
+            for item in self.sample_dataset:
+                image_path = item["image_path"]
+                question   = item["question"]
+                gt_score   = item.get("gt_score", "?")
+
+                messages = [
+                    {
+                        "role": "user",
+                        "content": [
+                            {"type": "image", "image": image_path},
+                            {"type": "text",  "text": question},
+                        ],
+                    }
+                ]
+                try:
+                    inputs = self.processor.apply_chat_template(
+                        messages, tokenize=True, return_dict=True,
+                        return_tensors="pt", add_generation_prompt=True,
+                    )
+                    inputs = {k: v.to(model.device) if hasattr(v, "to") else v
+                              for k, v in inputs.items()}
+
+                    out_ids = model.generate(
+                        **inputs,
+                        max_new_tokens=300,
+                        do_sample=False,
+                    )
+                    # decode only the generated tokens
+                    generated = out_ids[0][inputs["input_ids"].shape[1]:]
+                    text = self.processor.tokenizer.decode(generated, skip_special_tokens=False)
+
+                    lines.append(f"[Q] {question[:80]}")
+                    lines.append(f"[GT score] {gt_score}")
+                    lines.append(f"[Model output]\n{text}\n")
+
+                except Exception as e:
+                    lines.append(f"[ERROR] {e}\n")
+
+        with open(self._samples_path, "a") as f:
+            f.write("\n".join(lines) + "\n")
+
+        model.train()

qwenvl/train/score_loss.py

@@ -0,0 +1,247 @@
+"""
+Score loss for aesthetic scoring tasks.
+
+Total loss = ce_weight * CE_loss  +  reg_weight * regression_loss
+
+The regression branch is selected by `sample_type`:
+  - sample_type == 1  ("crit"):  use crit_reg_weight  + criteria  token ids/values
+  - sample_type == 0  ("total"): use total_reg_weight + total     token ids/values
+
+Because the dataset produces one criterion or one total-score question per sample,
+a batch is homogeneous in type and we route once per forward pass — no scanning
+of both heads.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import List, Optional
+
+from .score_tokens import ScoreTokenRegistry
+
+IGNORE_INDEX = -100
+
+SAMPLE_TYPE_TOTAL = 0
+SAMPLE_TYPE_CRIT  = 1
+
+
+class ScoreLoss(nn.Module):
+    """
+    Args:
+        crit_registry:    ScoreTokenRegistry for criteria scores (1–10).
+        total_registry:   ScoreTokenRegistry for total score (0–100).
+        tokenizer:        HF tokenizer with the custom score tokens already added.
+        ce_weight:        Weight for the CE (language-modelling) loss.
+        crit_reg_weight:  Weight for the criteria regression loss.
+        total_reg_weight: Weight for the total-score regression loss.
+        regression_loss:  "mse" or "huber".
+    """
+
+    def __init__(
+        self,
+        crit_registry: ScoreTokenRegistry,
+        total_registry: ScoreTokenRegistry,
+        tokenizer,
+        ce_weight: float = 1.0,
+        crit_reg_weight: float = 0.5,
+        total_reg_weight: float = 0.5,
+        regression_loss: str = "mse",
+    ):
+        super().__init__()
+        self.ce_weight = ce_weight
+        self.crit_reg_weight = crit_reg_weight
+        self.total_reg_weight = total_reg_weight
+        self.regression_loss_type = regression_loss
+
+        # ── criteria head ─────────────────────────────────────────────────────
+        self.crit_ids = crit_registry.get_token_ids(tokenizer)
+        self.register_buffer("crit_vals", crit_registry.get_value_tensor())
+        self.crit_range = crit_registry.score_max - crit_registry.score_min
+        self.crit_use_placeholder = crit_registry.use_placeholder
+        self.crit_ph_id = crit_registry.get_placeholder_id(tokenizer) if crit_registry.use_placeholder else None
+
+        # ── total head ────────────────────────────────────────────────────────
+        self.total_ids = total_registry.get_token_ids(tokenizer)
+        self.register_buffer("total_vals", total_registry.get_value_tensor())
+        self.total_range = total_registry.score_max - total_registry.score_min
+        self.total_use_placeholder = total_registry.use_placeholder
+        self.total_ph_id = total_registry.get_placeholder_id(tokenizer) if total_registry.use_placeholder else None
+
+    # ── helpers ───────────────────────────────────────────────────────────────
+
+    def _reg_loss(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        # if self.regression_loss_type == "huber":
+        #     return F.huber_loss(pred, target, delta=1.0)
+        return F.mse_loss(pred, target)
+
+    def _regression_for_head(
+        self,
+        flat_logits: torch.Tensor,   # (B*(T-1), V)
+        flat_labels: torch.Tensor,   # (B*(T-1),)
+        token_ids: List[int],
+        vals: torch.Tensor,          # (n_bins,)
+        score_range: float,
+        gt_scores: Optional[torch.Tensor] = None,  # (B,) continuous gt scores
+        placeholder_id: Optional[int] = None,  # placeholder token id
+        batch_size: int = None,      # needed for non-placeholder mode
+    ) -> Optional[torch.Tensor]:
+        """Compute regression loss at score positions.
+
+        If placeholder_id is set, find positions where label == placeholder_id,
+        then use logits over token_ids at those positions to compute expected score.
+        Otherwise, find the LAST score token position per sample to avoid matching
+        incidental occurrences of existing-vocab tokens in reasoning text.
+        """
+        if placeholder_id is not None:
+            mask = (flat_labels == placeholder_id)
+            if not mask.any():
+                return None
+            if gt_scores is None:
+                return None
+            selected_logits = flat_logits[mask]
+        else:
+            # Existing vocab tokens (aa, ab..) may appear in reasoning text.
+            # Find the LAST matching position per sample to get the intended score token.
+            ids_t = torch.tensor(token_ids, device=flat_labels.device)
+            token_match = (flat_labels.unsqueeze(1) == ids_t.unsqueeze(0)).any(dim=1)
+
+            if not token_match.any():
+                return None
+            if gt_scores is None:
+                return None
+
+            if batch_size is not None and batch_size > 0:
+                seq_len = flat_labels.shape[0] // batch_size
+                labels_2d = flat_labels.view(batch_size, seq_len)
+                match_2d = token_match.view(batch_size, seq_len)
+
+                # For each sample, find the last matching position
+                selected_indices = []
+                for b in range(batch_size):
+                    positions = match_2d[b].nonzero(as_tuple=False).squeeze(-1)
+                    if positions.numel() > 0:
+                        last_pos = positions[-1].item()
+                        selected_indices.append(b * seq_len + last_pos)
+
+                if not selected_indices:
+                    return None
+
+                idx_t = torch.tensor(selected_indices, device=flat_logits.device)
+                selected_logits = flat_logits[idx_t]
+                # Align gt_scores: only keep samples that had a match
+                gt_scores = gt_scores[:len(selected_indices)]
+            else:
+                # Fallback: use all matches (may cause size mismatch warning)
+                selected_logits = flat_logits[token_match]
+
+        ids_t = torch.tensor(token_ids, device=flat_labels.device)
+        with torch.amp.autocast('cuda', enabled=False):
+            logits_fp32 = selected_logits.float()
+            head_logits = logits_fp32[:, ids_t]               # (M, n_bins) fp32
+            head_logits = torch.clamp(head_logits, min=-50.0, max=50.0)
+            probs       = torch.softmax(head_logits, dim=-1)
+            pred_scores = (probs * vals.float()).sum(dim=-1)   # (M,) fp32
+            target_scores = gt_scores.float().to(pred_scores.device)
+            reg = self._reg_loss(pred_scores / score_range, target_scores / score_range)
+        return reg
+
+    # ── forward ───────────────────────────────────────────────────────────────
+
+    def forward(
+        self,
+        logits: torch.Tensor,           # (B, T, V)
+        labels: torch.Tensor,           # (B, T)
+        sample_type: Optional[torch.Tensor] = None,  # (B,) int, 1=crit 0=total 2=mcq
+        gt_score: Optional[torch.Tensor] = None,     # (B,) float, original continuous score
+        vocab_size: int = None,
+    ) -> torch.Tensor:
+        B, T, V = logits.shape
+
+        shift_logits = logits[:, :-1].contiguous()   # (B, T-1, V)
+        shift_labels = labels[:, 1:].contiguous()    # (B, T-1)
+
+        flat_logits = shift_logits.view(-1, V)
+        flat_labels = shift_labels.view(-1)
+
+        # ── CE loss ───────────────────────────────────────────────────────────
+        if torch.isnan(flat_logits).any() or torch.isinf(flat_logits).any():
+            print(f"[ScoreLoss] WARNING: logits contain nan/inf BEFORE ce_loss, "
+                  f"nan_count={torch.isnan(flat_logits).sum()}, inf_count={torch.isinf(flat_logits).sum()}")
+            return torch.tensor(0.0, device=logits.device, requires_grad=True)
+
+        ce_loss = F.cross_entropy(flat_logits, flat_labels, ignore_index=IGNORE_INDEX)
+        if torch.isnan(ce_loss):
+            ce_loss = torch.tensor(0.0, device=logits.device)
+        total_loss = self.ce_weight * ce_loss
+
+        # ── Regression loss: route by sample_type ────────────────────────────
+        if sample_type is not None:
+            has_crit  = (sample_type == SAMPLE_TYPE_CRIT).any()
+            has_total = (sample_type == SAMPLE_TYPE_TOTAL).any()
+        else:
+            has_crit  = True
+            has_total = True
+
+        seq_len = T - 1  # after shift
+
+        if has_crit and self.crit_reg_weight > 0.0:
+            crit_gt = None
+            if gt_score is not None and sample_type is not None:
+                crit_mask = (sample_type == SAMPLE_TYPE_CRIT)
+                if crit_mask.any():
+                    crit_gt = gt_score[crit_mask]
+                    crit_logits = shift_logits[crit_mask].view(-1, V)
+                    crit_labels = shift_labels[crit_mask].view(-1)
+                    crit_B = crit_mask.sum().item()
+                else:
+                    crit_gt = None
+            elif gt_score is not None:
+                crit_gt = gt_score
+                crit_logits = flat_logits
+                crit_labels = flat_labels
+                crit_B = B
+            else:
+                crit_gt = None
+
+            if crit_gt is not None:
+                reg = self._regression_for_head(
+                    crit_logits, crit_labels,
+                    self.crit_ids, self.crit_vals.to(logits.device),
+                    self.crit_range, gt_scores=crit_gt,
+                    placeholder_id=self.crit_ph_id,
+                    batch_size=crit_B,
+                )
+                if reg is not None:
+                    total_loss = total_loss + self.crit_reg_weight * reg
+
+        if has_total and self.total_reg_weight > 0.0:
+            total_gt = None
+            if gt_score is not None and sample_type is not None:
+                total_mask = (sample_type == SAMPLE_TYPE_TOTAL)
+                if total_mask.any():
+                    total_gt = gt_score[total_mask]
+                    total_logits = shift_logits[total_mask].view(-1, V)
+                    total_labels = shift_labels[total_mask].view(-1)
+                    total_B = total_mask.sum().item()
+                else:
+                    total_gt = None
+            elif gt_score is not None:
+                total_gt = gt_score
+                total_logits = flat_logits
+                total_labels = flat_labels
+                total_B = B
+            else:
+                total_gt = None
+
+            if total_gt is not None:
+                reg = self._regression_for_head(
+                    total_logits, total_labels,
+                    self.total_ids, self.total_vals.to(logits.device),
+                    self.total_range, gt_scores=total_gt,
+                    placeholder_id=self.total_ph_id,
+                    batch_size=total_B,
+                )
+                if reg is not None:
+                    total_loss = total_loss + self.total_reg_weight * reg
+
+        return total_loss

qwenvl/train/score_tokens.py

@@ -0,0 +1,222 @@
+"""
+Generic integer-index score tokens.
+
+A ScoreTokenRegistry holds N tokens <score1> ... <scoreN> that represent
+N evenly-spaced bins covering [score_min, score_max].
+
+The actual score value for bin k  (1-indexed) is:
+    score_min + (k - 1) / (N - 1) * (score_max - score_min)
+
+Usage
+-----
+    # criteria: 1–10 range, 20 bins  →  <score1> … <score20>
+    crit_reg = ScoreTokenRegistry(n_bins=20, score_min=1.0, score_max=10.0)
+
+    # total:    0–100 range, 101 bins →  <total1> … <total101>
+    total_reg = ScoreTokenRegistry(n_bins=101, score_min=0.0, score_max=100.0,
+                                   prefix="total")
+
+    all_tokens = crit_reg.tokens + total_reg.tokens
+"""
+
+import string
+import torch
+from typing import List, Optional
+
+
+def get_single_token_two_letter_combos(vocab: dict) -> List[str]:
+    """Return ordered list of two-letter combos that are single tokens in vocab."""
+    combos = []
+    for c1 in string.ascii_lowercase:
+        for c2 in string.ascii_lowercase:
+            t = c1 + c2
+            if t in vocab:
+                combos.append(t)
+    return combos
+
+
+class ScoreTokenRegistry:
+    """
+    Registry for a set of evenly-spaced score tokens.
+
+    Args:
+        n_bins:     Number of discrete bins / tokens.
+        score_min:  Real value corresponding to bin 1.
+        score_max:  Real value corresponding to bin n_bins.
+        prefix:     Token prefix string (default "score").
+                    Tokens are named <{prefix}1> … <{prefix}{n_bins}>.
+        use_placeholder: If True, add a placeholder token <{prefix}_ph> that
+                    appears in text instead of specific score tokens.
+    """
+
+    def __init__(
+        self,
+        n_bins: int,
+        score_min: float,
+        score_max: float,
+        prefix: str = "score",
+        use_placeholder: bool = False,
+    ):
+        assert n_bins >= 2
+        self.n_bins = n_bins
+        self.score_min = score_min
+        self.score_max = score_max
+        self.prefix = prefix
+        self.use_placeholder = use_placeholder
+
+        self.tokens: list[str] = [f"<{prefix}{k}>" for k in range(1, n_bins + 1)]
+        self.placeholder_token: str = f"<{prefix}_ph>"
+
+        # Real value for each bin (linearly spaced)
+        self.values: list[float] = [
+            score_min + (k / (n_bins - 1)) * (score_max - score_min)
+            for k in range(n_bins)
+        ]
+
+        # Pre-computed constants for score_to_token
+        self._scale = (n_bins - 1) / (score_max - score_min)
+
+        # O(1) lookup tables
+        self._token_to_idx: dict[str, int] = {t: i for i, t in enumerate(self.tokens)}
+
+    @property
+    def all_tokens(self) -> list[str]:
+        """All tokens to register (score bins + placeholder if enabled)."""
+        if self.use_placeholder:
+            return self.tokens + [self.placeholder_token]
+        return self.tokens
+
+    def get_placeholder_id(self, tokenizer) -> int:
+        """Return the token id of the placeholder token."""
+        return tokenizer.convert_tokens_to_ids(self.placeholder_token)
+
+    # ── conversion helpers ────────────────────────────────────────────────────
+
+    def score_to_token(self, score: float) -> str:
+        """Map a real score to the nearest bin token (or placeholder if enabled)."""
+        if self.use_placeholder:
+            return self.placeholder_token
+        k = int(round((score - self.score_min) * self._scale))
+        k = max(0, min(self.n_bins - 1, k))
+        return self.tokens[k]
+
+    def token_to_score(self, token: str) -> float:
+        return self.values[self._token_to_idx[token]]
+
+    # ── tensor helpers (used at inference / loss time) ────────────────────────
+
+    def get_token_ids(self, tokenizer) -> list[int]:
+        """Return token ids in token order (token1, token2, …, tokenN)."""
+        return [tokenizer.convert_tokens_to_ids(t) for t in self.tokens]
+
+    def get_value_tensor(
+        self, tokenizer=None, device: Optional[torch.device] = None
+    ) -> torch.Tensor:
+        """Return float tensor of real values aligned with token order."""
+        t = torch.tensor(self.values, dtype=torch.float32)
+        return t.to(device) if device else t
+
+    def expected_score(
+        self,
+        logits: torch.Tensor,
+        token_ids: list[int],
+        value_tensor: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Compute expected score from a logits slice.
+
+        Args:
+            logits:       (vocab_size,) raw logits for one position
+            token_ids:    ordered list of this registry's token ids
+            value_tensor: (n_bins,) real values matching token_ids order
+
+        Returns:
+            scalar expected score
+        """
+        ids_t = torch.tensor(token_ids, device=logits.device)
+        probs = torch.softmax(logits[ids_t], dim=-1)
+        return (probs * value_tensor.to(logits.device)).sum()
+
+
+class ExistingVocabScoreTokenRegistry:
+    """
+    Registry that reuses existing vocabulary tokens (two-letter combos)
+    as score bins. Only the placeholder token (if enabled) is newly added.
+
+    Args:
+        token_list:  Pre-selected tokens from the existing vocabulary.
+        score_min:   Real value corresponding to the first token.
+        score_max:   Real value corresponding to the last token.
+        prefix:      Used only for the placeholder token name <{prefix}_ph>.
+        use_placeholder: If True, add a placeholder token <{prefix}_ph>.
+    """
+
+    def __init__(
+        self,
+        token_list: List[str],
+        score_min: float,
+        score_max: float,
+        prefix: str = "score",
+        use_placeholder: bool = False,
+    ):
+        assert len(token_list) >= 2
+        self.n_bins = len(token_list)
+        self.score_min = score_min
+        self.score_max = score_max
+        self.prefix = prefix
+        self.use_placeholder = use_placeholder
+
+        self.tokens: List[str] = list(token_list)
+        self.placeholder_token: str = f"<{prefix}_ph>"
+
+        self.values: List[float] = [
+            score_min + (k / (self.n_bins - 1)) * (score_max - score_min)
+            for k in range(self.n_bins)
+        ]
+
+        self._scale = (self.n_bins - 1) / (score_max - score_min)
+        self._token_to_idx: dict[str, int] = {t: i for i, t in enumerate(self.tokens)}
+
+    @property
+    def all_tokens(self) -> List[str]:
+        """Only the placeholder is a new token; bin tokens already exist."""
+        if self.use_placeholder:
+            return [self.placeholder_token]
+        return []
+
+    @property
+    def new_tokens(self) -> List[str]:
+        """Tokens that need to be added to the tokenizer."""
+        return self.all_tokens
+
+    def get_placeholder_id(self, tokenizer) -> int:
+        return tokenizer.convert_tokens_to_ids(self.placeholder_token)
+
+    def score_to_token(self, score: float) -> str:
+        if self.use_placeholder:
+            return self.placeholder_token
+        k = int(round((score - self.score_min) * self._scale))
+        k = max(0, min(self.n_bins - 1, k))
+        return self.tokens[k]
+
+    def token_to_score(self, token: str) -> float:
+        return self.values[self._token_to_idx[token]]
+
+    def get_token_ids(self, tokenizer) -> List[int]:
+        return [tokenizer.convert_tokens_to_ids(t) for t in self.tokens]
+
+    def get_value_tensor(
+        self, tokenizer=None, device: Optional[torch.device] = None
+    ) -> torch.Tensor:
+        t = torch.tensor(self.values, dtype=torch.float32)
+        return t.to(device) if device else t
+
+    def expected_score(
+        self,
+        logits: torch.Tensor,
+        token_ids: List[int],
+        value_tensor: torch.Tensor,
+    ) -> torch.Tensor:
+        ids_t = torch.tensor(token_ids, device=logits.device)
+        probs = torch.softmax(logits[ids_t], dim=-1)
+        return (probs * value_tensor.to(logits.device)).sum()

qwenvl/train/train_qwen.py

@@ -0,0 +1,496 @@
+# Adopted from https://github.com/lm-sys/FastChat. Below is the original copyright:
+# Adopted from tatsu-lab@stanford_alpaca. Below is the original copyright:
+#    Copyright 2023 Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li
+#
+#    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.
+
+import os
+import logging
+import pathlib
+import torch
+import transformers
+import sys
+from pathlib import Path
+from PIL import Image
+Image.MAX_IMAGE_PIXELS = None  # disable decompression bomb check for large training images
+
+project_root = Path(__file__).parent.parent.parent
+sys.path.append(str(project_root))
+
+from trainer import replace_qwen2_vl_attention_class
+
+from transformers import (
+    Qwen2VLForConditionalGeneration,
+    Qwen2_5_VLForConditionalGeneration,
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLMoeForConditionalGeneration
+)
+from qwenvl.data.data_processor import make_supervised_data_module
+from qwenvl.data.aesthetic_dataset import AestheticScoringDataset
+from qwenvl.train.argument import (
+    ModelArguments,
+    DataArguments,
+    TrainingArguments,
+)
+from qwenvl.train.score_tokens import ScoreTokenRegistry, ExistingVocabScoreTokenRegistry, get_single_token_two_letter_combos
+from qwenvl.train.score_loss import ScoreLoss
+from qwenvl.train.logging_callback import LocalLoggingCallback
+from transformers import AutoProcessor, Trainer
+from deepspeed import zero
+
+local_rank = None
+
+
+def rank0_print(*args):
+    if local_rank == 0:
+        print(*args)
+
+
+def expand_vocab_for_scoring(model, tokenizer, registries):
+    """Add custom score tokens to tokenizer and resize model embeddings.
+
+    Score tokens are initialized along a bad→good semantic axis;
+    placeholder tokens get the mean embedding.
+    """
+    all_tokens = [t for reg in registries for t in reg.all_tokens]
+    num_added = tokenizer.add_tokens(all_tokens)
+    rank0_print(f"Added {num_added} score tokens to tokenizer (total vocab: {len(tokenizer)})")
+    model.resize_token_embeddings(len(tokenizer))
+
+    bad_words = ["bad", "terrible", "poor", "worst", "awful", "low", "weak"]
+    good_words = ["good", "excellent", "great", "best", "amazing", "high", "strong"]
+
+    with torch.no_grad():
+        w_in = model.get_input_embeddings().weight.data
+        w_out = model.get_output_embeddings().weight.data
+        base_vector = w_in[:-num_added].mean(dim=0, keepdim=True)
+
+        # Compute quality direction from semantic anchors
+        bad_ids = [tokenizer.convert_tokens_to_ids(w) for w in bad_words]
+        good_ids = [tokenizer.convert_tokens_to_ids(w) for w in good_words]
+        bad_vecs = torch.stack([w_in[i] for i in bad_ids if i != tokenizer.unk_token_id])
+        good_vecs = torch.stack([w_in[i] for i in good_ids if i != tokenizer.unk_token_id])
+
+        if len(bad_vecs) > 0 and len(good_vecs) > 0:
+            quality_dir = good_vecs.mean(0) - bad_vecs.mean(0)
+            q_norm = torch.norm(quality_dir)
+            if q_norm > 0:
+                quality_dir = quality_dir / q_norm
+            quality_scale = w_in[:-num_added].std() * 0.5
+
+            offset = 0
+            for reg in registries:
+                n_bins = reg.n_bins
+                for i in range(n_bins):
+                    weight = 2.0 * i / (n_bins - 1) - 1.0 if n_bins > 1 else 0.0
+                    vec = base_vector + quality_dir * weight * quality_scale
+                    w_in[-num_added + offset + i] = vec
+                    w_out[-num_added + offset + i] = vec
+                offset += n_bins
+                if reg.use_placeholder:
+                    w_in[-num_added + offset] = base_vector
+                    w_out[-num_added + offset] = base_vector
+                    offset += 1
+
+            rank0_print(f"[init] Score tokens initialized with quality axis (bad→good), scale={quality_scale:.4f}")
+        else:
+            w_in[-num_added:] = base_vector
+            w_out[-num_added:] = base_vector
+            rank0_print("[init] Score tokens initialized with mean embedding (fallback)")
+
+    return model, tokenizer
+
+
+def expand_vocab_for_existing_tokens(model, tokenizer, registries):
+    """Add only placeholder tokens when using existing vocab tokens as score bins.
+
+    Bin tokens (aa, ab, ...) already exist in the vocabulary — no embedding init needed.
+    Only placeholder tokens (if enabled) are newly added and initialized with mean embedding.
+    """
+    new_tokens = [t for reg in registries for t in reg.new_tokens]
+    if new_tokens:
+        num_added = tokenizer.add_tokens(new_tokens)
+        rank0_print(f"[existing-vocab] Added {num_added} placeholder tokens (total vocab: {len(tokenizer)})")
+        model.resize_token_embeddings(len(tokenizer))
+
+        with torch.no_grad():
+            w_in = model.get_input_embeddings().weight.data
+            w_out = model.get_output_embeddings().weight.data
+            base_vector = w_in[:-num_added].mean(dim=0, keepdim=True)
+            w_in[-num_added:] = base_vector
+            w_out[-num_added:] = base_vector
+            rank0_print(f"[existing-vocab] Placeholder tokens initialized with mean embedding")
+    else:
+        rank0_print("[existing-vocab] No placeholder — using direct score bin tokens only")
+
+    for reg in registries:
+        bin_ids = reg.get_token_ids(tokenizer)
+        rank0_print(f"[existing-vocab] {reg.prefix}: {reg.n_bins} bins using tokens "
+                     f"{reg.tokens[0]}..{reg.tokens[-1]} (ids {bin_ids[0]}..{bin_ids[-1]})")
+
+    return model, tokenizer
+
+
+class AestheticDataCollator:
+    """Wraps DataCollatorForSupervisedDataset and additionally collects sample_type."""
+
+    def __init__(self, base_collator):
+        self.base = base_collator
+
+    def __call__(self, instances):
+        # collect sample_type and gt_score before passing to base collator (which ignores unknown keys)
+        sample_types = [inst.pop("sample_type", None) for inst in instances]
+        gt_scores = [inst.pop("gt_score", None) for inst in instances]
+        batch = self.base(instances)
+        if any(t is not None for t in sample_types):
+            batch["sample_type"] = torch.tensor(
+                [t if t is not None else -1 for t in sample_types],
+                dtype=torch.long,
+            )
+        if any(s is not None for s in gt_scores):
+            batch["gt_score"] = torch.tensor(
+                [s if s is not None else 0.0 for s in gt_scores],
+                dtype=torch.float32,
+            )
+        return batch
+
+
+class AestheticTrainer(Trainer):
+    """Trainer subclass that uses ScoreLoss when score tokens are enabled."""
+
+    def __init__(self, *args, score_loss: ScoreLoss = None, tokenizer=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.score_loss = score_loss
+        self._decode_tokenizer = tokenizer  # kept separately for decoding
+
+    def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
+        if self.score_loss is None:
+            return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs)
+
+        labels      = inputs.pop("labels", None)
+        sample_type = inputs.pop("sample_type", None)   # (B,) int tensor or None
+        gt_score    = inputs.pop("gt_score", None)      # (B,) float tensor or None
+        if gt_score is not None:
+            gt_score = gt_score.float()
+        outputs = model(**inputs)
+        logits  = outputs.logits
+
+        loss = self.score_loss(logits, labels, sample_type=sample_type, gt_score=gt_score)
+
+        # ── per-batch text monitor (rank-0 only, every N steps) ───────────────
+        step = self.state.global_step
+        if (step != getattr(self, '_last_print_step', -1)
+                and self.args.local_rank in (-1, 0)
+                and self._decode_tokenizer is not None):
+            self._last_print_step = step
+            tok = self._decode_tokenizer
+            input_text = tok.decode(
+                inputs["input_ids"][0], skip_special_tokens=False
+            )
+            pred_ids = logits[0].argmax(dim=-1)
+            label_row = labels[0]
+            mask = label_row != -100
+            if mask.any():
+                gt_ids   = label_row[mask]
+                pred_at_label_pos = pred_ids[mask.roll(-1)][:-1] if mask.sum() > 1 else pred_ids[mask]
+                gt_text   = tok.decode(gt_ids,   skip_special_tokens=False)
+                pred_text = tok.decode(pred_at_label_pos, skip_special_tokens=False)
+                if torch.distributed.get_rank() == 0:
+                    print(f"\n[step {step}] INPUT : {input_text[-200:]}")
+                    print(f"[step {step}] GT    : {gt_text[-300:]}")
+                    print(f"[step {step}] PRED  : {pred_text[-300:]}\n", flush=True)
+
+        return (loss, outputs) if return_outputs else loss
+
+
+
+def safe_save_model_for_hf_trainer(trainer: transformers.Trainer, output_dir: str):
+    """Collects the state dict and dump to disk."""
+
+    if trainer.deepspeed:
+        torch.cuda.synchronize()
+        trainer.save_model(output_dir)
+        return
+
+    state_dict = trainer.model.state_dict()
+    if trainer.args.should_save:
+        cpu_state_dict = {key: value.cpu() for key, value in state_dict.items()}
+        del state_dict
+        trainer._save(output_dir, state_dict=cpu_state_dict)  # noqa
+
+
+def set_model(model_args, model):
+    if model_args.tune_mm_vision:
+        for n, p in model.visual.named_parameters():
+            p.requires_grad = True
+    else:
+        for n, p in model.visual.named_parameters():
+            p.requires_grad = False
+
+    if model_args.tune_mm_mlp:
+        for n, p in model.visual.merger.named_parameters():
+            p.requires_grad = True
+    else:
+        for n, p in model.visual.merger.named_parameters():
+            p.requires_grad = False
+
+    if model_args.tune_mm_llm:
+        for n, p in model.language_model.named_parameters():
+            p.requires_grad = True
+        model.lm_head.requires_grad = True
+    else:
+        for n, p in model.language_model.named_parameters():
+            p.requires_grad = False
+        model.lm_head.requires_grad = False
+
+
+def train(attn_implementation="flash_attention_2"):
+    global local_rank
+
+    parser = transformers.HfArgumentParser(
+        (ModelArguments, DataArguments, TrainingArguments)
+    )
+    model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    local_rank = training_args.local_rank
+    os.makedirs(training_args.output_dir, exist_ok=True)
+
+    # if "qwen3" in model_args.model_name_or_path.lower() and "a" in Path(model_args.model_name_or_path.rstrip("/")).name.lower():
+    #     model = Qwen3VLMoeForConditionalGeneration.from_pretrained(
+    #         model_args.model_name_or_path,
+    #         cache_dir=training_args.cache_dir,
+    #         attn_implementation=attn_implementation,
+    #         dtype=(torch.bfloat16 if training_args.bf16 else None),
+    #     )
+    #     data_args.model_type = "qwen3vl"
+    if "qwen3" in model_args.model_name_or_path.lower():
+        model = Qwen3VLForConditionalGeneration.from_pretrained(
+            model_args.model_name_or_path,
+            cache_dir=training_args.cache_dir,
+            attn_implementation=attn_implementation,
+            dtype=(torch.bfloat16 if training_args.bf16 else None),
+        )
+        data_args.model_type = "qwen3vl"
+    elif "qwen2.5" in model_args.model_name_or_path.lower():
+        model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            model_args.model_name_or_path,
+            cache_dir=training_args.cache_dir,
+            attn_implementation=attn_implementation,
+            dtype=(torch.bfloat16 if training_args.bf16 else None),
+        )
+        data_args.model_type = "qwen2.5vl"
+    else:
+        model = Qwen3VLForConditionalGeneration.from_pretrained(
+            model_args.model_name_or_path,
+            cache_dir=training_args.cache_dir,
+            attn_implementation=attn_implementation,
+            dtype=(torch.bfloat16 if training_args.bf16 else None),
+        )
+        data_args.model_type = "qwen3vl"
+
+    print(f'the initlized model is {model_args.model_name_or_path} the class is {model.__class__.__name__}')
+    print(f'attn_implementation: {model.config._attn_implementation}')
+    processor = AutoProcessor.from_pretrained(
+        model_args.model_name_or_path,
+    )
+
+    # Expand vocabulary with score tokens
+    score_loss_fn = None
+    crit_registry = None
+    total_registry = None
+    if data_args.use_score_tokens:
+        if data_args.use_existing_vocab_tokens:
+            import json as _json_vocab
+            vocab_path = os.path.join(model_args.model_name_or_path, "vocab.json")
+            with open(vocab_path, "r") as f:
+                vocab = _json_vocab.load(f)
+            valid_combos = get_single_token_two_letter_combos(vocab)
+            n_crit = data_args.crit_score_bins
+            n_total = data_args.total_score_bins
+            assert len(valid_combos) >= n_crit + n_total, (
+                f"Need {n_crit + n_total} single-token combos but only {len(valid_combos)} available"
+            )
+            crit_tokens = valid_combos[:n_crit]
+            total_tokens = valid_combos[n_crit:n_crit + n_total]
+
+            crit_registry = ExistingVocabScoreTokenRegistry(
+                token_list=crit_tokens,
+                score_min=1.0,
+                score_max=10.0,
+                prefix="crit",
+                use_placeholder=data_args.use_placeholder,
+            )
+            total_registry = ExistingVocabScoreTokenRegistry(
+                token_list=total_tokens,
+                score_min=0.0,
+                score_max=100.0,
+                prefix="total",
+                use_placeholder=data_args.use_placeholder,
+            )
+            model, processor.tokenizer = expand_vocab_for_existing_tokens(
+                model, processor.tokenizer, [crit_registry, total_registry]
+            )
+        else:
+            crit_registry = ScoreTokenRegistry(
+                n_bins=data_args.crit_score_bins,
+                score_min=1.0,
+                score_max=10.0,
+                prefix="score",
+                use_placeholder=data_args.use_placeholder,
+            )
+            total_registry = ScoreTokenRegistry(
+                n_bins=data_args.total_score_bins,
+                score_min=0.0,
+                score_max=100.0,
+                prefix="total",
+                use_placeholder=data_args.use_placeholder,
+            )
+            model, processor.tokenizer = expand_vocab_for_scoring(
+                model, processor.tokenizer, [crit_registry, total_registry]
+            )
+        score_loss_fn = ScoreLoss(
+            crit_registry=crit_registry,
+            total_registry=total_registry,
+            tokenizer=processor.tokenizer,
+            ce_weight=training_args.score_ce_weight,
+            crit_reg_weight=training_args.crit_reg_weight,
+            total_reg_weight=training_args.total_reg_weight,
+            regression_loss=training_args.score_regression_loss,
+        )
+
+    if data_args.data_flatten or data_args.data_packing:
+        replace_qwen2_vl_attention_class()
+    model.config.use_cache = False
+
+    if training_args.gradient_checkpointing:
+        if hasattr(model, "enable_input_require_grads"):
+            model.enable_input_require_grads()
+        else:
+
+            def make_inputs_require_grad(module, input, output):
+                output.requires_grad_(True)
+
+            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)
+
+    tokenizer = transformers.AutoTokenizer.from_pretrained(
+        model_args.model_name_or_path,
+        cache_dir=training_args.cache_dir,
+        model_max_length=training_args.model_max_length,
+        padding_side="right",
+        use_fast=False,
+    )
+    if data_args.use_score_tokens:
+        tokenizer = processor.tokenizer
+
+    if training_args.lora_enable:
+        from peft import LoraConfig, get_peft_model, TaskType
+        print("LoRA enabled")
+
+        for p in model.parameters():
+            p.requires_grad = False
+
+        lora_config = LoraConfig(
+            r=training_args.lora_r or 64,
+            lora_alpha=training_args.lora_alpha or 128,
+            lora_dropout=training_args.lora_dropout or 0.05,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],  # Qwen 的 attention 线性层
+            bias="none",
+            task_type=TaskType.CAUSAL_LM,
+        )
+        model = get_peft_model(model, lora_config)
+    else:
+        set_model(model_args, model)
+
+        if torch.distributed.get_rank() == 0:
+            model.visual.print_trainable_parameters()
+            model.model.print_trainable_parameters()
+    if torch.distributed.get_rank() == 0:
+        for name, param in model.named_parameters():
+            print(f"{name}:{param.requires_grad}")
+    
+    if data_args.use_score_tokens:
+        from qwenvl.data.data_processor import update_processor_pixels
+        processor = update_processor_pixels(processor, data_args)
+        from qwenvl.data import data_dict as _data_dict
+        _ds_cfg = _data_dict.get("aesthetic_track1", {})
+        annotation_path = os.environ.get("AESTHETIC_ANNOTATION", _ds_cfg["annotation_path"])
+        image_dir       = os.environ.get("AESTHETIC_IMAGE_DIR",  _ds_cfg["data_path"])
+        train_dataset = AestheticScoringDataset(
+            processor=processor,
+            annotation_path=annotation_path,
+            image_dir=image_dir,
+            data_args=data_args,
+            crit_registry=crit_registry,
+            total_registry=total_registry,
+        )
+        from qwenvl.data.data_processor import DataCollatorForSupervisedDataset
+        base_collator = DataCollatorForSupervisedDataset(processor.tokenizer)
+        data_collator = AestheticDataCollator(base_collator)
+        data_module = dict(train_dataset=train_dataset, eval_dataset=None, data_collator=data_collator)
+    else:
+        data_module = make_supervised_data_module(processor, data_args=data_args)
+
+    # Build a small fixed sample set for generation logging (rank-0 only)
+    _sample_items = []
+    if data_args.use_score_tokens and local_rank in (-1, 0):
+        import json as _json
+        _ann_path = os.environ.get("AESTHETIC_ANNOTATION", "")
+        _img_root = os.environ.get("AESTHETIC_IMAGE_DIR", "")
+        try:
+            _all_raw = _json.load(open(_ann_path))
+            _raw = [x for x in _all_raw if "question" not in x][:4]
+            from qwenvl.data.aesthetic_dataset import CRITERIA_ORDER, CRITERIA_QUESTIONS
+            for _item in _raw:
+                _img = str(Path(_img_root) / _item["image_path"])
+                _crit = CRITERIA_ORDER[0]
+                _sample_items.append({
+                    "image_path": _img,
+                    "question":   CRITERIA_QUESTIONS[_crit],
+                    "gt_score":   _item["criteria"].get(_crit, {}).get("score", "?"),
+                })
+        except Exception as _e:
+            rank0_print(f"[LocalLoggingCallback] Could not build sample items: {_e}")
+
+    _logging_cb = LocalLoggingCallback(
+        processor=processor,
+        sample_dataset=_sample_items,
+        sample_every_steps=training_args.sample_every_steps,
+        output_dir=training_args.output_dir,
+    )
+
+    trainer = AestheticTrainer(
+        model=model,
+        processing_class=tokenizer,
+        args=training_args,
+        score_loss=score_loss_fn,
+        tokenizer=tokenizer,
+        callbacks=[_logging_cb],
+        **data_module,
+    )
+
+    if list(pathlib.Path(training_args.output_dir).glob("checkpoint-*")):
+        logging.info("checkpoint found, resume training")
+        trainer.train(resume_from_checkpoint=True)
+    else:
+        trainer.train()
+    trainer.save_state()
+
+    model.config.use_cache = True
+
+    safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
+    
+    processor.save_pretrained(training_args.output_dir)
+
+
+if __name__ == "__main__":
+    train(attn_implementation="flash_attention_2")

scripts/sft_aesthetic.sh

@@ -0,0 +1,146 @@
+#!/bin/bash
+# ──────────────────────────────────────────────────────────────────────────────
+# Training script for aesthetic scoring with custom score tokens.
+#
+# Each image is expanded into 14 independent samples:
+#   - 13 per-criterion samples  (score range 1–10, bins=crit_score_bins)
+#   - 1  total-score sample     (score range 0–100, bins=total_score_bins)
+#
+# Assistant output format (per sample):
+#   <think>{reason}</think>
+#   The score is <scoreK>          ← criteria sample
+#
+#   The score is <totalK>          ← total sample
+#
+# Loss = CE_loss + crit_reg_weight * Huber(criteria) + total_reg_weight * Huber(total)
+# ──────────────────────────────────────────────────────────────────────────────
+# export CUDA_VISIBLE_DEVICES=0
+# ── Distributed configuration (auto-detect) ──────────────────────────────────
+MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+MASTER_PORT=${MASTER_PORT:-$(shuf -i 20001-29999 -n 1)}
+NNODES=${WORLD_SIZE:-1}
+
+if [ -z "$NPROC_PER_NODE" ]; then
+    if command -v nvidia-smi &> /dev/null; then
+        NPROC_PER_NODE=$(nvidia-smi --list-gpus | wc -l)
+    else
+        NPROC_PER_NODE=1
+    fi
+fi
+
+
+export DS_BUILD_SHM_COMM=0  # Disable building shared memory communication ops
+export DS_BUILD_OPS=0  # Disable building ops to avoid compilation errors
+# ── Paths ─────────────────────────────────────────────────────────────────────
+# Base model: Qwen3-VL-8B or a previously fine-tuned checkpoint
+# LLM=/mnt/hdfs/zw04mlnn01/checkpoint/llm_platform/hldy/model/Qwen/Qwen3-VL-8B-Instruct/main
+LLM=
+# Image root directory (must contain all image_path entries in track_1_train.json)
+IMAGE_DIR=zijielou/PortraitCraft/main
+
+# Annotation file (image_path uses subdir format: images_00/xxx.jpg)
+ANNOTATION=PortraitCraft/track_1_train_subdir.json
+
+# Output
+RUN_NAME="qwen3vl_aesthetic_score_tokens_0509"
+OUTPUT_DIR=result/portrait/${RUN_NAME}
+
+# DeepSpeed config
+DEEPSPEED=./scripts/zero2.json
+
+# ── Hyper-parameters ──────────────────────────────────────────────────────────
+LR=2e-5
+BATCH_SIZE=2               # per GPU
+GRAD_ACCUM=16                # effective batch = NPROC_PER_NODE * BATCH_SIZE * GRAD_ACCUM
+NUM_EPOCHS=3
+
+# ── Score token settings ──────────────────────────────────────────────────────
+CRIT_BINS=20                # <score1>…<score20> covering 1.0–10.0
+TOTAL_BINS=101              # <total1>…<total101> covering 0–100
+
+# ── Loss weights ──────────────────────────────────────────────────────────────
+CE_WEIGHT=1.0
+CRIT_REG_WEIGHT=0.5         # regression loss weight for criteria samples
+TOTAL_REG_WEIGHT=0.5        # regression loss weight for total-score samples
+REG_LOSS_TYPE=huber         # huber | mse
+
+# ── Image resolution ──────────────────────────────────────────────────────────
+MAX_PIXELS=$((32 * 32 * 1024))   # ~451k pixels (~672px side)
+MIN_PIXELS=$((32 * 32 * 224))    # ~12k  pixels
+IMAGE_MAX_TOKEN_NUM=1024
+# ── Entry point ───────────────────────────────────────────────────────────────
+ENTRY=qwenvl/train/train_qwen.py
+
+# ── Update aesthetic data path in data registry ───────────────────────────────
+# Patch the image_dir at runtime via environment variable so we don't hard-code
+# it inside __init__.py.  The AestheticScoringDataset reads IMAGE_DIR if set.
+export AESTHETIC_IMAGE_DIR=${IMAGE_DIR}
+export AESTHETIC_ANNOTATION=${ANNOTATION}
+export WANDB_DISABLED=true
+
+# ── Auto-compute save_steps = half epoch ──────────────────────────────────────
+NUM_IMAGES=$(python -c "import json; print(len(json.load(open('${ANNOTATION}'))))")
+SAMPLES_PER_IMAGE=14   # 13 criteria + 1 total
+TOTAL_SAMPLES=$((NUM_IMAGES * SAMPLES_PER_IMAGE))
+EFFECTIVE_BATCH=$((NPROC_PER_NODE * NNODES * BATCH_SIZE * GRAD_ACCUM))
+STEPS_PER_EPOCH=$(( (TOTAL_SAMPLES + EFFECTIVE_BATCH - 1) / EFFECTIVE_BATCH ))
+SAVE_STEPS=$(( STEPS_PER_EPOCH / 2 ))
+[ ${SAVE_STEPS} -lt 1 ] && SAVE_STEPS=1
+echo "[auto] ${NUM_IMAGES} images × ${SAMPLES_PER_IMAGE} = ${TOTAL_SAMPLES} samples"
+echo "[auto] effective_batch=${EFFECTIVE_BATCH}, steps/epoch=${STEPS_PER_EPOCH}, save_steps=${SAVE_STEPS} (half epoch)"
+
+# ── Training arguments ────────────────────────────────────────────────────────
+ARGS="
+    --model_name_or_path ${LLM} \
+    --use_score_tokens True \
+    --crit_score_bins ${CRIT_BINS} \
+    --total_score_bins ${TOTAL_BINS} \
+    --score_ce_weight ${CE_WEIGHT} \
+    --crit_reg_weight ${CRIT_REG_WEIGHT} \
+    --total_reg_weight ${TOTAL_REG_WEIGHT} \
+    --tune_mm_vision True \
+    --tune_mm_mlp True \
+    --tune_mm_llm True \
+    --bf16 True \
+    --output_dir ${OUTPUT_DIR} \
+    --num_train_epochs ${NUM_EPOCHS} \
+    --per_device_train_batch_size ${BATCH_SIZE} \
+    --per_device_eval_batch_size $((BATCH_SIZE * 2)) \
+    --gradient_accumulation_steps ${GRAD_ACCUM} \
+    --max_pixels ${MAX_PIXELS} \
+    --min_pixels ${MIN_PIXELS} \
+    --eval_strategy no \
+    --save_strategy steps \
+    --save_steps ${SAVE_STEPS} \
+    --save_total_limit 6 \
+    --learning_rate ${LR} \
+    --weight_decay 0.0 \
+    --warmup_ratio 0.03 \
+    --max_grad_norm 2.0 \
+    --lr_scheduler_type cosine \
+    --logging_steps 1 \
+    --model_max_length 4096 \
+    --gradient_checkpointing False \
+    --dataloader_num_workers 4 \
+    --run_name ${RUN_NAME} \
+    --sample_every_steps 500 \
+    --report_to none \
+    --deepspeed ${DEEPSPEED}
+"
+
+# ── Launch ────────────────────────────────────────────────────────────────────
+mkdir -p ${OUTPUT_DIR}
+
+echo "============================================================"
+echo " GPUs: ${NPROC_PER_NODE}, Nodes: ${NNODES}"
+echo " Master: ${MASTER_ADDR}:${MASTER_PORT}"
+echo " Output: ${OUTPUT_DIR}"
+echo "============================================================"
+
+torchrun \
+    --nproc_per_node=${NPROC_PER_NODE} \
+    --nnodes=${NNODES} \
+    --master_addr=${MASTER_ADDR} \
+    --master_port=${MASTER_PORT} \
+    ${ENTRY} ${ARGS} 2>&1 | tee ${OUTPUT_DIR}/train.log
+

scripts/zero2.json

@@ -0,0 +1,23 @@
+{
+    "fp16": {
+        "enabled": "auto",
+        "loss_scale": 0,
+        "loss_scale_window": 1000,
+        "initial_scale_power": 16,
+        "hysteresis": 2,
+        "min_loss_scale": 1
+    },
+    "bf16": {
+        "enabled": "auto"
+    },
+    "train_micro_batch_size_per_gpu": "auto",
+    "train_batch_size": "auto",
+    "gradient_accumulation_steps": "auto",
+    "zero_optimization": {
+        "stage": 2,
+        "overlap_comm": true,
+        "contiguous_gradients": true,
+        "sub_group_size": 1e9,
+        "reduce_bucket_size": "auto"
+    }
+}