From 389c0789332e3e3e3e0ed6bcc21f09212e7f2f25 Mon Sep 17 00:00:00 2001
From: Molly Xu <mollyxu@fb.com>
Date: Mon, 20 Oct 2025 14:03:01 -0700
Subject: [PATCH] fallback to container duration in approximate mode when
 stream duration is unavailable

---
 src/torchcodec/_core/SingleStreamDecoder.cpp | 10 ++++++++++
 test/test_decoders.py                        | 11 -----------
 test/utils.py                                |  7 -------
 3 files changed, 10 insertions(+), 18 deletions(-)

diff --git a/src/torchcodec/_core/SingleStreamDecoder.cpp b/src/torchcodec/_core/SingleStreamDecoder.cpp
index 2fbc111c1..e21cca49b 100644
--- a/src/torchcodec/_core/SingleStreamDecoder.cpp
+++ b/src/torchcodec/_core/SingleStreamDecoder.cpp
@@ -158,6 +158,16 @@ void SingleStreamDecoder::initializeDecoder() {
         ptsToSeconds(formatContext_->duration, defaultTimeBase);
   }
 
+  // Use container duration as fallback for streams missing stream duration
+  if (containerMetadata_.durationSecondsFromHeader.has_value()) {
+    for (auto& streamMetadata : containerMetadata_.allStreamMetadata) {
+      if (!streamMetadata.durationSecondsFromHeader.has_value()) {
+        streamMetadata.durationSecondsFromHeader =
+            containerMetadata_.durationSecondsFromHeader;
+      }
+    }
+  }
+
   if (formatContext_->bit_rate > 0) {
     containerMetadata_.bitRate = formatContext_->bit_rate;
   }
diff --git a/test/test_decoders.py b/test/test_decoders.py
index 098e4e969..2000a1572 100644
--- a/test/test_decoders.py
+++ b/test/test_decoders.py
@@ -46,7 +46,6 @@
     SINE_MONO_S32,
     SINE_MONO_S32_44100,
     SINE_MONO_S32_8000,
-    supports_approximate_mode,
     TEST_SRC_2_720P,
     TEST_SRC_2_720P_H265,
     TEST_SRC_2_720P_MPEG4,
@@ -1492,8 +1491,6 @@ def test_get_frames_at_tensor_indices(self):
     def test_beta_cuda_interface_get_frame_at(
         self, asset, contiguous_indices, seek_mode
     ):
-        if seek_mode == "approximate" and not supports_approximate_mode(asset):
-            pytest.skip("asset doesn't work with approximate mode")
 
         if in_fbcode() and asset is AV1_VIDEO:
             pytest.skip("AV1 CUDA not supported internally")
@@ -1540,8 +1537,6 @@ def test_beta_cuda_interface_get_frame_at(
     def test_beta_cuda_interface_get_frames_at(
         self, asset, contiguous_indices, seek_mode
     ):
-        if seek_mode == "approximate" and not supports_approximate_mode(asset):
-            pytest.skip("asset doesn't work with approximate mode")
         if in_fbcode() and asset is AV1_VIDEO:
             pytest.skip("AV1 CUDA not supported internally")
 
@@ -1585,8 +1580,6 @@ def test_beta_cuda_interface_get_frames_at(
     )
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_get_frame_played_at(self, asset, seek_mode):
-        if seek_mode == "approximate" and not supports_approximate_mode(asset):
-            pytest.skip("asset doesn't work with approximate mode")
         if in_fbcode() and asset is AV1_VIDEO:
             pytest.skip("AV1 CUDA not supported internally")
 
@@ -1627,8 +1620,6 @@ def test_beta_cuda_interface_get_frame_played_at(self, asset, seek_mode):
     )
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_get_frames_played_at(self, asset, seek_mode):
-        if seek_mode == "approximate" and not supports_approximate_mode(asset):
-            pytest.skip("asset doesn't work with approximate mode")
         if in_fbcode() and asset is AV1_VIDEO:
             pytest.skip("AV1 CUDA not supported internally")
 
@@ -1670,8 +1661,6 @@ def test_beta_cuda_interface_get_frames_played_at(self, asset, seek_mode):
     )
     @pytest.mark.parametrize("seek_mode", ("exact", "approximate"))
     def test_beta_cuda_interface_backwards(self, asset, seek_mode):
-        if seek_mode == "approximate" and not supports_approximate_mode(asset):
-            pytest.skip("asset doesn't work with approximate mode")
         if in_fbcode() and asset is AV1_VIDEO:
             pytest.skip("AV1 CUDA not supported internally")
 
diff --git a/test/utils.py b/test/utils.py
index b59681b37..08e120c06 100644
--- a/test/utils.py
+++ b/test/utils.py
@@ -798,10 +798,3 @@ def sample_format(self) -> str:
     },
     frames={0: {}},  # Not needed for now
 )
-
-
-def supports_approximate_mode(asset: TestVideo) -> bool:
-    # Those are missing the `duration` field so they fail in approximate mode (on all devices).
-    # TODO: we should address this, see
-    # https://github.com/meta-pytorch/torchcodec/issues/945
-    return asset not in (AV1_VIDEO, TEST_SRC_2_720P_VP9, TEST_SRC_2_720P_VP8)