diff --git a/cmds/screencap/screencap.cpp b/cmds/screencap/screencap.cpp
index 12de82a..d563ad3 100644
--- a/cmds/screencap/screencap.cpp
+++ b/cmds/screencap/screencap.cpp
@@ -416,7 +416,6 @@
         format = ANDROID_BITMAP_COMPRESS_FORMAT_PNG;
     } else if (jpeg) {
         format = ANDROID_BITMAP_COMPRESS_FORMAT_JPEG;
-        captureArgs.attachGainmap = true;
     }
 
     // setThreadPoolMaxThreadCount(0) actually tells the kernel it's
diff --git a/core/java/android/window/ScreenCapture.java b/core/java/android/window/ScreenCapture.java
index 5446428..fc41307 100644
--- a/core/java/android/window/ScreenCapture.java
+++ b/core/java/android/window/ScreenCapture.java
@@ -20,6 +20,7 @@
 import android.annotation.Nullable;
 import android.graphics.Bitmap;
 import android.graphics.ColorSpace;
+import android.graphics.Gainmap;
 import android.graphics.PixelFormat;
 import android.graphics.Rect;
 import android.hardware.HardwareBuffer;
@@ -184,17 +185,29 @@
      * @hide
      */
     public static class ScreenshotHardwareBuffer {
+        private static final float EPSILON = 1.0f / 64.0f;
+
         private final HardwareBuffer mHardwareBuffer;
         private final ColorSpace mColorSpace;
         private final boolean mContainsSecureLayers;
         private final boolean mContainsHdrLayers;
+        private final HardwareBuffer mGainmap;
+        private final float mHdrSdrRatio;
 
         public ScreenshotHardwareBuffer(HardwareBuffer hardwareBuffer, ColorSpace colorSpace,
                 boolean containsSecureLayers, boolean containsHdrLayers) {
+            this(hardwareBuffer, colorSpace, containsSecureLayers, containsHdrLayers, null, 1.0f);
+        }
+
+        public ScreenshotHardwareBuffer(HardwareBuffer hardwareBuffer, ColorSpace colorSpace,
+                boolean containsSecureLayers, boolean containsHdrLayers, HardwareBuffer gainmap,
+                float hdrSdrRatio) {
             mHardwareBuffer = hardwareBuffer;
             mColorSpace = colorSpace;
             mContainsSecureLayers = containsSecureLayers;
             mContainsHdrLayers = containsHdrLayers;
+            mGainmap = gainmap;
+            mHdrSdrRatio = hdrSdrRatio;
         }
 
         /**
@@ -209,13 +222,12 @@
          * @param containsHdrLayers    Indicates whether this graphic buffer contains HDR content.
          */
         private static ScreenshotHardwareBuffer createFromNative(HardwareBuffer hardwareBuffer,
-                int dataspace, boolean containsSecureLayers, boolean containsHdrLayers) {
+                int dataspace, boolean containsSecureLayers, boolean containsHdrLayers,
+                HardwareBuffer gainmap, float hdrSdrRatio) {
             ColorSpace colorSpace = ColorSpace.getFromDataSpace(dataspace);
-            return new ScreenshotHardwareBuffer(
-                    hardwareBuffer,
+            return new ScreenshotHardwareBuffer(hardwareBuffer,
                     colorSpace != null ? colorSpace : ColorSpace.get(ColorSpace.Named.SRGB),
-                    containsSecureLayers,
-                    containsHdrLayers);
+                    containsSecureLayers, containsHdrLayers, gainmap, hdrSdrRatio);
         }
 
         public ColorSpace getColorSpace() {
@@ -259,7 +271,22 @@
                 Log.w(TAG, "Failed to take screenshot. Null screenshot object");
                 return null;
             }
-            return Bitmap.wrapHardwareBuffer(mHardwareBuffer, mColorSpace);
+
+            Bitmap bitmap = Bitmap.wrapHardwareBuffer(mHardwareBuffer, mColorSpace);
+            if (mGainmap != null) {
+                Bitmap gainmapBitmap = Bitmap.wrapHardwareBuffer(mGainmap, null);
+                Gainmap gainmap = new Gainmap(gainmapBitmap);
+                gainmap.setRatioMin(1.0f, 1.0f, 1.0f);
+                gainmap.setRatioMax(mHdrSdrRatio, mHdrSdrRatio, mHdrSdrRatio);
+                gainmap.setGamma(1.0f, 1.0f, 1.0f);
+                gainmap.setEpsilonSdr(EPSILON, EPSILON, EPSILON);
+                gainmap.setEpsilonHdr(EPSILON, EPSILON, EPSILON);
+                gainmap.setMinDisplayRatioForHdrTransition(1.0f);
+                gainmap.setDisplayRatioForFullHdr(mHdrSdrRatio);
+                bitmap.setGainmap(gainmap);
+            }
+
+            return bitmap;
         }
     }
 
diff --git a/core/jni/android_window_ScreenCapture.cpp b/core/jni/android_window_ScreenCapture.cpp
index 1a52fb7..5657fa1 100644
--- a/core/jni/android_window_ScreenCapture.cpp
+++ b/core/jni/android_window_ScreenCapture.cpp
@@ -110,13 +110,19 @@
         captureResults.fenceResult.value()->waitForever(LOG_TAG);
         jobject jhardwareBuffer = android_hardware_HardwareBuffer_createFromAHardwareBuffer(
                 env, captureResults.buffer->toAHardwareBuffer());
+        jobject jGainmap = nullptr;
+        if (captureResults.optionalGainMap) {
+            jGainmap = android_hardware_HardwareBuffer_createFromAHardwareBuffer(
+                    env, captureResults.optionalGainMap->toAHardwareBuffer());
+        }
         jobject screenshotHardwareBuffer =
                 env->CallStaticObjectMethod(gScreenshotHardwareBufferClassInfo.clazz,
                                             gScreenshotHardwareBufferClassInfo.builder,
                                             jhardwareBuffer,
                                             static_cast<jint>(captureResults.capturedDataspace),
                                             captureResults.capturedSecureLayers,
-                                            captureResults.capturedHdrLayers);
+                                            captureResults.capturedHdrLayers, jGainmap,
+                                            captureResults.hdrSdrRatio);
         checkAndClearException(env, "builder");
         env->CallVoidMethod(consumer.get(), gConsumerClassInfo.accept, screenshotHardwareBuffer,
                             fenceStatus(captureResults.fenceResult));
@@ -340,7 +346,8 @@
             MakeGlobalRefOrDie(env, screenshotGraphicsBufferClazz);
     gScreenshotHardwareBufferClassInfo.builder =
             GetStaticMethodIDOrDie(env, screenshotGraphicsBufferClazz, "createFromNative",
-                                   "(Landroid/hardware/HardwareBuffer;IZZ)Landroid/window/"
+                                   "(Landroid/hardware/HardwareBuffer;IZZLandroid/hardware/"
+                                   "HardwareBuffer;F)Landroid/window/"
                                    "ScreenCapture$ScreenshotHardwareBuffer;");
 
     return err;
diff --git a/libs/hwui/Android.bp b/libs/hwui/Android.bp
index e2db2c9..677fd86 100644
--- a/libs/hwui/Android.bp
+++ b/libs/hwui/Android.bp
@@ -154,7 +154,10 @@
                 "libstatssocket_lazy",
                 "libtonemap",
             ],
-            whole_static_libs: ["hwui_flags_cc_lib"],
+            whole_static_libs: [
+                "hwui_flags_cc_lib",
+                "libsurfaceflingerflags",
+            ],
         },
         host: {
             static_libs: [
diff --git a/libs/hwui/hwui/Bitmap.cpp b/libs/hwui/hwui/Bitmap.cpp
index cc292d9..b1550b0 100644
--- a/libs/hwui/hwui/Bitmap.cpp
+++ b/libs/hwui/hwui/Bitmap.cpp
@@ -19,6 +19,7 @@
 #include "HardwareBitmapUploader.h"
 #include "Properties.h"
 #ifdef __ANDROID__  // Layoutlib does not support render thread
+#include <com_android_graphics_surfaceflinger_flags.h>
 #include <private/android/AHardwareBufferHelpers.h>
 #include <ui/GraphicBuffer.h>
 #include <ui/GraphicBufferMapper.h>
@@ -562,7 +563,7 @@
 
 bool Bitmap::compress(JavaCompressFormat format, int32_t quality, SkWStream* stream) {
 #ifdef __ANDROID__  // TODO: This isn't built for host for some reason?
-    if (hasGainmap() && format == JavaCompressFormat::Jpeg) {
+    if (hasGainmap()) {
         SkBitmap baseBitmap = getSkBitmap();
         SkBitmap gainmapBitmap = gainmap()->bitmap->getSkBitmap();
         if (gainmapBitmap.colorType() == SkColorType::kAlpha_8_SkColorType) {
@@ -572,12 +573,27 @@
             greyGainmap.setPixelRef(sk_ref_sp(gainmapBitmap.pixelRef()), 0, 0);
             gainmapBitmap = std::move(greyGainmap);
         }
-        SkJpegEncoder::Options options{.fQuality = quality};
-        return SkJpegGainmapEncoder::EncodeHDRGM(stream, baseBitmap.pixmap(), options,
-                                                 gainmapBitmap.pixmap(), options, gainmap()->info);
+        switch (format) {
+            case JavaCompressFormat::Jpeg: {
+                SkJpegEncoder::Options options{.fQuality = quality};
+                return SkJpegGainmapEncoder::EncodeHDRGM(stream, baseBitmap.pixmap(), options,
+                                                         gainmapBitmap.pixmap(), options,
+                                                         gainmap()->info);
+            }
+            case JavaCompressFormat::Png: {
+                if (com::android::graphics::surfaceflinger::flags::true_hdr_screenshots()) {
+                    SkGainmapInfo info = gainmap()->info;
+                    SkPngEncoder::Options options{.fGainmap = &gainmapBitmap.pixmap(),
+                                                  .fGainmapInfo = &info};
+                    return SkPngEncoder::Encode(stream, baseBitmap.pixmap(), options);
+                }
+                // fallthrough if we're not supporting HDR screenshots
+            }
+            default:
+                ALOGI("Format: %d doesn't support gainmap compression!", format);
+        }
     }
 #endif
-
     SkBitmap skbitmap;
     getSkBitmap(&skbitmap);
     return compress(skbitmap, format, quality, stream);