Fix transformers version

README.md

@@ -1917,6 +1917,13 @@ open-stream replay path uses version-specific Thinker internals. On Linux,
 compatible NVIDIA GPU, optionally install `flash-attn` and pass
 `--flash-attention-2`.
 
+The separate B1 local extractor, `extract_turn_local_features.py`, is tested
+with both Transformers 4.52.3 and 4.57.6. An existing 4.57.6 GPU container does
+not need to be downgraded for local extraction. This does not change the replay
+API's dependency assumptions. Keep all local sidecars for one experiment on
+the same version; their metadata records the actual runtime version. See
+[the local feature guide](docs/local_feature_baseline.md).
+
 ## Ego4D Audio-Visual Diarization download
 
 Do not download the complete Ego4D release for this experiment. Download only

docs/local_feature_baseline.md

@@ -28,9 +28,21 @@ Qwen 3B 的 **2048 是编码器原生输出维度，256 只是新增 head 的可
 这保证相对于 readout 时刻的输入因果性，但不保证编码器在窗口内部逐 token 因果。
 
 实现依据：[Qwen 3B 配置](https://huggingface.co/Qwen/Qwen2.5-Omni-3B/blob/main/config.json)、
-[Transformers 4.52.3 编码器 API](https://github.com/huggingface/transformers/blob/v4.52.3/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py)。
+[Transformers 4.52.3 编码器 API](https://github.com/huggingface/transformers/blob/v4.52.3/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py)、
+[4.57.6 编码器 API](https://github.com/huggingface/transformers/blob/v4.57.6/src/transformers/models/qwen2_5_omni/modeling_qwen2_5_omni.py)。
 脚本加载 Thinker checkpoint 后释放语言解码器，仅将冻结的视觉／音频塔移到 GPU；没有 Thinker forward。
 
+局部提取已适配 **Transformers 4.52.3 和 4.57.6**，现有 L20 / ms-swift 容器若是 4.57.6，
+无需降级，也不要为这一步重新安装整个 `requirements.txt`。旧脚本只接受精确的 4.52.3，
+因此即使接口存在也会在加载模型前主动报错；更新脚本后重跑原 smoke 命令即可。
+两个版本都校验原生 token 输出形状与权重完整性，并显式使用 `use_fast=False` 固定图像预处理路径。
+未知版本仍在媒体解码前报错；`--preview-only` 不受编码器版本检查影响。
+该兼容范围只针对本局部提取脚本，不代表 Thinker KV replay 或 ROMA 的全部接口可互换。
+
+输出的 `config.transformers` 记录实际运行版本，不再写死 4.52.3。
+同一局部特征实验的 train/test sessions 应统一版本、权重和提取配置；加载器会拒绝混用版本。
+接口兼容不等于不同版本的浮点特征逐位相同，已有全局 Omni 特征无需重提。
+
 ## 1. EasyCom：先检查框，再跑一个 4s smoke test
 
 在已有项目依赖的 GPU 机器运行，以下路径沿用之前实验；raw 数据根按实际挂载修改。
@@ -83,6 +95,8 @@ python - <<'PY'
 import torch
 p = torch.load('features/local_smoke/Session_1.pt', map_location='cpu', weights_only=False)
 print('partial:', p['partial'])
+print('transformers:', p['config']['transformers'])
+print('box coordinates:', p['config']['box_coordinate_space'])
 for key in ('visual', 'audio'):
     x = p[key]
     print(key, tuple(x.shape), x.dtype, 'finite:', bool(torch.isfinite(x).all()),
@@ -94,9 +108,12 @@ PY
 
 3B 预期 V/A 各 `[20,2048]`，时间 0.2–4.0s；没有人脸的时刻允许 V 无效。
 本机已验证 CPU 对齐／融合／合成数据训练推理，真实 checkpoint + GPU 提取仍需此 smoke test。
-另有不下载权重的小型随机 Qwen 编码器 API 测试。可在已有依赖的环境运行：
+另有不下载权重的小型随机 Qwen 编码器 API 和本地 checkpoint 加载测试，
+在 4.52.3 / 4.57.6 上均适用；checkpoint 加载测试需要 `accelerate`。
+可先检查本机编码器接口，再跑完整回归：
 
 ```bash
+python -m unittest -v test_local_turn_features.TinyQwenEncoderTests
 python -m unittest -v test_easycom_head_objectives test_local_turn_features
 ```
 

extract_turn_local_features.py

@@ -26,6 +26,20 @@ from local_turn_features import (
 )
 
 
+TESTED_TRANSFORMERS_VERSIONS = ("4.52.3", "4.57.6")
+
+
+def validate_local_transformers_version(version: str) -> str:
+    """Keep the local encoder's tested API separate from the replay KV API."""
+    if version not in TESTED_TRANSFORMERS_VERSIONS:
+        supported = ", ".join(TESTED_TRANSFORMERS_VERSIONS)
+        raise RuntimeError(
+            f"Local encoder API is tested with transformers versions {supported}; found {version}. "
+            "Use a tested version or validate the encoder APIs before adding support."
+        )
+    return version
+
+
 def model_config_digest(config: dict[str, Any]) -> str:
     """Ignore location fields so an explicitly identified checkpoint can move."""
     def stable(value):
@@ -45,12 +59,13 @@ class FrozenQwenLocalEncoder:
         import transformers
         from transformers import Qwen2_5OmniProcessor, Qwen2_5OmniThinkerForConditionalGeneration
 
-        if transformers.__version__ != "4.52.3":
-            raise RuntimeError("Local encoder API is tested against transformers==4.52.3; use requirements.txt")
+        self.transformers_version = validate_local_transformers_version(transformers.__version__)
+        print(f"Local encoders: transformers={self.transformers_version}", flush=True)
         self.device = torch.device(args.device)
         if self.device.type == "cuda" and not torch.cuda.is_available():
             raise RuntimeError(f"CUDA requested but unavailable: {args.device}")
-        kwargs = {"torch_dtype": getattr(torch, args.dtype), "device_map": {"": "cpu"},
+        dtype_key = "dtype" if self.transformers_version == "4.57.6" else "torch_dtype"
+        kwargs = {dtype_key: getattr(torch, args.dtype), "device_map": {"": "cpu"},
                   "output_loading_info": True}
         if args.flash_attention_2:
             kwargs["attn_implementation"] = "flash_attention_2"
@@ -70,7 +85,9 @@ class FrozenQwenLocalEncoder:
         model.visual.to(self.device)
         model.audio_tower.to(self.device)
         self.model = model
-        self.processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id)
+        # Keep PIL/slow image preprocessing consistent with the original B1
+        # definition, independently of newer Transformers' fast defaults.
+        self.processor = Qwen2_5OmniProcessor.from_pretrained(args.model_id, use_fast=False)
         self.max_pixels = args.max_pixels
         self.image_batch_size = args.image_batch_size
 
@@ -87,8 +104,9 @@ class FrozenQwenLocalEncoder:
             tokens = self.model.get_image_features(pixels, grid)
             merge = int(self.model.config.vision_config.spatial_merge_size)
             counts = (grid.prod(-1) // (merge * merge)).tolist()
-            if tokens.ndim != 2 or tokens.shape != (sum(counts), self.visual_dim):
-                raise RuntimeError(f"Unexpected image encoder shape: {tuple(tokens.shape)}")
+            if not isinstance(tokens, torch.Tensor) or tokens.shape != (sum(counts), self.visual_dim):
+                raise RuntimeError(f"Unexpected image encoder output: {type(tokens).__name__}, "
+                                   f"shape={getattr(tokens, 'shape', None)}")
             # Split BEFORE spatial mean: each person has one vote regardless of
             # box size / number of visual tokens.
             pieces.extend(part.float().mean(0).cpu() for part in tokens.split(counts))
@@ -108,8 +126,10 @@ class FrozenQwenLocalEncoder:
         mask = batch["attention_mask"].to(self.device)
         tokens = self.model.get_audio_features(features, feature_attention_mask=mask)
         _, lengths = self.model.audio_tower._get_feat_extract_output_lengths(mask.sum(-1))
-        if tokens.ndim != 2 or tokens.shape != (int(lengths.sum()), self.audio_dim) or not len(tokens):
-            raise RuntimeError(f"Unexpected valid audio encoder shape: {tuple(tokens.shape)}")
+        if (not isinstance(tokens, torch.Tensor)
+                or tokens.shape != (int(lengths.sum()), self.audio_dim) or not len(tokens)):
+            raise RuntimeError(f"Unexpected valid audio encoder output: {type(tokens).__name__}, "
+                               f"shape={getattr(tokens, 'shape', None)}")
         # get_audio_features removes padding and verifies the output length.
         return tokens.float().mean(0).cpu()
 
@@ -371,6 +391,11 @@ def parse_args(argv=None):
 
 def main() -> None:
     args = parse_args()
+    if not args.preview_only:
+        import transformers
+
+        # Fail before indexing/decoding source media, not after expensive I/O.
+        validate_local_transformers_version(transformers.__version__)
     names = [normalize_sample_name(value, args.dataset) for value in args.samples]
     if len(set(names)) != len(names):
         raise ValueError("Duplicate samples")
@@ -417,7 +442,7 @@ def main() -> None:
             "config": {
                 "encoder_id": args.encoder_id or args.model_id,
                 "model_config_sha256": encoder.model_config_sha256,
-                "transformers": "4.52.3", "dtype": args.dtype,
+                "transformers": encoder.transformers_version, "dtype": args.dtype,
                 "attention": "flash_attention_2" if args.flash_attention_2 else "default",
                 "visual_dim": encoder.visual_dim, "audio_dim": encoder.audio_dim,
                 "visual_fps": args.video_fps, "visual_window_seconds": args.visual_window_seconds,

requirements.txt

@@ -1,5 +1,7 @@
 torch
 numpy
+# Default for Thinker replay. The B1 local extractor also supports 4.57.6;
+# keep that version in existing GPU containers (see docs/local_feature_baseline.md).
 transformers==4.52.3
 accelerate
 qwen-omni-utils[decord]

test_local_turn_features.py

@@ -12,6 +12,7 @@ import sys
 import tempfile
 import unittest
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import numpy as np
 import torch
@@ -23,6 +24,8 @@ from extract_turn_local_features import (
     validate_source_timeline,
     model_config_digest,
     box_coordinate_offset, frame_crops, parse_args, preview_boxes,
+    TESTED_TRANSFORMERS_VERSIONS, validate_local_transformers_version,
+    main as extract_main,
 )
 from local_turn_features import (
     LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes, load_local_features,
@@ -52,6 +55,13 @@ class LocalTests(unittest.TestCase):
     def setUp(self):
         torch.set_num_threads(1)
 
+    def test_tested_transformers_versions(self):
+        for version in ("4.52.3", "4.57.6"):
+            self.assertEqual(validate_local_transformers_version(version), version)
+        for version in ("4.50.0.dev0", "4.57.6.dev0", "5.0.0"):
+            with self.subTest(version=version), self.assertRaisesRegex(RuntimeError, "found " + version):
+                validate_local_transformers_version(version)
+
     def test_half_open_grid_and_pooling(self):
         self.assertEqual(list(causal_frame_ticks(.2, 2)), [0])
         self.assertEqual(list(causal_frame_ticks(1., 2)), [0, 1])
@@ -263,6 +273,67 @@ class LocalTests(unittest.TestCase):
             get_image_features=lambda *_: torch.tensor([[1., 1.], [3., 3.], [5., 5.]]))
         actual = encoder.images([object(), object()])
         self.assertTrue(torch.equal(actual, torch.tensor([[1., 1.], [4., 4.]])))
+        # A future API returning per-image tensors must not silently change
+        # the pooling rule or crash with an uninformative AttributeError.
+        encoder.model.get_image_features = lambda *_: (torch.ones(1, 2), torch.ones(2, 2))
+        with self.assertRaisesRegex(RuntimeError, "Unexpected image encoder output"):
+            encoder.images([object(), object()])
+
+    def test_extraction_saves_runtime_version_and_rejects_mixed_versions(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            for folder in ("features", "targets", "local"):
+                (root / folder).mkdir()
+            reference, side, labels = fixture()
+            torch.save(reference, root / "features/Session_4.pt")
+            np.save(root / "targets/Session_4.npy", labels)
+            media = SimpleNamespace(clips=[], metadata=lambda: {})
+            argv = ["extract_turn_local_features.py", "--data-root", str(root),
+                    "--feature-root", str(root / "features"), "--target-root", str(root / "targets"),
+                    "--output-root", str(root / "local"), "--samples", "4", "--overwrite"]
+            contracts = []
+            for version in TESTED_TRANSFORMERS_VERSIONS:
+                fake_encoder = SimpleNamespace(transformers_version=version, visual_dim=6, audio_dim=5,
+                                               model_config_sha256="same-checkpoint")
+                with patch.object(sys, "argv", argv), \
+                     patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__=version)}), \
+                     patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])), \
+                     patch("extract_turn_local_features.FrozenQwenLocalEncoder", return_value=fake_encoder), \
+                     patch("extract_turn_local_features.extract_local", return_value=copy.deepcopy(side)):
+                    extract_main()
+                output = root / "local/Session_4.pt"
+                _, contract = load_local_features(output, reference, sample_id="Session_4",
+                                                   dataset="easycom", modalities="av")
+                self.assertEqual(contract["config"]["transformers"], version)
+                contracts.append(contract)
+            with self.assertRaisesRegex(ValueError, "contract differs"):
+                load_local_features(output, reference, sample_id="Session_4", dataset="easycom",
+                                    modalities="av", expected_contract=contracts[0])
+
+    def test_unsupported_version_fails_before_media_and_preview_still_works(self):
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            reference, _, labels = fixture()
+            torch.save(reference, root / "Session_4.pt")
+            np.save(root / "Session_4.npy", labels)
+            argv = ["extract_turn_local_features.py", "--data-root", str(root),
+                    "--feature-root", str(root), "--target-root", str(root),
+                    "--output-root", str(root / "preview"), "--samples", "4"]
+            with patch.object(sys, "argv", argv), \
+                 patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="5.0.0")}), \
+                 patch("extract_turn_local_features.build_media_and_boxes") as build:
+                with self.assertRaisesRegex(RuntimeError, "found 5.0.0"):
+                    extract_main()
+                build.assert_not_called()
+            media = SimpleNamespace(clips=[], metadata=lambda: {})
+            with patch.object(sys, "argv", argv + ["--preview-only"]), \
+                 patch.dict(sys.modules, {"transformers": SimpleNamespace(__version__="5.0.0")}), \
+                 patch("extract_turn_local_features.build_media_and_boxes", return_value=(media, [])), \
+                 patch("extract_turn_local_features.preview_boxes") as preview, \
+                 patch("extract_turn_local_features.FrozenQwenLocalEncoder") as encoder:
+                extract_main()
+                preview.assert_called_once()
+                encoder.assert_not_called()
 
     def test_extraction_never_encodes_future_audio_and_saves_ragged_people(self):
         class Media:
@@ -294,17 +365,19 @@ class LocalTests(unittest.TestCase):
         self.assertTrue(result["audio_valid"].all())
 
 
-@unittest.skipUnless(importlib.util.find_spec("transformers"), "optional: install transformers==4.52.3")
+@unittest.skipUnless(importlib.util.find_spec("transformers"), "optional: install a tested Transformers version")
 class TinyQwenEncoderTests(unittest.TestCase):
-    def test_real_encoder_apis_without_downloading_weights(self):
+    def setUp(self):
         import transformers
-        if transformers.__version__ != "4.52.3":
-            self.skipTest("Pinned encoder API test requires transformers==4.52.3")
+        if transformers.__version__ not in TESTED_TRANSFORMERS_VERSIONS:
+            self.skipTest(f"Encoder API tests require one of {TESTED_TRANSFORMERS_VERSIONS}")
+        torch.set_num_threads(1)
+
+    def tiny_model_and_processor(self):
         from transformers import (
             Qwen2_5OmniThinkerConfig, Qwen2_5OmniThinkerForConditionalGeneration,
             Qwen2VLImageProcessor, WhisperFeatureExtractor,
         )
-        torch.set_num_threads(1)
         config = Qwen2_5OmniThinkerConfig(
             audio_config=dict(d_model=16, encoder_layers=1, encoder_attention_heads=2,
                               encoder_ffn_dim=32, output_dim=8),
@@ -313,14 +386,18 @@ class TinyQwenEncoderTests(unittest.TestCase):
             text_config=dict(vocab_size=32, hidden_size=8, intermediate_size=16,
                              num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=2),
         )
+        model = Qwen2_5OmniThinkerForConditionalGeneration(config).eval().requires_grad_(False)
+        processor = SimpleNamespace(image_processor=Qwen2VLImageProcessor(),
+                                    feature_extractor=WhisperFeatureExtractor(feature_size=128))
+        return model, processor
+
+    def test_real_encoder_apis_without_downloading_weights(self):
         encoder = FrozenQwenLocalEncoder.__new__(FrozenQwenLocalEncoder)
-        encoder.model = Qwen2_5OmniThinkerForConditionalGeneration(config).eval().requires_grad_(False)
+        encoder.model, encoder.processor = self.tiny_model_and_processor()
         encoder.model.model = None
         encoder.model.lm_head = None
         encoder.visual_dim = encoder.audio_dim = 8
         encoder.device, encoder.max_pixels, encoder.image_batch_size = torch.device("cpu"), 224 * 224, 8
-        encoder.processor = SimpleNamespace(image_processor=Qwen2VLImageProcessor(),
-                                            feature_extractor=WhisperFeatureExtractor(feature_size=128))
         crops = [Image.new("RGB", (56, 84), "red"), Image.new("RGB", (112, 56), "blue")]
         batch = encoder.images(crops)
         single = encoder.images(crops[:1])
@@ -332,6 +409,44 @@ class TinyQwenEncoderTests(unittest.TestCase):
             self.assertTrue(torch.isfinite(audio).all())
             self.assertFalse(audio.requires_grad)
 
+    @unittest.skipUnless(importlib.util.find_spec("accelerate"), "optional: install accelerate for checkpoint loading")
+    def test_constructor_loads_local_checkpoint_and_preserves_encoder_weights(self):
+        import transformers
+
+        model, processor = self.tiny_model_and_processor()
+        expected = {key: value.clone() for key, value in model.state_dict().items()
+                    if key.startswith(("visual.", "audio_tower."))}
+        with tempfile.TemporaryDirectory() as directory:
+            model.save_pretrained(directory)
+            args = SimpleNamespace(model_id=directory, device="cpu", dtype="float32",
+                                   flash_attention_2=False, max_pixels=224 * 224, image_batch_size=8)
+            with patch.object(transformers.Qwen2_5OmniProcessor, "from_pretrained", return_value=processor) as load:
+                encoder = FrozenQwenLocalEncoder(args)
+            load.assert_called_once_with(directory, use_fast=False)
+            self.assertEqual(encoder.transformers_version, transformers.__version__)
+            self.assertIsNone(encoder.model.model)
+            self.assertIsNone(encoder.model.lm_head)
+            self.assertFalse(encoder.model.training)
+            self.assertFalse(any(p.requires_grad for p in encoder.model.parameters()))
+            for key, value in encoder.model.state_dict().items():
+                self.assertTrue(torch.equal(value, expected[key]), key)
+            self.assertEqual(encoder.images([Image.new("RGB", (56, 84), "red")]).shape, (1, 8))
+            self.assertEqual(encoder.audio(np.zeros(3200, dtype=np.float32)).shape, (8,))
+
+    def test_constructor_rejects_incomplete_encoder_weights(self):
+        import transformers
+
+        args = SimpleNamespace(model_id="unused", device="cpu", dtype="float32", flash_attention_2=False)
+        for info in ({"missing_keys": ["visual.patch_embed.proj.weight"]},
+                     {"missing_keys": ["audio_tower.conv1.weight"]},
+                     {"mismatched_keys": [("visual.weight", (2,), (3,))]},
+                     {"error_msgs": ["broken checkpoint"]}):
+            with self.subTest(info=info), \
+                 patch.object(transformers.Qwen2_5OmniThinkerForConditionalGeneration,
+                              "from_pretrained", return_value=(None, info)), \
+                 self.assertRaisesRegex(RuntimeError, "incompletely loaded weights"):
+                FrozenQwenLocalEncoder(args)
+
 
 class RoundTripTests(unittest.TestCase):
     def test_train_infer_easycom_and_ego4d(self):