Fix transformers version

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,