diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachine.java b/java/framework/src/android/system/virtualmachine/VirtualMachine.java
index 72a2429..9b1fcd1 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachine.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachine.java
@@ -310,7 +310,7 @@
     /** Running instance of virtmgr that hosts VirtualizationService for this VM. */
     @NonNull private final VirtualizationService mVirtualizationService;
 
-    @NonNull private final MemoryManagementCallbacks mMemoryManagementCallbacks;
+    private final MemoryManagementCallbacks mMemoryManagementCallbacks;
 
     @NonNull private final Context mContext;
 
@@ -441,7 +441,6 @@
         mInstanceFilePath = new File(thisVmDir, INSTANCE_IMAGE_FILE);
         mIdsigFilePath = new File(thisVmDir, IDSIG_FILE);
         mExtraApks = setupExtraApks(context, config, thisVmDir);
-        mMemoryManagementCallbacks = new MemoryManagementCallbacks();
         mContext = context;
         mEncryptedStoreFilePath =
                 (config.isEncryptedStorageEnabled())
@@ -451,6 +450,14 @@
         mVmOutputCaptured = config.isVmOutputCaptured();
         mVmConsoleInputSupported = config.isVmConsoleInputSupported();
         mConnectVmConsole = config.isConnectVmConsole();
+
+        VirtualMachineCustomImageConfig customImageConfig;
+        customImageConfig = config.getCustomImageConfig();
+        if (customImageConfig == null || customImageConfig.useAutoMemoryBalloon()) {
+            mMemoryManagementCallbacks = new MemoryManagementCallbacks();
+        } else {
+            mMemoryManagementCallbacks = null;
+        }
     }
 
     /**
@@ -820,7 +827,9 @@
      */
     @GuardedBy("mLock")
     private void dropVm() {
-        mContext.unregisterComponentCallbacks(mMemoryManagementCallbacks);
+        if (mMemoryManagementCallbacks != null) {
+            mContext.unregisterComponentCallbacks(mMemoryManagementCallbacks);
+        }
         mVirtualMachine = null;
     }
 
@@ -1166,6 +1175,46 @@
                         new InputEvent(EV_SYN, SYN_REPORT, 0)));
     }
 
+    /** @hide */
+    public long getMemoryBalloon() {
+        long bytes = 0;
+
+        if (mMemoryManagementCallbacks != null) {
+            Log.d(TAG, "Auto balloon enabled in getMemoryBalloon");
+            return bytes;
+        }
+
+        synchronized (mLock) {
+            try {
+                if (mVirtualMachine != null) {
+                    bytes = mVirtualMachine.getMemoryBalloon();
+                }
+            } catch (RemoteException e) {
+                Log.w(TAG, "Cannot getMemoryBalloon", e);
+            }
+        }
+
+        return bytes;
+    }
+
+    /** @hide */
+    public void setMemoryBalloon(long bytes) {
+        if (mMemoryManagementCallbacks != null) {
+            Log.d(TAG, "Auto balloon enabled in setMemoryBalloon");
+            return;
+        }
+
+        synchronized (mLock) {
+            try {
+                if (mVirtualMachine != null) {
+                    mVirtualMachine.setMemoryBalloon(bytes);
+                }
+            } catch (RemoteException e) {
+                Log.w(TAG, "Cannot setMemoryBalloon", e);
+            }
+        }
+    }
+
     private boolean writeEventsToSock(ParcelFileDescriptor sock, List<InputEvent> evtList) {
         ByteBuffer byteBuffer =
                 ByteBuffer.allocate(8 /* (type: u16 + code: u16 + value: i32) */ * evtList.size());
@@ -1317,7 +1366,9 @@
                 mVirtualMachine =
                         service.createVm(vmConfigParcel, consoleOutFd, consoleInFd, mLogWriter);
                 mVirtualMachine.registerCallback(new CallbackTranslator(service));
-                mContext.registerComponentCallbacks(mMemoryManagementCallbacks);
+                if (mMemoryManagementCallbacks != null) {
+                    mContext.registerComponentCallbacks(mMemoryManagementCallbacks);
+                }
                 if (mConnectVmConsole) {
                     mVirtualMachine.setHostConsoleName(getHostConsoleName());
                 }
diff --git a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
index 3a1c784..2da83a0 100644
--- a/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
+++ b/java/framework/src/android/system/virtualmachine/VirtualMachineCustomImageConfig.java
@@ -45,6 +45,7 @@
     private static final String KEY_GPU = "gpu";
     private static final String KEY_AUDIO_CONFIG = "audio_config";
     private static final String KEY_TRACKPAD = "trackpad";
+    private static final String KEY_AUTO_MEMORY_BALLOON = "auto_memory_balloon";
 
     @Nullable private final String name;
     @Nullable private final String kernelPath;
@@ -61,6 +62,7 @@
     private final boolean network;
     @Nullable private final GpuConfig gpuConfig;
     private final boolean trackpad;
+    private final boolean autoMemoryBalloon;
 
     @Nullable
     public Disk[] getDisks() {
@@ -112,6 +114,10 @@
         return mouse;
     }
 
+    public boolean useAutoMemoryBalloon() {
+        return autoMemoryBalloon;
+    }
+
     public boolean useNetwork() {
         return network;
     }
@@ -132,7 +138,8 @@
             boolean network,
             GpuConfig gpuConfig,
             AudioConfig audioConfig,
-            boolean trackpad) {
+            boolean trackpad,
+            boolean autoMemoryBalloon) {
         this.name = name;
         this.kernelPath = kernelPath;
         this.initrdPath = initrdPath;
@@ -148,6 +155,7 @@
         this.gpuConfig = gpuConfig;
         this.audioConfig = audioConfig;
         this.trackpad = trackpad;
+        this.autoMemoryBalloon = autoMemoryBalloon;
     }
 
     static VirtualMachineCustomImageConfig from(PersistableBundle customImageConfigBundle) {
@@ -199,6 +207,7 @@
                 customImageConfigBundle.getPersistableBundle(KEY_AUDIO_CONFIG);
         builder.setAudioConfig(AudioConfig.from(audioConfigPb));
         builder.useTrackpad(customImageConfigBundle.getBoolean(KEY_TRACKPAD));
+        builder.useAutoMemoryBalloon(customImageConfigBundle.getBoolean(KEY_AUTO_MEMORY_BALLOON));
         return builder.build();
     }
 
@@ -258,6 +267,7 @@
                 KEY_AUDIO_CONFIG,
                 Optional.ofNullable(audioConfig).map(ac -> ac.toPersistableBundle()).orElse(null));
         pb.putBoolean(KEY_TRACKPAD, trackpad);
+        pb.putBoolean(KEY_AUTO_MEMORY_BALLOON, autoMemoryBalloon);
         return pb;
     }
 
@@ -352,6 +362,7 @@
         private boolean network;
         private GpuConfig gpuConfig;
         private boolean trackpad;
+        private boolean autoMemoryBalloon = true;
 
         /** @hide */
         public Builder() {}
@@ -435,6 +446,12 @@
         }
 
         /** @hide */
+        public Builder useAutoMemoryBalloon(boolean autoMemoryBalloon) {
+            this.autoMemoryBalloon = autoMemoryBalloon;
+            return this;
+        }
+
+        /** @hide */
         public Builder useNetwork(boolean network) {
             this.network = network;
             return this;
@@ -463,7 +480,8 @@
                     network,
                     gpuConfig,
                     audioConfig,
-                    trackpad);
+                    trackpad,
+                    autoMemoryBalloon);
         }
     }