Request specific AudioFormat from AudioGraphInputs on creation.

Also adds an alternate way to configure the AudioGraph.

Apps should no longer need to ensure that inputs have the same sample
rate.

PiperOrigin-RevId: 588747431

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

@@ -25,6 +25,7 @@ import androidx.media3.common.Format;
 import androidx.media3.common.audio.AudioProcessor.AudioFormat;
 import androidx.media3.common.audio.AudioProcessor.UnhandledAudioFormatException;
 import java.nio.ByteBuffer;
+import java.util.Objects;
 
 /** Processes raw audio samples. */
 /* package */ final class AudioGraph {
@@ -57,17 +58,38 @@ import java.nio.ByteBuffer;
     return true;
   }
 
-  /** Returns a new {@link AudioGraphInput} instance. */
+  /**
+   * Configures the graph.
+   *
+   * <p>Must be called before {@linkplain #getOutput() accessing output}.
+   *
+   * <p>Should be called at most once, before {@link #registerInput registering input}.
+   *
+   * @param requestedAudioFormat The {@link AudioFormat} requested for output from the mixer.
+   * @throws UnhandledAudioFormatException If the audio format is not supported by the {@link
+   *     AudioMixer}.
+   */
+  public void configure(AudioFormat requestedAudioFormat) throws UnhandledAudioFormatException {
+    this.outputAudioFormat = requestedAudioFormat;
+    mixer.configure(requestedAudioFormat, /* bufferSizeMs= */ C.LENGTH_UNSET, /* startTimeUs= */ 0);
+  }
+
+  /**
+   * Returns a new {@link AudioGraphInput} instance.
+   *
+   * <p>Calls {@link #configure} if not already configured, using the {@linkplain
+   * AudioGraphInput#getOutputAudioFormat() outputAudioFormat} of the input.
+   */
   public AudioGraphInput registerInput(EditedMediaItem editedMediaItem, Format format)
       throws ExportException {
     checkArgument(format.pcmEncoding != Format.NO_VALUE);
     try {
-      AudioGraphInput audioGraphInput = new AudioGraphInput(editedMediaItem, format);
+      AudioGraphInput audioGraphInput =
+          new AudioGraphInput(outputAudioFormat, editedMediaItem, format);
 
-      if (inputs.size() == 0) {
-        outputAudioFormat = audioGraphInput.getOutputAudioFormat();
-        mixer.configure(
-            outputAudioFormat, /* bufferSizeMs= */ C.LENGTH_UNSET, /* startTimeUs= */ 0);
+      if (Objects.equals(outputAudioFormat, AudioFormat.NOT_SET)) {
+        // Graph not configured, configure before doing anything else.
+        configure(audioGraphInput.getOutputAudioFormat());
       }
 
       int sourceId = mixer.addSource(audioGraphInput.getOutputAudioFormat(), /* startTimeUs= */ 0);
@@ -79,10 +101,8 @@ import java.nio.ByteBuffer;
   }
 
   /**
-   * Returns the {@link AudioFormat} of the {@linkplain #getOutput() output}.
-   *
-   * <p>{@link AudioFormat#NOT_SET} is returned if no inputs have been {@linkplain #registerInput
-   * registered}.
+   * Returns the {@link AudioFormat} of the {@linkplain #getOutput() output}, or {@link
+   * AudioFormat#NOT_SET} if not {@linkplain #configure configured}.
    */
   public AudioFormat getOutputAudioFormat() {
     return outputAudioFormat;

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

@@ -67,7 +67,8 @@ import java.util.concurrent.atomic.AtomicReference;
   private boolean receivedEndOfStreamFromInput;
   private boolean queueEndOfStreamAfterSilence;
 
-  public AudioGraphInput(EditedMediaItem editedMediaItem, Format inputFormat)
+  public AudioGraphInput(
+      AudioFormat requestedOutputAudioFormat, EditedMediaItem editedMediaItem, Format inputFormat)
       throws UnhandledAudioFormatException {
     AudioFormat inputAudioFormat = new AudioFormat(inputFormat);
     checkArgument(isInputAudioFormatValid(inputAudioFormat), /* errorMessage= */ inputAudioFormat);
@@ -84,10 +85,7 @@ import java.util.concurrent.atomic.AtomicReference;
     silentAudioGenerator = new SilentAudioGenerator(inputAudioFormat);
     audioProcessingPipeline =
         configureProcessing(
-            editedMediaItem,
-            inputFormat,
-            inputAudioFormat,
-            /* requiredOutputAudioFormat= */ AudioFormat.NOT_SET);
+            editedMediaItem, inputFormat, inputAudioFormat, requestedOutputAudioFormat);
     // APP configuration not active until flush called. getOutputAudioFormat based on active config.
     audioProcessingPipeline.flush();
     outputAudioFormat = audioProcessingPipeline.getOutputAudioFormat();
@@ -336,24 +334,25 @@ import java.util.concurrent.atomic.AtomicReference;
           new SpeedChangingAudioProcessor(new SegmentSpeedProvider(inputFormat.metadata)));
     }
     audioProcessors.addAll(editedMediaItem.effects.audioProcessors);
-    // Ensure the output from APP matches what the encoder is configured to receive.
-    if (!requiredOutputAudioFormat.equals(AudioFormat.NOT_SET)) {
+
+    if (requiredOutputAudioFormat.sampleRate != Format.NO_VALUE) {
       SonicAudioProcessor sampleRateChanger = new SonicAudioProcessor();
       sampleRateChanger.setOutputSampleRateHz(requiredOutputAudioFormat.sampleRate);
       audioProcessors.add(sampleRateChanger);
+    }
 
-      // TODO(b/262706549): Handle channel mixing with AudioMixer.
-      if (requiredOutputAudioFormat.channelCount <= 2) {
-        // ChannelMixingMatrix.create only has defaults for mono/stereo input/output.
-        ChannelMixingAudioProcessor channelCountChanger = new ChannelMixingAudioProcessor();
-        channelCountChanger.putChannelMixingMatrix(
-            ChannelMixingMatrix.create(
-                /* inputChannelCount= */ 1, requiredOutputAudioFormat.channelCount));
-        channelCountChanger.putChannelMixingMatrix(
-            ChannelMixingMatrix.create(
-                /* inputChannelCount= */ 2, requiredOutputAudioFormat.channelCount));
-        audioProcessors.add(channelCountChanger);
-      }
+    // TODO(b/262706549): Handle channel mixing with AudioMixer.
+    // ChannelMixingMatrix.create only has defaults for mono/stereo input/output.
+    if (requiredOutputAudioFormat.channelCount == 1
+        || requiredOutputAudioFormat.channelCount == 2) {
+      ChannelMixingAudioProcessor channelCountChanger = new ChannelMixingAudioProcessor();
+      channelCountChanger.putChannelMixingMatrix(
+          ChannelMixingMatrix.create(
+              /* inputChannelCount= */ 1, requiredOutputAudioFormat.channelCount));
+      channelCountChanger.putChannelMixingMatrix(
+          ChannelMixingMatrix.create(
+              /* inputChannelCount= */ 2, requiredOutputAudioFormat.channelCount));
+      audioProcessors.add(channelCountChanger);
     }
 
     AudioProcessingPipeline audioProcessingPipeline =

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

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.media3.transformer;
+
+import static androidx.media3.common.util.Util.getPcmFormat;
+import static com.google.common.truth.Truth.assertThat;
+
+import androidx.media3.common.C;
+import androidx.media3.common.Format;
+import androidx.media3.common.MediaItem;
+import androidx.media3.common.MimeTypes;
+import androidx.media3.common.audio.AudioProcessor.AudioFormat;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+/** Unit tests for {@link AudioGraphInput}. */
+@RunWith(AndroidJUnit4.class)
+public class AudioGraphInputTest {
+  private static final EditedMediaItem FAKE_ITEM =
+      new EditedMediaItem.Builder(MediaItem.EMPTY).build();
+
+  @Test
+  public void getOutputAudioFormat_withUnsetRequestedFormat_matchesInputFormat() throws Exception {
+    AudioFormat requestedAudioFormat = AudioFormat.NOT_SET;
+    AudioFormat inputAudioFormat =
+        new AudioFormat(/* sampleRate= */ 48_000, /* channelCount= */ 1, C.ENCODING_PCM_16BIT);
+    Format inputFormat =
+        getPcmFormat(inputAudioFormat).buildUpon().setSampleMimeType(MimeTypes.AUDIO_RAW).build();
+
+    AudioGraphInput audioGraphInput =
+        new AudioGraphInput(requestedAudioFormat, FAKE_ITEM, inputFormat);
+
+    assertThat(audioGraphInput.getOutputAudioFormat()).isEqualTo(inputAudioFormat);
+  }
+
+  @Test
+  public void getOutputAudioFormat_withRequestedFormat_matchesRequestedFormat() throws Exception {
+    AudioFormat requestedAudioFormat =
+        new AudioFormat(/* sampleRate= */ 44_100, /* channelCount= */ 2, C.ENCODING_PCM_16BIT);
+    AudioFormat inputAudioFormat =
+        new AudioFormat(/* sampleRate= */ 48_000, /* channelCount= */ 1, C.ENCODING_PCM_16BIT);
+    Format inputFormat =
+        getPcmFormat(inputAudioFormat).buildUpon().setSampleMimeType(MimeTypes.AUDIO_RAW).build();
+
+    AudioGraphInput audioGraphInput =
+        new AudioGraphInput(requestedAudioFormat, FAKE_ITEM, inputFormat);
+
+    assertThat(audioGraphInput.getOutputAudioFormat()).isEqualTo(requestedAudioFormat);
+  }
+}

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

@@ -15,12 +15,14 @@
  */
 package androidx.media3.transformer;
 
+import static androidx.media3.common.MimeTypes.AUDIO_RAW;
+import static androidx.media3.common.util.Util.getPcmFormat;
 import static com.google.common.truth.Truth.assertThat;
 
 import androidx.media3.common.C;
 import androidx.media3.common.Format;
 import androidx.media3.common.MediaItem;
-import androidx.media3.common.MimeTypes;
+import androidx.media3.common.audio.AudioProcessor.AudioFormat;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import java.nio.ByteBuffer;
 import org.junit.Test;
@@ -29,22 +31,24 @@ import org.junit.runner.RunWith;
 /** Unit tests for {@link AudioGraph}. */
 @RunWith(AndroidJUnit4.class)
 public class AudioGraphTest {
+  private static final EditedMediaItem FAKE_ITEM =
+      new EditedMediaItem.Builder(MediaItem.EMPTY).build();
+
   @Test
   public void silentItem_outputsCorrectAmountOfBytes() throws Exception {
-    EditedMediaItem item = new EditedMediaItem.Builder(MediaItem.EMPTY).build();
     Format format =
         new Format.Builder()
             .setSampleRate(50_000)
             .setChannelCount(6)
             .setPcmEncoding(C.ENCODING_PCM_16BIT)
-            .setSampleMimeType(MimeTypes.AUDIO_RAW)
+            .setSampleMimeType(AUDIO_RAW)
             .build();
 
     AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
-    GraphInput input = audioGraph.registerInput(item, format);
+    GraphInput input = audioGraph.registerInput(FAKE_ITEM, format);
 
     input.onMediaItemChanged(
-        item, /* durationUs= */ 3_000_000, /* trackFormat= */ null, /* isLast= */ true);
+        FAKE_ITEM, /* durationUs= */ 3_000_000, /* trackFormat= */ null, /* isLast= */ true);
     int bytesOutput = drainAudioGraph(audioGraph);
 
     // 3 second stream with 50_000 frames per second.
@@ -52,6 +56,84 @@ public class AudioGraphTest {
     assertThat(bytesOutput).isEqualTo(3 * 50_000 * 2 * 6);
   }
 
+  @Test
+  public void getOutputAudioFormat_afterInitialization_isNotSet() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+
+    assertThat(audioGraph.getOutputAudioFormat()).isEqualTo(AudioFormat.NOT_SET);
+  }
+
+  @Test
+  public void getOutputAudioFormat_afterRegisterInput_matchesInputFormat() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    AudioFormat inputAudioFormat =
+        new AudioFormat(/* sampleRate= */ 48_000, /* channelCount= */ 1, C.ENCODING_PCM_16BIT);
+
+    audioGraph.registerInput(
+        FAKE_ITEM, getPcmFormat(inputAudioFormat).buildUpon().setSampleMimeType(AUDIO_RAW).build());
+
+    assertThat(audioGraph.getOutputAudioFormat()).isEqualTo(inputAudioFormat);
+  }
+
+  @Test
+  public void getOutputAudioFormat_afterConfigure_matchesConfiguredFormat() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    AudioFormat configuredAudioFormat =
+        new AudioFormat(/* sampleRate= */ 44_100, /* channelCount= */ 6, C.ENCODING_PCM_16BIT);
+
+    audioGraph.configure(configuredAudioFormat);
+
+    assertThat(audioGraph.getOutputAudioFormat()).isEqualTo(configuredAudioFormat);
+  }
+
+  @Test
+  public void registerInput_afterConfigure_doesNotChangeOutputFormat() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    AudioFormat configuredAudioFormat =
+        new AudioFormat(/* sampleRate= */ 44_100, /* channelCount= */ 2, C.ENCODING_PCM_16BIT);
+
+    audioGraph.configure(configuredAudioFormat);
+    audioGraph.registerInput(
+        FAKE_ITEM,
+        getPcmFormat(
+                new AudioFormat(
+                    /* sampleRate= */ 48_000, /* channelCount= */ 2, C.ENCODING_PCM_16BIT))
+            .buildUpon()
+            .setSampleMimeType(AUDIO_RAW)
+            .build());
+    audioGraph.registerInput(
+        FAKE_ITEM,
+        getPcmFormat(
+                new AudioFormat(
+                    /* sampleRate= */ 44_100, /* channelCount= */ 1, C.ENCODING_PCM_16BIT))
+            .buildUpon()
+            .setSampleMimeType(AUDIO_RAW)
+            .build());
+
+    assertThat(audioGraph.getOutputAudioFormat()).isEqualTo(configuredAudioFormat);
+  }
+
+  @Test
+  public void registerInput_afterRegisterInput_doesNotChangeOutputFormat() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    AudioFormat firstInputAudioFormat =
+        new AudioFormat(/* sampleRate= */ 48_000, /* channelCount= */ 2, C.ENCODING_PCM_16BIT);
+
+    audioGraph.registerInput(
+        FAKE_ITEM,
+        getPcmFormat(firstInputAudioFormat).buildUpon().setSampleMimeType(AUDIO_RAW).build());
+    audioGraph.registerInput(
+        FAKE_ITEM,
+        getPcmFormat(
+                new AudioFormat(
+                    /* sampleRate= */ 44_100, /* channelCount= */ 1, C.ENCODING_PCM_16BIT))
+            .buildUpon()
+            .setSampleMimeType(AUDIO_RAW)
+            .build());
+
+    assertThat(audioGraph.getOutputAudioFormat()).isEqualTo(firstInputAudioFormat);
+  }
+
   /** Drains the graph and returns the number of bytes output. */
   private static int drainAudioGraph(AudioGraph audioGraph) throws ExportException {
     int bytesOutput = 0;