sync custom_st.py from harness/src/hf_sources/nano_variant

custom_st.py

@@ -419,34 +419,6 @@ def _build_audio_model_inputs(owner, audio_input, device, prefix: str = ""):
     return inputs
 
 
-def _extract_audio_from_video(video_path):
-    """Return mono float32 audio @ 16 kHz decoded from the video's audio track, or
-    None if no audio stream is present. PyAV is already a dep for video decoding."""
-    try:
-        import av
-        import numpy as np
-        from av.audio.resampler import AudioResampler
-    except ImportError:
-        return None
-    container = av.open(video_path)
-    try:
-        audio_stream = next((s for s in container.streams if s.type == "audio"), None)
-        if audio_stream is None:
-            return None
-        resampler = AudioResampler(format="flt", layout="mono", rate=16000)
-        samples = []
-        for frame in container.decode(audio=0):
-            for rf in resampler.resample(frame):
-                samples.append(rf.to_ndarray().flatten())
-        for rf in resampler.resample(None):
-            samples.append(rf.to_ndarray().flatten())
-        if not samples:
-            return None
-        return np.concatenate(samples).astype(np.float32)
-    finally:
-        container.close()
-
-
 def _eval_video_frames(video_path):
     if not isinstance(video_path, str):
         return video_path
@@ -692,7 +664,7 @@ class Transformer(nn.Module):
         if first_kind == "image":
             return {"_images": values, "_is_image_batch": True}
         if first_kind == "video":
-            return {"_video_paths": values, "_is_video_batch": True}
+            return {"_multipart_batch": [(v,) for v in values], "_is_multipart_batch": True}
         if first_kind == "audio":
             return {"_audio_paths": values, "_is_audio_batch": True}
         if first_kind == "pdf":
@@ -727,18 +699,6 @@ class Transformer(nn.Module):
             hidden = self.model(**inputs).last_hidden_state
         return self._last_token_pool(hidden, inputs["attention_mask"]).squeeze(0)
 
-    def _encode_single_video(self, video_path, device) -> torch.Tensor:
-        video = _eval_video_frames(video_path)
-        inputs = self.processor(videos=video, text=VIDEO_PROMPT, return_tensors="pt", truncation=False)
-        inputs = {k: v.to(device) for k, v in inputs.items() if torch.is_tensor(v)}
-        if "pixel_values_videos" in inputs:
-            inputs["pixel_values"] = inputs.pop("pixel_values_videos")
-        if "video_grid_thw" in inputs:
-            inputs["image_grid_thw"] = inputs.pop("video_grid_thw")
-        with torch.no_grad():
-            hidden = self.model(**inputs).last_hidden_state
-        return self._last_token_pool(hidden, inputs["attention_mask"]).squeeze(0)
-
     def _encode_single_audio(self, audio_input, device, prefix: str = "") -> torch.Tensor:
         inputs = _build_audio_model_inputs(self, audio_input, device, prefix=prefix)
         with torch.no_grad():
@@ -875,9 +835,9 @@ class Transformer(nn.Module):
 
         Each part may be a URL, a local path (sniffed by magic bytes if no
         extension), a PIL.Image, a 1-D numpy audio array, a PDF (rasterised
-        to one image per page), or plain text. A video with an audio track
-        is auto-expanded to [extracted_audio, video] so the audio tokens
-        precede the video tokens.
+        to one image per page), or plain text. A video contributes its frames
+        only; pass the soundtrack as its own part to fuse both, e.g.
+        `encode(("narration.wav", "clip.mp4"))`.
         """
         import numpy as np
         from transformers import WhisperFeatureExtractor
@@ -885,17 +845,10 @@ class Transformer(nn.Module):
         # Normalize every part first (URL -> path, content-sniff if needed).
         resolved = [_resolve_input(p) for p in parts]
 
-        # Expand videos-with-audio: prepend extracted audio.
         # Expand PDFs: rasterise into one image-part per page.
         expanded = []
         for kind, value in resolved:
-            if kind == "video":
-                if isinstance(value, str):
-                    aud = _extract_audio_from_video(value)
-                    if aud is not None and aud.size > 0:
-                        expanded.append(("audio", aud))
-                expanded.append(("video", value))
-            elif kind == "pdf":
+            if kind == "pdf":
                 for page in _pdf_to_images(value):
                     expanded.append(("image", page))
             else:
@@ -937,11 +890,6 @@ class Transformer(nn.Module):
             features["sentence_embedding"] = torch.stack(embs)
             return self._maybe_truncate(features, truncate_dim)
 
-        if features.get("_is_video_batch"):
-            embs = [self._encode_single_video(p, device) for p in features["_video_paths"]]
-            features["sentence_embedding"] = torch.stack(embs)
-            return self._maybe_truncate(features, truncate_dim)
-
         if features.get("_is_audio_batch"):
             audio_items = features.get("_audio_paths") or features.get("_audio_wrappers", [])
             embs = [self._encode_single_audio(a, device) for a in audio_items]