Views
No views yet
en)1"""
2Single-clip inference: backbone and head both loaded from HuggingFace Hub
3(or local HF_HOME cache).
4
5Usage:
6 python vq_infer.py --video path/to/clip.mp4
7
8 Backbone : cyankiwi/Qwen3.5-27B-AWQ-4bit (vLLM pooling, AWQ-Marlin)
9 Head : Nastooh/vq_head (PyTorchModelHubMixin)
10"""
11
12from __future__ import annotations
13
14import argparse
15import os
16import time
17import warnings
18
19import numpy as np
20import torch
21import torch.nn as nn
22from huggingface_hub import PyTorchModelHubMixin
23
24import vllm
25from vllm import LLM, EngineArgs
26from vllm.assets.video import video_get_metadata, video_to_ndarrays
27
28# ---------------------------------------------------------------------------
29# Configuration
30# ---------------------------------------------------------------------------
31BACKBONE_REPO = "cyankiwi/Qwen3.5-27B-AWQ-4bit"
32HEAD_REPO = "Nastooh/vq_head"
33MODEL_DTYPE = "float16"
34MAX_NUM_SEQS = 1
35
36VQ_PROMPT = (
37 "Carefully assess the perceptual visual quality of this video, "
38 "considering sharpness, noise, compression artifacts, motion blur, "
39 "color fidelity, and overall fidelity."
40)
41
42os.environ["VLLM_CONFIGURE_LOGGING"] = "0"
43os.environ["VLLM_LOGGING_LEVEL"] = "ERROR"
44os.environ["TQDM_DISABLE"] = "1"
45
46
47# ---------------------------------------------------------------------------
48# Backbone (vLLM pooling mode — runs AWQ-Marlin kernels, no decompression)
49# ---------------------------------------------------------------------------
50def build_extractor() -> LLM:
51 engine_args = EngineArgs(
52 model = BACKBONE_REPO,
53 runner = "pooling",
54 max_model_len = -1,
55 max_num_seqs = MAX_NUM_SEQS,
56 limit_mm_per_prompt = {"video": MAX_NUM_SEQS},
57 dtype = MODEL_DTYPE,
58 trust_remote_code = True,
59 enforce_eager = False,
60 async_scheduling = True,
61 tensor_parallel_size = 1,
62 enable_prefix_caching = True,
63 gpu_memory_utilization= 0.8,
64 )
65 print(f"vLLM {vllm.__version__}")
66 print(f"Loading backbone: {BACKBONE_REPO}")
67 return LLM.from_engine_args(engine_args)
68
69
70def embed(llm: LLM, video_path: str) -> np.ndarray:
71 """Return a (D,) float32 pooled embedding for one video."""
72 prompt = (
73 "<|im_start|>system\nYou are a strict video quality grader.<|im_end|>\n"
74 "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
75 f"{VQ_PROMPT}<|im_end|>\n"
76 "<|im_start|>assistant\n"
77 )
78 frames = video_to_ndarrays(video_path)
79 meta = video_get_metadata(video_path)
80 out = llm.embed([{"prompt": prompt,
81 "multi_modal_data": {"video": (frames, meta)}}],
82 use_tqdm=False)[0]
83 vec = getattr(out.outputs, "embedding", None) or getattr(out.outputs, "data", None)
84 if vec is None:
85 raise RuntimeError("Could not locate embedding in RequestOutput")
86 return np.asarray(vec, dtype=np.float32)
87
88
89# ---------------------------------------------------------------------------
90# Head (PyTorchModelHubMixin — loaded from Hub or local HF_HOME cache)
91# ---------------------------------------------------------------------------
92class QualityHead(nn.Module, PyTorchModelHubMixin):
93 def __init__(self,
94 embed_dim: int,
95 hidden_dims: list[int] | None = None,
96 dropout: float = 0.0,
97 y_mean: float = 0.0,
98 y_std: float = 1.0):
99 super().__init__()
100
101 self.y_mean = y_mean
102 self.y_std = y_std
103 layers, prev = [], embed_dim
104 for h in hidden_dims:
105 layers += [nn.Linear(prev, h), nn.ReLU()]
106 if dropout > 0.0:
107 layers.append(nn.Dropout(dropout))
108 prev = h
109 layers.append(nn.Linear(prev, 1))
110 self.net = nn.Sequential(*layers)
111
112 def forward(self, x: torch.Tensor) -> torch.Tensor:
113 return self.net(x).squeeze(-1)
114
115 def predict_mos(self, x: torch.Tensor) -> torch.Tensor:
116 return self.forward(x) * self.y_std + self.y_mean
117
118
119def load_head(embed_dim: int, device: str) -> QualityHead:
120 # Prefer flat local cache (HF_HOME/hub/models--<org>--<name>/) to avoid
121 # a network round-trip when the snapshot layout isn't present.
122 hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
123 _slug = HEAD_REPO.replace("/", "--")
124 _local = os.path.join(hf_home, "hub", f"models--{_slug}")
125 head_src = _local if os.path.isfile(os.path.join(_local, "model.safetensors")) else HEAD_REPO
126 print(f"Loading head from: {head_src}")
127 head = QualityHead.from_pretrained(head_src, embed_dim=embed_dim).to(device)
128 head.eval()
129 return head
130
131
132# ---------------------------------------------------------------------------
133# Entry point
134# ---------------------------------------------------------------------------
135def main():
136 ap = argparse.ArgumentParser(description=__doc__,
137 formatter_class=argparse.RawDescriptionHelpFormatter)
138 ap.add_argument("--video", required=True, help="Path to the video clip")
139 args = ap.parse_args()
140
141 device = "cuda" if torch.cuda.is_available() else "cpu"
142 llm = build_extractor()
143
144 t0 = time.perf_counter()
145 vec = embed(llm, args.video) # (D,) float32
146 head = load_head(embed_dim=vec.shape[-1], device=device)
147 with torch.no_grad():
148 score = head.predict_mos(
149 torch.from_numpy(vec).unsqueeze(0).to(device)
150 ).item()
151 dt = time.perf_counter() - t0
152
153 print(f"video : {args.video}")
154 print(f"score : {score:.4f}")
155 print(f"time : {dt:.2f}s")
156
157
158if __name__ == "__main__":
159 warnings.simplefilter("ignore", FutureWarning)
160 torch.set_float32_matmul_precision("high")
161
162 main()python3 -Xfrozen_modules=off vq_infer.py --video <video file path>
...
video : <video file path>
score : 19.2252
time : 52.68s