DKmode22/YuE2-3B-FP8

1#!/usr/bin/env python32"""Idempotent patch for the installed yue2_infer `fast.py` (vLLM backend).3 4Adds ONE lever: if the env var YUE2_AR_CHECKPOINT names a directory, the vLLM5worker serves THAT checkpoint instead of the bf16 AR checkpoint it derives from6model.safetensors. That is how an NVFP4 (compressed-tensors) requant of the7derived Qwen3-shaped AR checkpoint is put on the real generation path. Nothing8else changes (dtype, KV sizing from config.json, logits processor, max_num_seqs=1).9 10Usage: patch_fast.py <site-packages>/yue2/fast.py [--check]11"""12from __future__ import annotations13 14import re15import sys16 17MARK = "# s5-patch: YUE2_AR_CHECKPOINT override (services/yue2-nvfp4/patch_fast.py)"18OLD = 'derived = derive_ar_checkpoint(setup["model_dir"])'19NEW = (20    MARK + "\n"21    '    _override = os.environ.get("YUE2_AR_CHECKPOINT")\n'22    '    derived = Path(_override) if _override else derive_ar_checkpoint(setup["model_dir"])\n'23    '    if _override and not (derived / "config.json").exists():\n'24    '        raise FileNotFoundError(f"YUE2_AR_CHECKPOINT has no config.json: {derived}")\n'25    '    print(f"s5-patch: AR checkpoint = {derived} (override={bool(_override)})", file=sys.stderr, flush=True)'26)27 28 29def main() -> int:30    path = sys.argv[1]31    check = "--check" in sys.argv32    src = open(path).read()33    if MARK in src:34        print("already patched")35        return 036    if check:37        print("NOT patched")38        return 139    if src.count(OLD) != 1:40        print(f"expected exactly one occurrence of {OLD!r}, found {src.count(OLD)}")41        return 242    # the derive call is indented 4 spaces inside _worker_main43    new_src = re.sub(r"^(\s+)" + re.escape(OLD) + r"$",44                     lambda m: m.group(1) + NEW.replace("\n    ", "\n" + m.group(1)), src, count=1, flags=re.M)45    if new_src == src:46        print("substitution failed")47        return 348    if "from pathlib import Path" not in new_src and "import Path" not in new_src:49        print("fast.py has no Path import; refusing")50        return 451    open(path, "w").write(new_src)52    print("patched", path)53    return 054 55 56if __name__ == "__main__":57    sys.exit(main())58