diff --git a/libraries/effect/src/main/java/androidx/media3/effect/OverlayShaderProgram.java b/libraries/effect/src/main/java/androidx/media3/effect/OverlayShaderProgram.java
index fef584b180..014fc72106 100644
--- a/libraries/effect/src/main/java/androidx/media3/effect/OverlayShaderProgram.java
+++ b/libraries/effect/src/main/java/androidx/media3/effect/OverlayShaderProgram.java
@@ -142,7 +142,7 @@ import java.io.IOException;
               glProgram.setSamplerTexIdUniform(
                   "uGainmapTexSampler" + texUnitIndex,
                   gainmapTexIds.get(texUnitIndex),
-                  texUnitIndex);
+                  /* texUnitIndex= */ overlays.size() + texUnitIndex);
               GainmapUtil.setGainmapUniforms(
                   glProgram, lastGainmaps.get(texUnitIndex), texUnitIndex);
             }
diff --git a/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_hlg.png b/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_hlg.png
index f136eda01c..83815c8fb4 100644
Binary files a/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_hlg.png and b/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_hlg.png differ
diff --git a/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_pq.png b/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_pq.png
index 4bb602c3c6..a3db1d15d0 100644
Binary files a/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_pq.png and b/libraries/test_data/src/test/assets/test-generated-goldens/hdr-goldens/ultrahdr_overlay_pq.png differ
diff --git a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/DefaultVideoFrameProcessorTextureOutputPixelTest.java b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/DefaultVideoFrameProcessorTextureOutputPixelTest.java
index 570a97633b..28bcfff350 100644
--- a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/DefaultVideoFrameProcessorTextureOutputPixelTest.java
+++ b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/mh/DefaultVideoFrameProcessorTextureOutputPixelTest.java
@@ -359,12 +359,7 @@ public final class DefaultVideoFrameProcessorTextureOutputPixelTest {
     videoFrameProcessorTestRunner.processFirstFrameAndEnd();
     Bitmap actualBitmap = videoFrameProcessorTestRunner.getOutputBitmap();
 
-    // TODO(b/207848601): Switch to using proper tooling for testing against golden data.
-    float averagePixelAbsoluteDifference =
-        BitmapPixelTestUtil.getBitmapAveragePixelAbsoluteDifferenceFp16(
-            expectedBitmap, actualBitmap);
-    assertThat(averagePixelAbsoluteDifference)
-        .isAtMost(MAXIMUM_AVERAGE_PIXEL_ABSOLUTE_DIFFERENCE_DIFFERENT_DEVICE_FP16);
+    assertBitmapsAreSimilar(expectedBitmap, actualBitmap, HDR_PSNR_THRESHOLD);
   }
 
   @Test
@@ -394,12 +389,9 @@ public final class DefaultVideoFrameProcessorTextureOutputPixelTest {
     videoFrameProcessorTestRunner.processFirstFrameAndEnd();
     Bitmap actualBitmap = videoFrameProcessorTestRunner.getOutputBitmap();
 
-    // TODO(b/207848601): Switch to using proper tooling for testing against golden data.
-    float averagePixelAbsoluteDifference =
-        BitmapPixelTestUtil.getBitmapAveragePixelAbsoluteDifferenceFp16(
-            expectedBitmap, actualBitmap);
-    assertThat(averagePixelAbsoluteDifference)
-        .isAtMost(MAXIMUM_AVERAGE_PIXEL_ABSOLUTE_DIFFERENCE_DIFFERENT_DEVICE_FP16);
+    // TODO: b/376016540 - use HDR_PSNR_THRESHOLD when HDR videos are processed with high floating
+    //  point precision.
+    assertBitmapsAreSimilar(expectedBitmap, actualBitmap, /* psnrThresholdDb= */ 42.5);
   }
 
   @Test