Add VEFX-Bench reference code

.gitignore

@@ -0,0 +1,37 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+*.egg
+
+# Environment
+.env
+*.env
+
+# IDE
+.vscode/
+.idea/
+
+# Model weights
+*.safetensors
+*.pth
+*.bin
+*.ckpt
+*.onnx
+
+# Data
+*.avi
+*.mov
+
+# Keep sample videos for examples
+!examples/sample_videos/*.mp4
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Outputs
+merged_model/
+results/

LICENSE

@@ -0,0 +1,176 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to the Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by the Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding any notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS

examples/batch_scoring.py

@@ -0,0 +1,51 @@
+"""
+Batch scoring: Evaluate multiple video edits from a CSV file.
+
+Expected CSV format:
+    original_video,edited_video,instruction
+    path/to/orig1.mp4,path/to/edit1.mp4,"make it snowy"
+    path/to/orig2.mp4,path/to/edit2.mp4,"add a red hat"
+
+Usage:
+    python examples/batch_scoring.py \
+        --csv edits.csv \
+        --output results.csv
+"""
+
+import argparse
+import csv
+
+import torch
+from vefx_reward import VEFXReward
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Batch score video edits")
+    parser.add_argument("--csv", required=True, help="Input CSV with columns: original_video, edited_video, instruction")
+    parser.add_argument("--output", default="results.csv", help="Output CSV path")
+    parser.add_argument("--model", default="VEFX-Reward/VEFX-Reward-4B")
+    parser.add_argument("--device", default="cuda")
+    args = parser.parse_args()
+
+    model = VEFXReward(args.model, device=args.device)
+
+    with open(args.csv) as f:
+        rows = list(csv.DictReader(f))
+    print(f"Loaded {len(rows)} samples from {args.csv}")
+
+    results = []
+    for i, row in enumerate(rows):
+        scores = model.score(row["original_video"], row["edited_video"], row["instruction"])
+        results.append({**row, **scores})
+        print(f"[{i+1}/{len(rows)}] IF={scores['IF']:.2f}  RQ={scores['RQ']:.2f}  EE={scores['EE']:.2f}  Overall={scores['Overall']:.2f}")
+
+    fieldnames = list(rows[0].keys()) + ["IF", "RQ", "EE", "Overall"]
+    with open(args.output, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames)
+        writer.writeheader()
+        writer.writerows(results)
+    print(f"\nResults saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

examples/multi_gpu_scoring.py

@@ -0,0 +1,119 @@
+"""
+Multi-GPU parallel scoring using subprocess workers.
+
+Splits a CSV of video edits across multiple GPUs for faster inference.
+
+Usage:
+    python examples/multi_gpu_scoring.py \
+        --csv edits.csv \
+        --output results.csv \
+        --num_gpus 4
+"""
+
+import argparse
+import csv
+import json
+import os
+import subprocess
+import sys
+import tempfile
+
+
+def worker_main(args):
+    """Single-GPU worker: load model, score shard, write results."""
+    import torch
+    from vefx_reward import VEFXReward
+
+    with open(args.shard_file) as f:
+        shard = json.load(f)
+
+    model = VEFXReward(args.model, device="cuda:0")
+
+    results = []
+    for i, item in enumerate(shard):
+        try:
+            scores = model.score(item["original_video"], item["edited_video"], item["instruction"])
+            results.append({**item, **scores})
+            print(f"[GPU {args.gpu_id}] [{i+1}/{len(shard)}] "
+                  f"IF={scores['IF']:.2f} RQ={scores['RQ']:.2f} EE={scores['EE']:.2f}", flush=True)
+        except Exception as e:
+            print(f"[GPU {args.gpu_id}] [{i+1}/{len(shard)}] ERROR: {e}", flush=True)
+            results.append({**item, "IF": None, "RQ": None, "EE": None, "Overall": None, "error": str(e)})
+
+    with open(args.output_file, "w") as f:
+        json.dump(results, f)
+    print(f"[GPU {args.gpu_id}] Done — {len(results)} results", flush=True)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Multi-GPU video edit scoring")
+    parser.add_argument("--csv", required=True, help="Input CSV")
+    parser.add_argument("--output", default="results.csv", help="Output CSV")
+    parser.add_argument("--model", default="VEFX-Reward/VEFX-Reward-4B")
+    parser.add_argument("--num_gpus", type=int, default=4)
+    # Internal worker args
+    parser.add_argument("--_worker", action="store_true", help=argparse.SUPPRESS)
+    parser.add_argument("--gpu_id", type=int, default=0, help=argparse.SUPPRESS)
+    parser.add_argument("--shard_file", type=str, default="", help=argparse.SUPPRESS)
+    parser.add_argument("--output_file", type=str, default="", help=argparse.SUPPRESS)
+    args = parser.parse_args()
+
+    if args._worker:
+        worker_main(args)
+        return
+
+    # --- Launcher mode ---
+    with open(args.csv) as f:
+        rows = list(csv.DictReader(f))
+    print(f"Loaded {len(rows)} samples, distributing across {args.num_gpus} GPUs")
+
+    items = [dict(row) for row in rows]
+    shards = [[] for _ in range(args.num_gpus)]
+    for i, item in enumerate(items):
+        shards[i % args.num_gpus].append(item)
+
+    tmpdir = tempfile.mkdtemp(prefix="vefx_multi_")
+    script = os.path.abspath(__file__)
+    procs = []
+    for gid in range(args.num_gpus):
+        if not shards[gid]:
+            continue
+        sf = os.path.join(tmpdir, f"shard_{gid}.json")
+        of = os.path.join(tmpdir, f"result_{gid}.json")
+        with open(sf, "w") as f:
+            json.dump(shards[gid], f)
+        env = os.environ.copy()
+        env["CUDA_VISIBLE_DEVICES"] = str(gid)
+        env["TOKENIZERS_PARALLELISM"] = "false"
+        p = subprocess.Popen(
+            [sys.executable, script,
+             "--_worker", "--gpu_id", str(gid),
+             "--shard_file", sf, "--output_file", of,
+             "--model", args.model],
+            env=env, stdout=sys.stdout, stderr=sys.stderr,
+        )
+        procs.append((p, of))
+
+    for p, _ in procs:
+        p.wait()
+
+    # Merge results
+    all_results = []
+    for _, of in procs:
+        if os.path.exists(of):
+            with open(of) as f:
+                all_results.extend(json.load(f))
+
+    fieldnames = list(rows[0].keys()) + ["IF", "RQ", "EE", "Overall"]
+    with open(args.output, "w", newline="") as f:
+        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
+        writer.writeheader()
+        writer.writerows(all_results)
+    print(f"\nAll done — {len(all_results)} results saved to {args.output}")
+
+    import shutil
+    shutil.rmtree(tmpdir, ignore_errors=True)
+
+
+if __name__ == "__main__":
+    main()

examples/quick_start.py

@@ -0,0 +1,39 @@
+"""
+Quick start: Score a single video edit with VEFX-Reward.
+
+Usage:
+    python examples/quick_start.py \
+        --original path/to/original.mp4 \
+        --edited path/to/edited.mp4 \
+        --instruction "add a hat to the person"
+"""
+
+import argparse
+import torch
+from vefx_reward import VEFXReward
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Score a video edit with VEFX-Reward")
+    parser.add_argument("--original", required=True, help="Path to original video")
+    parser.add_argument("--edited", required=True, help="Path to edited video")
+    parser.add_argument("--instruction", required=True, help="Editing instruction")
+    parser.add_argument("--model", default="VEFX-Reward/VEFX-Reward-4B", help="Model path or HF ID")
+    parser.add_argument("--device", default="cuda", help="Device (cuda / cpu)")
+    args = parser.parse_args()
+
+    model = VEFXReward(args.model, device=args.device)
+    scores = model.score(args.original, args.edited, args.instruction)
+
+    print("\n" + "=" * 50)
+    print("VEFX-Reward Scores")
+    print("=" * 50)
+    print(f"  Instructional Following (IF): {scores['IF']:.2f}")
+    print(f"  Render Quality          (RQ): {scores['RQ']:.2f}")
+    print(f"  Edit Exclusivity        (EE): {scores['EE']:.2f}")
+    print(f"  Overall                     : {scores['Overall']:.2f}")
+    print("=" * 50)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,11 @@
+torch>=2.1.0
+torchvision>=0.16.0
+transformers>=4.51.0
+accelerate>=0.30.0
+safetensors>=0.4.0
+huggingface_hub>=0.20.0
+Pillow>=10.0.0
+numpy>=1.24.0
+requests
+packaging
+decord>=0.6.0

scripts/prepare_and_upload.py

@@ -0,0 +1,214 @@
+"""
+Prepare and upload VEFX-Reward model to HuggingFace Hub.
+
+This script:
+1. Loads the Qwen3-VL-4B base model
+2. Manually merges LoRA weights into the base model
+3. Loads non-LoRA weights (rm_head, merger, special token embeddings)
+4. Saves and uploads the complete model to HuggingFace
+
+Prerequisites:
+    pip install huggingface_hub safetensors
+    huggingface-cli login
+
+Usage:
+    python scripts/prepare_and_upload.py \
+        --checkpoint_dir /path/to/training/logs/v4/ord_4B_lora_2stage_promptv2_res399k \
+        --checkpoint_step 1050 \
+        --hf_repo VEFX-Reward/VEFX-Reward-4B \
+        --output_dir ./merged_model
+"""
+
+import argparse
+import json
+import os
+
+import safetensors.torch as st
+import torch
+from transformers import AutoProcessor, AutoTokenizer
+
+import sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
+from vefx_reward.model import Qwen3VLRewardModelBT
+
+SPECIAL_TOKENS = [
+    "<|VQ_reward|>", "<|MQ_reward|>", "<|TA_reward|>",
+    "<|IF_reward|>", "<|RQ_reward|>", "<|EE_reward|>",
+]
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Prepare and upload VEFX-Reward model")
+    parser.add_argument("--checkpoint_dir", required=True,
+                        help="Training output directory containing model_config.json and checkpoint-*/")
+    parser.add_argument("--checkpoint_step", type=int, default=-1,
+                        help="Checkpoint step to use (-1 = latest)")
+    parser.add_argument("--hf_repo", default="VEFX-Reward/VEFX-Reward-4B",
+                        help="HuggingFace repo ID to upload to")
+    parser.add_argument("--output_dir", default="./merged_model",
+                        help="Local directory to save merged model before upload")
+    parser.add_argument("--upload", action="store_true",
+                        help="Actually upload to HuggingFace (otherwise just save locally)")
+    args = parser.parse_args()
+
+    # 1. Load training config
+    config_path = os.path.join(args.checkpoint_dir, "model_config.json")
+    with open(config_path) as f:
+        config_dict = json.load(f)
+
+    model_config = config_dict["model_config"]
+    data_config = config_dict["data_config"]
+
+    base_model_path = model_config["model_name_or_path"]
+    output_dim = model_config["output_dim"]
+    use_ordinal = model_config["use_ordinal"]
+    num_classes = model_config["num_classes"]
+    reward_token = model_config["reward_token"]
+
+    print(f"Base model: {base_model_path}")
+    print(f"Output dim: {output_dim}, Ordinal: {use_ordinal}, Num classes: {num_classes}")
+
+    # 2. Find checkpoint
+    import glob as globmod
+    ckpt_dirs = sorted(globmod.glob(os.path.join(args.checkpoint_dir, "checkpoint-*")),
+                       key=lambda x: int(x.split("-")[-1]))
+    if args.checkpoint_step == -1:
+        ckpt_path = ckpt_dirs[-1]
+    else:
+        ckpt_path = os.path.join(args.checkpoint_dir, f"checkpoint-{args.checkpoint_step}")
+    print(f"Using checkpoint: {ckpt_path}")
+
+    # 3. Load processor from base model with checkpoint's tokenizer
+    processor = AutoProcessor.from_pretrained(base_model_path, padding_side="right")
+    ckpt_tokenizer_path = os.path.join(ckpt_path, "tokenizer")
+    if os.path.isdir(ckpt_tokenizer_path):
+        processor.tokenizer = AutoTokenizer.from_pretrained(ckpt_tokenizer_path)
+    else:
+        processor.tokenizer.add_special_tokens({"additional_special_tokens": SPECIAL_TOKENS})
+    special_token_ids = processor.tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS)
+    print(f"Tokenizer vocab size: {len(processor.tokenizer)}")
+
+    # 4. Load base model with reward head
+    print("Loading base model...")
+    model = Qwen3VLRewardModelBT.from_pretrained(
+        base_model_path,
+        torch_dtype=torch.bfloat16,
+        output_dim=output_dim,
+        reward_token=reward_token,
+        special_token_ids=special_token_ids,
+        use_ordinal=use_ordinal,
+        num_classes=num_classes,
+        use_cache=True,
+        device_map="cpu",
+    )
+    model.resize_token_embeddings(len(processor.tokenizer))
+    print(f"Model embeddings resized to {len(processor.tokenizer)}")
+
+    # 5. Manual LoRA merge (bypasses PEFT tie_word_embeddings issues)
+    print("Loading and merging adapter weights...")
+    adapter_weights = st.load_file(os.path.join(ckpt_path, "adapter_model.safetensors"), device="cpu")
+
+    with open(os.path.join(ckpt_path, "adapter_config.json")) as f:
+        lora_cfg = json.load(f)
+    scaling = lora_cfg["lora_alpha"] / lora_cfg["r"]
+
+    # Categorize adapter keys
+    base_layers, lora_As, lora_Bs, emb_As, emb_Bs = {}, {}, {}, {}, {}
+    for k, v in adapter_weights.items():
+        ck = k.replace("base_model.model.", "")
+        if ".base_layer.weight" in ck:
+            base_layers[ck.replace(".base_layer.weight", "")] = v
+        elif ".lora_A.weight" in ck:
+            lora_As[ck.replace(".lora_A.weight", "")] = v
+        elif ".lora_B.weight" in ck:
+            lora_Bs[ck.replace(".lora_B.weight", "")] = v
+        elif ".lora_embedding_A" in ck:
+            emb_As[ck.replace(".lora_embedding_A", "")] = v
+        elif ".lora_embedding_B" in ck:
+            emb_Bs[ck.replace(".lora_embedding_B", "")] = v
+
+    model_state = model.state_dict()
+
+    # Replace base layer weights (for resized lm_head / embed_tokens)
+    for mod, w in base_layers.items():
+        key = mod + ".weight"
+        if key in model_state:
+            model_state[key] = w.to(model_state[key].dtype)
+            print(f"  Replaced base layer: {key}")
+
+    # Merge LoRA: W_merged = W + B @ A * scaling
+    merged_count = 0
+    for mod in lora_As:
+        if mod in lora_Bs:
+            A, B = lora_As[mod].float(), lora_Bs[mod].float()
+            delta = (B @ A) * scaling
+            key = mod + ".weight"
+            if key in model_state:
+                model_state[key] = (model_state[key].float() + delta).to(torch.bfloat16)
+                merged_count += 1
+    print(f"  Merged {merged_count} LoRA modules")
+
+    # Merge embedding LoRA
+    for mod in emb_As:
+        if mod in emb_Bs:
+            A, B = emb_As[mod].float(), emb_Bs[mod].float()
+            delta = (B @ A).T * scaling
+            key = mod + ".weight"
+            if key in model_state:
+                model_state[key] = (model_state[key].float() + delta).to(torch.bfloat16)
+                print(f"  Merged embedding LoRA: {key}")
+
+    # 6. Load non-LoRA weights (rm_head, merger, special embeddings)
+    non_lora_path = os.path.join(ckpt_path, "non_lora_state_dict.pth")
+    if os.path.exists(non_lora_path):
+        print("Loading non-LoRA weights...")
+        non_lora_weights = torch.load(non_lora_path, map_location="cpu")
+        for k, v in non_lora_weights.items():
+            ck = k.replace("base_model.model.", "")
+            if ck in model_state:
+                model_state[ck] = v.to(model_state[ck].dtype)
+                print(f"  Loaded: {ck}")
+
+    model.load_state_dict(model_state)
+    print(f"All weights loaded. rm_head shape: {model.rm_head.weight.shape}")
+
+    # 7. Save merged model
+    os.makedirs(args.output_dir, exist_ok=True)
+    print(f"Saving merged model to {args.output_dir}...")
+    model.save_pretrained(args.output_dir, safe_serialization=True)
+    processor.save_pretrained(args.output_dir)
+
+    # Save VEFX-specific config
+    vefx_config = {
+        "output_dim": output_dim,
+        "use_ordinal": use_ordinal,
+        "num_classes": num_classes,
+        "reward_token": reward_token,
+        "fps": data_config.get("fps", 4.0),
+        "max_frame_pixels": data_config.get("max_frame_pixels", 399360),
+        "eval_dim": data_config.get("eval_dim", ["IF", "RQ", "EE"]),
+        "prompt_template_type": data_config.get("prompt_template_type", "editreward_v2_special"),
+    }
+    with open(os.path.join(args.output_dir, "vefx_config.json"), "w") as f:
+        json.dump(vefx_config, f, indent=2)
+    print("Saved vefx_config.json")
+
+    # 8. Upload to HuggingFace
+    if args.upload:
+        from huggingface_hub import HfApi
+        api = HfApi()
+        print(f"Uploading to {args.hf_repo}...")
+        api.upload_folder(
+            folder_path=args.output_dir,
+            repo_id=args.hf_repo,
+            repo_type="model",
+        )
+        print(f"Upload complete: https://huggingface.co/{args.hf_repo}")
+    else:
+        print(f"\nModel saved to {args.output_dir}")
+        print(f"To upload, run again with --upload flag, or manually:")
+        print(f"  huggingface-cli upload {args.hf_repo} {args.output_dir} .")
+
+
+if __name__ == "__main__":
+    main()

setup.py

@@ -0,0 +1,30 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="vefx-reward",
+    version="0.1.0",
+    description="VEFX-Reward: A reward model for video editing quality assessment",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    url="",
+    packages=find_packages(),
+    python_requires=">=3.10",
+    install_requires=[
+        "torch>=2.1.0",
+        "torchvision>=0.16.0",
+        "transformers>=4.51.0",
+        "accelerate>=0.30.0",
+        "safetensors>=0.4.0",
+        "huggingface_hub>=0.20.0",
+        "Pillow>=10.0.0",
+        "numpy>=1.24.0",
+        "requests",
+        "packaging",
+        "decord>=0.6.0",
+    ],
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: Apache Software License",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+)

vefx_reward/__init__.py

@@ -0,0 +1,14 @@
+"""
+VEFX-Reward: A reward model for video editing quality assessment.
+
+Evaluates video edits on three dimensions (1–4 scale):
+- IF (Instructional Following)
+- RQ (Render Quality)
+- EE (Edit Exclusivity)
+"""
+
+__version__ = "0.1.0"
+
+from .inference import VEFXReward
+
+__all__ = ["VEFXReward"]

vefx_reward/inference.py

@@ -0,0 +1,226 @@
+"""
+VEFX-Reward: Video editing quality assessment inference API.
+
+Usage:
+    from vefx_reward import VEFXReward
+
+    model = VEFXReward("VEFX-Reward/VEFX-Reward-4B", device="cuda")
+    scores = model.score("original.mp4", "edited.mp4", "add a hat to the person")
+    # {'IF': 3.21, 'RQ': 2.85, 'EE': 3.54, 'Overall': 9.60}
+"""
+
+import json
+import os
+from collections.abc import Mapping
+from typing import Optional
+
+import numpy as np
+import torch
+from transformers import AutoProcessor
+
+from .model import Qwen3VLRewardModelBT, ordinal_predict
+from .prompt_template import build_prompt
+from .vision_process import process_vision_info
+
+# Default model hyperparameters (matching the released VEFX-Reward-4B)
+DEFAULT_FPS = 4.0
+DEFAULT_MAX_FRAME_PIXELS = 399360
+DEFAULT_NUM_CLASSES = 4
+DEFAULT_OUTPUT_DIM = 3
+DIMS = ["IF", "RQ", "EE"]
+
+SPECIAL_TOKENS = [
+    "<|VQ_reward|>", "<|MQ_reward|>", "<|TA_reward|>",
+    "<|IF_reward|>", "<|RQ_reward|>", "<|EE_reward|>",
+]
+
+
+class VEFXReward:
+    """VEFX-Reward model for video editing quality assessment.
+
+    Scores video edits on three dimensions (1–4 scale):
+    - **IF** (Instructional Following): How well the edit follows the instruction.
+    - **RQ** (Render Quality): Visual and temporal quality of the edited video.
+    - **EE** (Edit Exclusivity): Whether only the intended region was modified.
+
+    Args:
+        model_path: HuggingFace model ID or local path
+            (e.g., ``"VEFX-Reward/VEFX-Reward-4B"``).
+        device: Device string (default ``"cuda"``).
+        dtype: Torch dtype (default ``torch.bfloat16``).
+        fps: Frames per second for video sampling (default 4.0).
+        max_frame_pixels: Maximum pixels per frame (default 399360).
+
+    Example::
+
+        model = VEFXReward("VEFX-Reward/VEFX-Reward-4B")
+        scores = model.score("original.mp4", "edited.mp4", "make it snowy")
+        print(scores)
+        # {'IF': 3.2, 'RQ': 2.9, 'EE': 3.5, 'Overall': 9.6}
+    """
+
+    def __init__(
+        self,
+        model_path: str = "VEFX-Reward/VEFX-Reward-4B",
+        device: str = "cuda",
+        dtype: torch.dtype = torch.bfloat16,
+        fps: float = DEFAULT_FPS,
+        max_frame_pixels: int = DEFAULT_MAX_FRAME_PIXELS,
+    ):
+        self.device = device
+        self.dtype = dtype
+        self.fps = fps
+        self.max_frame_pixels = max_frame_pixels
+
+        # Load config
+        vefx_config_path = os.path.join(model_path, "vefx_config.json") if os.path.isdir(model_path) else None
+        if vefx_config_path and os.path.exists(vefx_config_path):
+            with open(vefx_config_path) as f:
+                vefx_config = json.load(f)
+        else:
+            # Try to download from HF hub
+            try:
+                from huggingface_hub import hf_hub_download
+                vefx_config_path = hf_hub_download(model_path, "vefx_config.json")
+                with open(vefx_config_path) as f:
+                    vefx_config = json.load(f)
+            except Exception:
+                vefx_config = {}
+
+        self.num_classes = vefx_config.get("num_classes", DEFAULT_NUM_CLASSES)
+        self.output_dim = vefx_config.get("output_dim", DEFAULT_OUTPUT_DIM)
+        self.use_ordinal = vefx_config.get("use_ordinal", True)
+        reward_token = vefx_config.get("reward_token", "special")
+
+        # Load processor and add special tokens
+        self.processor = AutoProcessor.from_pretrained(model_path, padding_side="right")
+        existing_tokens = set(self.processor.tokenizer.get_vocab().keys())
+        tokens_to_add = [t for t in SPECIAL_TOKENS if t not in existing_tokens]
+        if tokens_to_add:
+            self.processor.tokenizer.add_special_tokens({"additional_special_tokens": tokens_to_add})
+        special_token_ids = self.processor.tokenizer.convert_tokens_to_ids(SPECIAL_TOKENS)
+
+        # Load model
+        self.model = Qwen3VLRewardModelBT.from_pretrained(
+            model_path,
+            torch_dtype=dtype,
+            output_dim=self.output_dim,
+            reward_token=reward_token,
+            special_token_ids=special_token_ids,
+            use_ordinal=self.use_ordinal,
+            num_classes=self.num_classes,
+            use_cache=True,
+        )
+        self.model.resize_token_embeddings(len(self.processor.tokenizer))
+
+        self.model.eval().to(self.device)
+        print(f"VEFX-Reward loaded on {self.device} ({dtype})")
+
+    def _prepare_input(self, data):
+        if isinstance(data, Mapping):
+            return type(data)({k: self._prepare_input(v) for k, v in data.items()})
+        elif isinstance(data, (tuple, list)):
+            return type(data)(self._prepare_input(v) for v in data)
+        elif isinstance(data, torch.Tensor):
+            return data.to(device=self.device)
+        return data
+
+    def _build_batch(self, original_video: str, edited_video: str, instruction: str):
+        """Build a single-sample batch from video paths and instruction."""
+        content = [
+            {
+                "type": "video",
+                "video": f"file://{os.path.abspath(original_video)}",
+                "max_pixels": self.max_frame_pixels,
+                "fps": self.fps,
+                "sample_type": "uniform",
+            },
+            {
+                "type": "video",
+                "video": f"file://{os.path.abspath(edited_video)}",
+                "max_pixels": self.max_frame_pixels,
+                "fps": self.fps,
+                "sample_type": "uniform",
+            },
+            {"type": "text", "text": build_prompt(instruction)},
+        ]
+        messages = [[{"role": "user", "content": content}]]
+        image_inputs, video_inputs, video_metadata_list = process_vision_info(messages)
+        video_inputs = [v.float() / 255.0 for v in video_inputs]
+
+        texts = self.processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        processor_kwargs = dict(
+            text=texts,
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+            videos_kwargs={"do_rescale": False, "do_sample_frames": False},
+        )
+        if video_metadata_list:
+            processor_kwargs["videos_kwargs"]["video_metadata"] = video_metadata_list
+            processor_kwargs["videos_kwargs"]["return_metadata"] = True
+
+        batch = self.processor(**processor_kwargs)
+        return self._prepare_input(batch)
+
+    def _logits_to_scores(self, logits: torch.Tensor) -> dict:
+        """Convert raw ordinal logits to IF/RQ/EE scores."""
+        logits_np = logits.float().cpu().numpy()
+        if self.use_ordinal:
+            num_dims = self.output_dim
+            num_thresholds = self.num_classes - 1
+            logits_reshaped = logits_np.reshape(1, num_dims, num_thresholds)
+            hard, soft = ordinal_predict(logits_reshaped, self.num_classes)
+            scores = {DIMS[j]: round(float(soft[0, j]), 3) for j in range(num_dims)}
+        else:
+            scores = {DIMS[j]: round(float(logits_np[0, j]), 3) for j in range(self.output_dim)}
+        scores["Overall"] = round(sum(scores[d] for d in DIMS), 3)
+        return scores
+
+    @torch.no_grad()
+    def score(
+        self,
+        original_video: str,
+        edited_video: str,
+        instruction: str,
+    ) -> dict:
+        """Score a single video edit.
+
+        Args:
+            original_video: Path to the original (source) video.
+            edited_video: Path to the edited video.
+            instruction: The editing instruction text.
+
+        Returns:
+            Dictionary with keys ``'IF'``, ``'RQ'``, ``'EE'``, ``'Overall'``.
+            Each dimension is scored on a continuous 1–4 scale.
+        """
+        batch = self._build_batch(original_video, edited_video, instruction)
+        logits = self.model(**batch, return_dict=True)["logits"]
+        return self._logits_to_scores(logits)
+
+    @torch.no_grad()
+    def score_batch(
+        self,
+        original_videos: list[str],
+        edited_videos: list[str],
+        instructions: list[str],
+    ) -> list[dict]:
+        """Score multiple video edits (processed sequentially to avoid OOM).
+
+        Args:
+            original_videos: List of paths to original videos.
+            edited_videos: List of paths to edited videos.
+            instructions: List of editing instruction texts.
+
+        Returns:
+            List of score dictionaries, one per sample.
+        """
+        assert len(original_videos) == len(edited_videos) == len(instructions)
+        results = []
+        for orig, edit, inst in zip(original_videos, edited_videos, instructions):
+            results.append(self.score(orig, edit, inst))
+        return results

vefx_reward/model.py

@@ -0,0 +1,149 @@
+"""
+VEFX-Reward: Qwen3-VL based reward model for video editing quality assessment.
+
+Extends Qwen3VLForConditionalGeneration with an rm_head for ordinal regression,
+scoring video edits on Instructional Following (IF), Render Quality (RQ),
+and Edit Exclusivity (EE) on a 1–4 scale.
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+from typing import List, Optional
+from transformers import Qwen3VLForConditionalGeneration
+
+
+class Qwen3VLRewardModelBT(Qwen3VLForConditionalGeneration):
+    """Qwen3-VL with a reward head for ordinal video edit quality scoring."""
+
+    def __init__(self, config, output_dim=3, reward_token="special",
+                 special_token_ids=None, use_ordinal=True, num_classes=4, **kwargs):
+        if 'use_cache' in kwargs:
+            config.use_cache = kwargs.pop('use_cache')
+        super().__init__(config, **kwargs)
+        self.output_dim = output_dim
+        self.rm_head = nn.Linear(config.text_config.hidden_size, output_dim, bias=False)
+        nn.init.normal_(self.rm_head.weight, mean=0.0, std=1.0 / config.text_config.hidden_size)
+        self.reward_token = reward_token
+        self.use_ordinal = use_ordinal
+        self.num_classes = num_classes
+        self.special_token_ids = special_token_ids
+        if self.special_token_ids is not None:
+            self.reward_token = "special"
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        pixel_values: Optional[torch.Tensor] = None,
+        pixel_values_videos: Optional[torch.FloatTensor] = None,
+        image_grid_thw: Optional[torch.LongTensor] = None,
+        video_grid_thw: Optional[torch.LongTensor] = None,
+        mm_token_type_ids: Optional[torch.IntTensor] = None,
+        **kwargs,
+    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.model(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            pixel_values=pixel_values,
+            pixel_values_videos=pixel_values_videos,
+            image_grid_thw=image_grid_thw,
+            video_grid_thw=video_grid_thw,
+            mm_token_type_ids=mm_token_type_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            **kwargs,
+        )
+
+        hidden_states = outputs[0]  # [B, L, D]
+        logits = self.rm_head(hidden_states)  # [B, L, output_dim]
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        pad_token_id = self.config.text_config.pad_token_id
+        if pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                sequence_lengths = torch.eq(input_ids, pad_token_id).int().argmax(-1) - 1
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(logits.device)
+            else:
+                sequence_lengths = -1
+
+        if self.reward_token == "last":
+            pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+        elif self.reward_token == "mean":
+            valid_lengths = torch.clamp(sequence_lengths, min=0, max=logits.size(1) - 1)
+            pooled_logits = torch.stack([logits[i, :valid_lengths[i]].mean(dim=0) for i in range(batch_size)])
+        elif self.reward_token == "special":
+            special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+            for special_token_id in self.special_token_ids:
+                special_token_mask = special_token_mask | (input_ids == special_token_id)
+            pooled_logits = logits[special_token_mask, ...]
+            num_matched = special_token_mask.sum(dim=1)
+            num_dims = num_matched[0].item()
+            pooled_logits = pooled_logits.view(batch_size, num_dims, -1)
+            if self.use_ordinal:
+                pooled_logits = pooled_logits.view(batch_size, -1)
+            else:
+                if self.output_dim == num_dims:
+                    pooled_logits = pooled_logits.diagonal(dim1=1, dim2=2)
+                pooled_logits = pooled_logits.view(batch_size, -1)
+        else:
+            raise ValueError(f"Invalid reward_token: {self.reward_token}")
+
+        return {"logits": pooled_logits}
+
+
+def ordinal_predict(logits: np.ndarray, num_classes: int):
+    """
+    Convert CORN ordinal logits to predicted scores.
+
+    Args:
+        logits: [B, D, K-1] raw threshold logits
+        num_classes: K (number of ordinal classes)
+
+    Returns:
+        hard_preds: [B, D] integer predictions in {1..K}
+        soft_preds: [B, D] continuous expected value E[Y]
+    """
+    probs = 1.0 / (1.0 + np.exp(-logits))  # sigmoid → P(Y>k | Y>=k)
+    cum_probs = np.cumprod(probs, axis=-1)  # P(Y>k) = prod_{j<=k} P(Y>j|Y>=j)
+
+    hard_preds = (cum_probs > 0.5).sum(axis=-1) + 1  # [B, D]
+
+    cum_ext = np.concatenate([
+        np.ones((*cum_probs.shape[:-1], 1)),
+        cum_probs,
+        np.zeros((*cum_probs.shape[:-1], 1)),
+    ], axis=-1)
+    p_class = cum_ext[..., :-1] - cum_ext[..., 1:]
+    p_class = np.maximum(p_class, 0)
+    class_values = np.arange(1, num_classes + 1)
+    soft_preds = (p_class * class_values).sum(axis=-1)
+
+    return hard_preds, soft_preds

vefx_reward/prompt_template.py

@@ -0,0 +1,95 @@
+"""
+Prompt templates for VEFX-Reward video editing quality evaluation.
+"""
+
+EDITREWARD_V2_SPECIAL = """You are an expert evaluator assessing the quality of AI-generated video edits. You will be provided with two videos:
+- **Video 1**: The Original Video (before editing)
+- **Video 2**: The Edited Video (after editing)
+
+The editing instruction is:
+"{text_prompt}"
+
+Your task is to evaluate the Edited Video across THREE independent dimensions. Each dimension is scored on a 1–4 integer scale. **Scores across dimensions are independent** — a failure in one dimension must NOT affect scores in another.
+
+---
+
+## Dimension 1: Instructional Following (IF)
+**Core question:** Does the edited video accurately reflect the semantic requirements of the editing instruction?
+
+Evaluation criteria:
+- Object replacement: If the instruction says "replace apple with orange," did the model actually generate an orange (not a lemon or tomato)?
+- Action/attribute changes: If the instruction involves motion or attribute changes (e.g., "make it night"), was this correctly executed?
+- Completeness: Were ALL parts of the instruction addressed, not just partial execution?
+
+Scoring rubric:
+- **4 (Perfect):** The edit precisely and completely executes all instructions. Object categories, attributes (color, shape), actions, and styles all match the instruction with no ambiguity.
+- **3 (High):** The main instruction was executed, but minor details deviate. E.g., instruction asks for "red sports car" but a "red truck" was generated — the main concept "red car" is correct.
+- **2 (Low):** The main instruction was partially executed but with significant deviations, or completely irrelevant modifications were made.
+- **1 (Failed):** The edit has no relation to the instruction. E.g., instruction asks for "night scene" but the video remains daytime, or no change occurred at all.
+
+**Important notes:**
+- If the edit instruction asks for a camera perspective change (e.g., "shift to high angle") and the video shows no actual perspective change, score 1.
+- If the instruction asks for adding/increasing objects and no new objects appear, score 1.
+- A video that looks identical to the original (no edit happened) always scores 1 for IF.
+
+Instructional Following score (integer 1-4): <|IF_reward|>
+
+---
+
+## Dimension 2: Render Quality (RQ)
+**Core question:** What is the visual and temporal quality of the edited video?
+
+Evaluation criteria:
+- Naturalness and clarity: Are all parts of the video natural and sharp? Any blurriness, noise, or artifacts?
+- Physical plausibility: Does object motion obey physics? Any flickering, jittering, objects disappearing/morphing unexpectedly?
+- Temporal consistency: Is the video smooth frame-to-frame? Any sudden jumps, abrupt texture/color changes between frames?
+
+Scoring rubric:
+- **4 (Excellent):** Video clarity is very high with no visible defects, or only extremely minor artifacts detectable on very close inspection. Object motion fully obeys physics, smooth and natural. Visual quality is on par with or better than the original.
+- **3 (Medium):** Some quality degradation exists (e.g., slight blurring, localized flickering), but all objects remain clearly identifiable. The video's overall structure is intact despite imperfections.
+- **2 (Poor):** Significant quality degradation with obvious artifacts, distortion, or frame-to-frame inconsistency. Some object outlines deform, motion appears unnatural, affecting viewing experience.
+- **1 (Unusable):** Video quality completely breaks down. Objects are severely deformed or unrecognizable, serious physics violations (e.g., person walking through walls, objects shattering spontaneously), heavy noise or complete blur.
+
+**Important notes:**
+- A sudden scene transition mid-video (e.g., white background abruptly becoming a construction site) counts as a physics/consistency violation — score ≤ 3.
+- If the edit did NOT happen (original is preserved), RQ can still be high if the video itself looks fine — evaluate the video's visual quality independently.
+- Evaluate temporal artifacts carefully: a single frame of flickering is minor (score 3), persistent warping or morphing is severe (score 1-2).
+
+Render Quality score (integer 1-4): <|RQ_reward|>
+
+---
+
+## Dimension 3: Edit Exclusivity (EE)
+**Core question:** Did the model ONLY perform the specified edit, without making unintended changes to other parts of the video?
+
+Evaluation criteria:
+- Over-editing: When editing a foreground object, did the background, lighting, or other unrelated objects change?
+- Scene consistency: Are pixels, textures, and structures in non-edited regions preserved?
+- Camera trajectory: Was the original camera movement preserved? (Changing camera motion when not instructed is over-editing.)
+- Identity preservation: Do unedited people maintain their facial features, expressions, and body movements?
+
+Scoring rubric:
+- **4 (Perfect):** Strict exclusivity maintained. Only the target region specified by the instruction changed. All other regions (background, unrelated objects) remain identical to the original. Tiny pixel-level differences invisible to the eye are acceptable.
+- **3 (Medium):** Visible over-editing occurred. Non-target areas show noticeable changes, but overall scene layout and unrelated object consistency are still preserved. E.g., replaced a cup on a table but the table style also changed, or a background window disappeared.
+- **2 (Poor):** The overall scene or multiple unrelated objects changed significantly.
+- **1 (Complete failure):** No exclusivity at all. The entire video looks like a completely new video. The surrounding scene changed drastically, or more than three unrelated objects underwent serious alterations.
+
+**Important notes:**
+- Camera trajectory changes (when not instructed) are over-editing — if the original video had camera motion and the edited video is static (or vice versa), penalize EE.
+- For style transfer instructions (e.g., "turn into cyberpunk style"), it is expected that the entire visual style changes — this is NOT over-editing. But if text content or distinct object identities are destroyed during style transfer, that IS over-editing (score ≤ 3).
+- If the edit failed (IF=1) but the rest of the video also changed, EE should still be scored low.
+
+Edit Exclusivity score (integer 1-4): <|EE_reward|>
+"""
+
+
+def build_prompt(instruction: str) -> str:
+    """Build the evaluation prompt from an editing instruction.
+
+    Args:
+        instruction: The video editing instruction text.
+
+    Returns:
+        The formatted prompt string with special reward tokens.
+    """
+    return EDITREWARD_V2_SPECIAL.format(text_prompt=instruction)

vefx_reward/vision_process.py

@@ -0,0 +1,283 @@
+"""
+Video processing utilities for VEFX-Reward.
+Handles video loading, frame sampling, and resizing for Qwen3-VL input.
+Adapted from qwen-vl-utils (https://github.com/kq-chen/qwen-vl-utils).
+"""
+
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import sys
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+
+def round_by_factor(number: int, factor: int) -> int:
+    return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+    return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+    return math.floor(number / factor) * factor
+
+
+def smart_resize(
+    height: int, width: int, factor: int = IMAGE_FACTOR,
+    min_pixels: int = MIN_PIXELS, max_pixels: int = MAX_PIXELS,
+) -> tuple[int, int]:
+    """Resize dimensions to be divisible by factor while respecting pixel budget."""
+    if max(height, width) / min(height, width) > MAX_RATIO:
+        raise ValueError(
+            f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
+            f"got {max(height, width) / min(height, width)}"
+        )
+    h_bar = max(factor, round_by_factor(height, factor))
+    w_bar = max(factor, round_by_factor(width, factor))
+    if h_bar * w_bar > max_pixels:
+        beta = math.sqrt((height * width) / max_pixels)
+        h_bar = floor_by_factor(height / beta, factor)
+        w_bar = floor_by_factor(width / beta, factor)
+    elif h_bar * w_bar < min_pixels:
+        beta = math.sqrt(min_pixels / (height * width))
+        h_bar = ceil_by_factor(height * beta, factor)
+        w_bar = ceil_by_factor(width * beta, factor)
+    return h_bar, w_bar
+
+
+def smart_nframes(ele: dict, total_frames: int, video_fps: int | float) -> int:
+    """Calculate the number of frames to extract based on fps or nframes config."""
+    assert not ("fps" in ele and "nframes" in ele), "Only accept either `fps` or `nframes`"
+    if "nframes" in ele:
+        nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+    else:
+        fps = ele.get("fps", FPS)
+        min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+        max_frames = floor_by_factor(
+            ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)), FRAME_FACTOR
+        )
+        nframes = total_frames / video_fps * fps
+        nframes = min(max(nframes, min_frames), max_frames)
+        nframes = round_by_factor(nframes, FRAME_FACTOR)
+    if nframes > total_frames:
+        nframes = total_frames
+    if not (FRAME_FACTOR <= nframes <= total_frames):
+        raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
+    return nframes
+
+
+def _read_video_torchvision(ele: dict) -> tuple[torch.Tensor, dict]:
+    """Read video using torchvision.io.read_video. Returns (T, C, H, W) tensor."""
+    video_path = ele["video"]
+    if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+        if "http://" in video_path or "https://" in video_path:
+            warnings.warn("torchvision < 0.19.0 does not support http/https video path.")
+        if "file://" in video_path:
+            video_path = video_path[7:]
+    video, audio, info = io.read_video(
+        video_path,
+        start_pts=ele.get("video_start", 0.0),
+        end_pts=ele.get("video_end", None),
+        pts_unit="sec",
+        output_format="TCHW",
+    )
+    total_frames, video_fps = video.size(0), info["video_fps"]
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+    video = video[idx]
+    metadata = {
+        "total_num_frames": total_frames,
+        "fps": video_fps,
+        "frames_indices": idx,
+    }
+    return video, metadata
+
+
+def is_decord_available() -> bool:
+    import importlib.util
+    return importlib.util.find_spec("decord") is not None
+
+
+def _read_video_decord(ele: dict) -> tuple[torch.Tensor, dict]:
+    """Read video using decord.VideoReader. Returns (T, C, H, W) tensor."""
+    import decord
+    video_path = ele["video"]
+    vr = decord.VideoReader(video_path)
+    total_frames, video_fps = len(vr), vr.get_avg_fps()
+    nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+    idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+    video = vr.get_batch(idx).asnumpy()
+    video = torch.tensor(video).permute(0, 3, 1, 2)  # NHWC → TCHW
+    metadata = {
+        "total_num_frames": total_frames,
+        "fps": video_fps,
+        "frames_indices": idx,
+    }
+    return video, metadata
+
+
+VIDEO_READER_BACKENDS = {
+    "decord": _read_video_decord,
+    "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+    if FORCE_QWENVL_VIDEO_READER is not None:
+        video_reader_backend = FORCE_QWENVL_VIDEO_READER
+    elif is_decord_available():
+        video_reader_backend = "decord"
+    else:
+        video_reader_backend = "torchvision"
+    print(f"vefx-reward using {video_reader_backend} to read video.", file=sys.stderr)
+    return video_reader_backend
+
+
+def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
+    if "image" in ele:
+        image = ele["image"]
+    else:
+        image = ele["image_url"]
+    image_obj = None
+    if isinstance(image, Image.Image):
+        image_obj = image
+    elif image.startswith("http://") or image.startswith("https://"):
+        image_obj = Image.open(requests.get(image, stream=True).raw)
+    elif image.startswith("file://"):
+        image_obj = Image.open(image[7:])
+    elif image.startswith("data:image"):
+        if "base64," in image:
+            _, base64_data = image.split("base64,", 1)
+            data = base64.b64decode(base64_data)
+            image_obj = Image.open(BytesIO(data))
+    else:
+        image_obj = Image.open(image)
+    if image_obj is None:
+        raise ValueError(f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}")
+    image = image_obj.convert("RGB")
+    if "resized_height" in ele and "resized_width" in ele:
+        resized_height, resized_width = smart_resize(
+            ele["resized_height"], ele["resized_width"], factor=size_factor,
+        )
+    else:
+        width, height = image.size
+        min_pixels = ele.get("min_pixels", MIN_PIXELS)
+        max_pixels = ele.get("max_pixels", MAX_PIXELS)
+        resized_height, resized_width = smart_resize(
+            height, width, factor=size_factor, min_pixels=min_pixels, max_pixels=max_pixels,
+        )
+    image = image.resize((resized_width, resized_height))
+    return image
+
+
+def fetch_video(ele: dict, image_factor: int = IMAGE_FACTOR) -> tuple[torch.Tensor | list[Image.Image], dict | None]:
+    if isinstance(ele["video"], str):
+        video_reader_backend = get_video_reader_backend()
+        video, metadata = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+        nframes, _, height, width = video.shape
+        min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+        total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+        max_pixels = max(min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR), int(min_pixels * 1.05))
+        max_pixels = ele.get("max_pixels", max_pixels)
+        if "resized_height" in ele and "resized_width" in ele:
+            resized_height, resized_width = smart_resize(
+                ele["resized_height"], ele["resized_width"], factor=image_factor,
+            )
+        else:
+            resized_height, resized_width = smart_resize(
+                height, width, factor=image_factor,
+                min_pixels=min_pixels, max_pixels=max_pixels,
+            )
+        video = transforms.functional.resize(
+            video, [resized_height, resized_width],
+            interpolation=InterpolationMode.BICUBIC, antialias=True,
+        ).float()
+        return video, metadata
+    else:
+        assert isinstance(ele["video"], (list, tuple))
+        process_info = ele.copy()
+        process_info.pop("type", None)
+        process_info.pop("video", None)
+        images = [
+            fetch_image({"image": video_element, **process_info}, size_factor=image_factor)
+            for video_element in ele["video"]
+        ]
+        nframes = ceil_by_factor(len(images), FRAME_FACTOR)
+        if len(images) < nframes:
+            images.extend([images[-1]] * (nframes - len(images)))
+        return images, None
+
+
+def extract_vision_info(conversations: list[dict] | list[list[dict]]) -> list[dict]:
+    vision_infos = []
+    if isinstance(conversations[0], dict):
+        conversations = [conversations]
+    for conversation in conversations:
+        for message in conversation:
+            if isinstance(message["content"], list):
+                for ele in message["content"]:
+                    if (
+                        "image" in ele
+                        or "image_url" in ele
+                        or "video" in ele
+                        or ele["type"] in ("image", "image_url", "video")
+                    ):
+                        vision_infos.append(ele)
+    return vision_infos
+
+
+def process_vision_info(
+    conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] | None, list[dict] | None]:
+    """Process vision info from conversation messages, loading images and videos."""
+    vision_infos = extract_vision_info(conversations)
+    image_inputs = []
+    video_inputs = []
+    video_metadata_list = []
+    for vision_info in vision_infos:
+        if "image" in vision_info or "image_url" in vision_info:
+            image_inputs.append(fetch_image(vision_info))
+        elif "video" in vision_info:
+            video, metadata = fetch_video(vision_info)
+            video_inputs.append(video)
+            video_metadata_list.append(metadata)
+        else:
+            raise ValueError("image, image_url or video should in content.")
+    if len(image_inputs) == 0:
+        image_inputs = None
+    if len(video_inputs) == 0:
+        video_inputs = None
+        video_metadata_list = None
+    return image_inputs, video_inputs, video_metadata_list