Nastooh/vq_head

YAML Front-Matter Metadata Section

language:

• en license: mit tags:
• video quality assessment
• mos datasets:
• Beyond8Bits metrics:
• pearson linear correlation coefficient
• spearman rank correlation coefficient

pipeline_tag: video quality assessment

Model Card for MOS VQ

This model is a Mean Opinion Score (MOS) visual quality assessment model that uses cyankiwi/Qwen3.5-27B-AWQ-4bit as its backbone. The trained head is a 400-node single layer, which is followed by a single regression node.

Model Details

• Developed by: Nastooh Avesta
• Model type: Transfer Learning
• Language(s) (NLP): English (en)
• License: MIT
• Finetuned from model: cyankiwi/Qwen3.5-27B-AWQ-4bit

Uses

Direct Use

This model can be used directly assess the MOS of an input video.

Out-of-Scope Use

N/A

Bias, Risks, and Limitations

N/A

How to Get Started with the Model

A vllm based example looks as follows:

"""
Single-clip inference: backbone and head both loaded from HuggingFace Hub
(or local HF_HOME cache).

Usage:
    python vq_infer.py --video path/to/clip.mp4

    Backbone : cyankiwi/Qwen3.5-27B-AWQ-4bit   (vLLM pooling, AWQ-Marlin)
    Head     : Nastooh/vq_head                  (PyTorchModelHubMixin)
"""

from __future__ import annotations

import argparse
import os
import time
import warnings

import numpy as np
import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin

import vllm
from vllm import LLM, EngineArgs
from vllm.assets.video import video_get_metadata, video_to_ndarrays

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
BACKBONE_REPO = "cyankiwi/Qwen3.5-27B-AWQ-4bit"
HEAD_REPO     = "Nastooh/vq_head"
MODEL_DTYPE   = "float16"
MAX_NUM_SEQS  = 1

VQ_PROMPT = (
    "Carefully assess the perceptual visual quality of this video, "
    "considering sharpness, noise, compression artifacts, motion blur, "
    "color fidelity, and overall fidelity."
)

os.environ["VLLM_CONFIGURE_LOGGING"] = "0"
os.environ["VLLM_LOGGING_LEVEL"]     = "ERROR"
os.environ["TQDM_DISABLE"]           = "1"


# ---------------------------------------------------------------------------
# Backbone (vLLM pooling mode — runs AWQ-Marlin kernels, no decompression)
# ---------------------------------------------------------------------------
def build_extractor() -> LLM:
    engine_args = EngineArgs(
        model                 = BACKBONE_REPO,
        runner                = "pooling",
        max_model_len         = -1,
        max_num_seqs          = MAX_NUM_SEQS,
        limit_mm_per_prompt   = {"video": MAX_NUM_SEQS},
        dtype                 = MODEL_DTYPE,
        trust_remote_code     = True,
        enforce_eager         = False,
        async_scheduling      = True,
        tensor_parallel_size  = 1,
        enable_prefix_caching = True,
        gpu_memory_utilization= 0.8,
    )
    print(f"vLLM {vllm.__version__}")
    print(f"Loading backbone: {BACKBONE_REPO}")
    return LLM.from_engine_args(engine_args)


def embed(llm: LLM, video_path: str) -> np.ndarray:
    """Return a (D,) float32 pooled embedding for one video."""
    prompt = (
        "<|im_start|>system\nYou are a strict video quality grader.<|im_end|>\n"
        "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
        f"{VQ_PROMPT}<|im_end|>\n"
        "<|im_start|>assistant\n"
    )
    frames = video_to_ndarrays(video_path)
    meta   = video_get_metadata(video_path)
    out    = llm.embed([{"prompt": prompt,
                         "multi_modal_data": {"video": (frames, meta)}}],
                       use_tqdm=False)[0]
    vec = getattr(out.outputs, "embedding", None) or getattr(out.outputs, "data", None)
    if vec is None:
        raise RuntimeError("Could not locate embedding in RequestOutput")
    return np.asarray(vec, dtype=np.float32)


# ---------------------------------------------------------------------------
# Head (PyTorchModelHubMixin — loaded from Hub or local HF_HOME cache)
# ---------------------------------------------------------------------------
class QualityHead(nn.Module, PyTorchModelHubMixin):
    def __init__(self,
                 embed_dim: int,
                 hidden_dims: list[int] | None = None,
                 dropout: float = 0.0,
                 y_mean: float = 0.0,
                 y_std: float = 1.0):
        super().__init__()

        self.y_mean = y_mean
        self.y_std  = y_std
        layers, prev = [], embed_dim
        for h in hidden_dims:
            layers += [nn.Linear(prev, h), nn.ReLU()]
            if dropout > 0.0:
                layers.append(nn.Dropout(dropout))
            prev = h
        layers.append(nn.Linear(prev, 1))
        self.net = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x).squeeze(-1)

    def predict_mos(self, x: torch.Tensor) -> torch.Tensor:
        return self.forward(x) * self.y_std + self.y_mean


def load_head(embed_dim: int, device: str) -> QualityHead:
    # Prefer flat local cache (HF_HOME/hub/models--<org>--<name>/) to avoid
    # a network round-trip when the snapshot layout isn't present.
    hf_home  = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
    _slug    = HEAD_REPO.replace("/", "--")
    _local   = os.path.join(hf_home, "hub", f"models--{_slug}")
    head_src = _local if os.path.isfile(os.path.join(_local, "model.safetensors")) else HEAD_REPO
    print(f"Loading head from: {head_src}")
    head = QualityHead.from_pretrained(head_src, embed_dim=embed_dim).to(device)
    head.eval()
    return head


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--video", required=True, help="Path to the video clip")
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    llm    = build_extractor()

    t0  = time.perf_counter()
    vec = embed(llm, args.video)                          # (D,) float32
    head = load_head(embed_dim=vec.shape[-1], device=device)
    with torch.no_grad():
        score = head.predict_mos(
            torch.from_numpy(vec).unsqueeze(0).to(device)
        ).item()
    dt = time.perf_counter() - t0

    print(f"video  : {args.video}")
    print(f"score  : {score:.4f}")
    print(f"time   : {dt:.2f}s")


if __name__ == "__main__":
    warnings.simplefilter("ignore", FutureWarning)
    torch.set_float32_matmul_precision("high")

    main()

Sample Run

python3 -Xfrozen_modules=off vq_infer.py --video <video file path>
...
video  : <video file path>
score  : 19.2252
time   : 52.68s

Performance

• PLCC: 0.953
• SRCC: 0.953

Training Details

Training Data

The model was trained on the Beyond8Bits Dataset, which consists of 25,917 source videos that expand into 41,419 transcoded clips.

Evaluation

Model shows a good fit on the validation set:

Testing Data & Metrics