Pass initializationData to EncoderFactory when requested

PiperOrigin-RevId: 580098432

libraries/transformer/src/main/java/androidx/media3/transformer/AudioSampleExporter.java

@@ -55,7 +55,8 @@ import org.checkerframework.dataflow.qual.Pure;
       AudioMixer.Factory mixerFactory,
       Codec.EncoderFactory encoderFactory,
       MuxerWrapper muxerWrapper,
-      FallbackListener fallbackListener)
+      FallbackListener fallbackListener,
+      boolean matchInitializationData)
       throws ExportException {
     super(firstAssetLoaderTrackFormat, muxerWrapper);
     audioGraph = new AudioGraph(mixerFactory);
@@ -74,6 +75,8 @@ import org.checkerframework.dataflow.qual.Pure;
             .setChannelCount(encoderInputAudioFormat.channelCount)
             .setPcmEncoding(encoderInputAudioFormat.encoding)
             .setAverageBitrate(DEFAULT_ENCODER_BITRATE)
+            .setInitializationData(
+                matchInitializationData ? firstInputFormat.initializationData : null)
             .build();
 
     encoder =

libraries/transformer/src/main/java/androidx/media3/transformer/Transformer.java

@@ -1206,7 +1206,8 @@ public final class Transformer {
             applicationHandler,
             debugViewProvider,
             clock,
-            initialTimestampOffsetUs);
+            initialTimestampOffsetUs,
+            /* matchInitializationData= */ false);
     transformerInternal.start();
   }
 

libraries/transformer/src/main/java/androidx/media3/transformer/TransformerInternal.java

@@ -129,6 +129,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   private final List<SampleExporter> sampleExporters;
   private final Object setMaxSequenceDurationUsLock;
   private final MuxerWrapper muxerWrapper;
+  private final boolean matchInitializationData;
   private final ConditionVariable transformerConditionVariable;
 
   private boolean isDrainingExporters;
@@ -153,7 +154,8 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
       HandlerWrapper applicationHandler,
       DebugViewProvider debugViewProvider,
       Clock clock,
-      long videoSampleTimestampOffsetUs) {
+      long videoSampleTimestampOffsetUs,
+      boolean matchInitializationData) {
     this.context = context;
     this.composition = composition;
     this.encoderFactory = new CapturingEncoderFactory(encoderFactory);
@@ -162,6 +164,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     this.clock = clock;
     this.videoSampleTimestampOffsetUs = videoSampleTimestampOffsetUs;
     this.muxerWrapper = muxerWrapper;
+    this.matchInitializationData = matchInitializationData;
     internalHandlerThread = new HandlerThread("Transformer:Internal");
     internalHandlerThread.start();
     sequenceAssetLoaders = new ArrayList<>();
@@ -582,7 +585,8 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
                 audioMixerFactory,
                 encoderFactory,
                 muxerWrapper,
-                fallbackListener));
+                fallbackListener,
+                matchInitializationData));
       } else {
         // TODO(b/267301878): Pass firstAssetLoaderOutputFormat once surface creation not in VSP.
         assetLoaderInputTracker.registerSampleExporter(
@@ -600,8 +604,8 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
                 fallbackListener,
                 debugViewProvider,
                 videoSampleTimestampOffsetUs,
-                /* hasMultipleInputs= */ assetLoaderInputTracker
+                /* hasMultipleInputs= */ assetLoaderInputTracker.hasMultipleConcurrentVideoTracks(),
-                    .hasMultipleConcurrentVideoTracks()));
+                matchInitializationData));
       }
     }
 

libraries/transformer/src/main/java/androidx/media3/transformer/VideoSampleExporter.java

@@ -88,7 +88,8 @@ import org.checkerframework.dataflow.qual.Pure;
       FallbackListener fallbackListener,
       DebugViewProvider debugViewProvider,
       long initialTimestampOffsetUs,
-      boolean hasMultipleInputs)
+      boolean hasMultipleInputs,
+      boolean matchInitializationData)
       throws ExportException {
     // TODO(b/278259383) Consider delaying configuration of VideoSampleExporter to use the decoder
     //  output format instead of the extractor output format, to match AudioSampleExporter behavior.
@@ -109,7 +110,8 @@ import org.checkerframework.dataflow.qual.Pure;
             firstInputFormat.buildUpon().setColorInfo(decoderInputColor).build(),
             muxerWrapper.getSupportedSampleMimeTypes(C.TRACK_TYPE_VIDEO),
             transformationRequest,
-            fallbackListener);
+            fallbackListener,
+            matchInitializationData);
     encoderOutputBuffer =
         new DecoderInputBuffer(DecoderInputBuffer.BUFFER_REPLACEMENT_MODE_DISABLED);
 
@@ -239,6 +241,7 @@ import org.checkerframework.dataflow.qual.Pure;
     private final FallbackListener fallbackListener;
     private final String requestedOutputMimeType;
     private final @Composition.HdrMode int hdrModeAfterFallback;
+    private final boolean matchInitializationData;
 
     private @MonotonicNonNull SurfaceInfo encoderSurfaceInfo;
 
@@ -251,13 +254,15 @@ import org.checkerframework.dataflow.qual.Pure;
         Format inputFormat,
         List<String> muxerSupportedMimeTypes,
         TransformationRequest transformationRequest,
-        FallbackListener fallbackListener) {
+        FallbackListener fallbackListener,
+        boolean matchInitializationData) {
       checkArgument(inputFormat.colorInfo != null);
       this.encoderFactory = encoderFactory;
       this.inputFormat = inputFormat;
       this.muxerSupportedMimeTypes = muxerSupportedMimeTypes;
       this.transformationRequest = transformationRequest;
       this.fallbackListener = fallbackListener;
+      this.matchInitializationData = matchInitializationData;
       Pair<String, Integer> outputMimeTypeAndHdrModeAfterFallback =
           getRequestedOutputMimeTypeAndHdrModeAfterFallback(inputFormat, transformationRequest);
       requestedOutputMimeType = outputMimeTypeAndHdrModeAfterFallback.first;
@@ -332,6 +337,8 @@ import org.checkerframework.dataflow.qual.Pure;
               .setFrameRate(inputFormat.frameRate)
               .setSampleMimeType(requestedOutputMimeType)
               .setColorInfo(getSupportedInputColor())
+              .setInitializationData(
+                  matchInitializationData ? inputFormat.initializationData : null)
               .build();
 
       encoder =

libraries/transformer/src/test/java/androidx/media3/transformer/VideoEncoderWrapperTest.java

@@ -68,7 +68,8 @@ public final class VideoEncoderWrapperTest {
               .build(),
           /* muxerSupportedMimeTypes= */ ImmutableList.of(MimeTypes.VIDEO_H264),
           emptyTransformationRequest,
-          fallbackListener);
+          fallbackListener,
+          /* matchInitializationData= */ false);
 
   @Before
   public void setUp() {
@@ -134,6 +135,52 @@ public final class VideoEncoderWrapperTest {
     assertThat(surfaceInfo.height).isEqualTo(fallbackHeight);
   }
 
+  @Test
+  public void matchInitializationData_setToFalse_initializationDataNotPassedToEncoderFactory()
+      throws Exception {
+    VideoSampleExporter.EncoderWrapper encoderWrapper =
+        new VideoSampleExporter.EncoderWrapper(
+            fakeEncoderFactory,
+            /* inputFormat= */ new Format.Builder()
+                .setSampleMimeType(MimeTypes.VIDEO_H264)
+                .setColorInfo(ColorInfo.SDR_BT709_LIMITED)
+                .setInitializationData(ImmutableList.of(new byte[] {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}))
+                .build(),
+            /* muxerSupportedMimeTypes= */ ImmutableList.of(MimeTypes.VIDEO_H264),
+            emptyTransformationRequest,
+            fallbackListener,
+            /* matchInitializationData= */ false);
+
+    // Create the video encoder.
+    encoderWrapper.getSurfaceInfo(/* requestedWidth= */ 150, /* requestedHeight= */ 200);
+
+    assertThat(fakeEncoderFactory.format.initializationData).isEmpty();
+  }
+
+  @Test
+  public void matchInitializationData_setToTrue_initializationDataIsPassedToEncoderFactory()
+      throws Exception {
+    Format inputFormat =
+        new Format.Builder()
+            .setSampleMimeType(MimeTypes.VIDEO_H264)
+            .setColorInfo(ColorInfo.SDR_BT709_LIMITED)
+            .setInitializationData(ImmutableList.of(new byte[] {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}))
+            .build();
+    VideoSampleExporter.EncoderWrapper encoderWrapper =
+        new VideoSampleExporter.EncoderWrapper(
+            fakeEncoderFactory,
+            inputFormat,
+            /* muxerSupportedMimeTypes= */ ImmutableList.of(MimeTypes.VIDEO_H264),
+            emptyTransformationRequest,
+            fallbackListener,
+            /* matchInitializationData= */ true);
+
+    // Create the video encoder.
+    encoderWrapper.getSurfaceInfo(/* requestedWidth= */ 150, /* requestedHeight= */ 200);
+
+    assertThat(fakeEncoderFactory.format.initializationDataEquals(inputFormat)).isTrue();
+  }
+
   private static void createShadowH264Encoder() {
     MediaFormat avcFormat = new MediaFormat();
     avcFormat.setString(MediaFormat.KEY_MIME, MediaFormat.MIMETYPE_VIDEO_AVC);
@@ -169,6 +216,8 @@ public final class VideoEncoderWrapperTest {
 
   private static class FakeVideoEncoderFactory implements Codec.EncoderFactory {
 
+    public Format format;
+
     private int fallbackWidth;
     private int fallbackHeight;
 
@@ -189,6 +238,7 @@ public final class VideoEncoderWrapperTest {
 
     @Override
     public Codec createForVideoEncoding(Format format) {
+      this.format = format;
       Codec mockEncoder = mock(Codec.class);
       if (fallbackWidth != C.LENGTH_UNSET) {
         format = format.buildUpon().setWidth(fallbackWidth).build();