diff --git a/docs/interactive-vm.md b/docs/interactive-vm.md
index 951ea3a..17df62e 100644
--- a/docs/interactive-vm.md
+++ b/docs/interactive-vm.md
@@ -12,6 +12,16 @@ afterwards you can run the interactive VM with:
 result/bin/disko-vm
 ```
 
+You can configure the VM using the `virtualisation.vmVariantWithDisko` NixOS option:
+
+```nix
+{
+  virtualisation.vmVariantWithDisko = {
+    virtualisation.fileSystems."/persist".neededForBoot = true;
+  };
+}
+```
+
 extraConfig that is set in disko.tests.extraConfig is also applied to the interactive VMs.
 imageSize of the VMs will be determined by the imageSize in the disk type in your disko config.
 memorySize is set by disko.memSize
diff --git a/lib/default.nix b/lib/default.nix
index 08a5e56..31acc9c 100644
--- a/lib/default.nix
+++ b/lib/default.nix
@@ -16,8 +16,6 @@ let
     # a version of makeDiskImage which runs outside of the store
     makeDiskImagesScript = args: (import ./make-disk-image.nix ({ inherit diskoLib; } // args)).impure;
 
-    makeVMRunner = args: (import ./interactive-vm.nix ({ inherit diskoLib; } // args)).pure;
-
     testLib = import ./tests.nix { inherit lib makeTest eval-config; };
     # like lib.types.oneOf but instead of a list takes an attrset
     # uses the field "type" to find the correct type in the attrset
diff --git a/lib/interactive-vm.nix b/lib/interactive-vm.nix
index 6fb0098..fd3ac05 100644
--- a/lib/interactive-vm.nix
+++ b/lib/interactive-vm.nix
@@ -1,12 +1,9 @@
-{ nixosConfig
-, diskoLib
-, pkgs ? nixosConfig.pkgs
-, name ? "${nixosConfig.config.networking.hostName}-disko-images"
-, extraConfig ? { }
-}:
+# We need to specify extendModules here to ensure that it is available
+# in args for makeDiskImages
+{ diskoLib, modulesPath, config, pkgs, lib, extendModules, ... }@args:
+
 let
-  lib = pkgs.lib;
-  vm_disko = (diskoLib.testLib.prepareDiskoConfig nixosConfig.config diskoLib.testLib.devices).disko;
+  vm_disko = (diskoLib.testLib.prepareDiskoConfig config diskoLib.testLib.devices).disko;
   cfg_ = (lib.evalModules {
     modules = lib.singleton {
       # _file = toString input;
@@ -24,7 +21,7 @@ let
   }).config;
   disks = lib.attrValues cfg_.disko.devices.disk;
   diskoImages = diskoLib.makeDiskImages {
-    nixosConfig = nixosConfig;
+    nixosConfig = args;
     copyNixStore = false;
     extraConfig = {
       disko.devices = cfg_.disko.devices;
@@ -47,34 +44,29 @@ let
       driveExtraOpts.werror = "report";
     })
     (builtins.tail disks);
-  vm = (nixosConfig.extendModules {
-    modules = [
-      ({ modulesPath, ... }: {
-        imports = [
-          (modulesPath + "/virtualisation/qemu-vm.nix")
-        ];
-      })
-      {
-        virtualisation.useEFIBoot = nixosConfig.config.disko.tests.efi;
-        virtualisation.memorySize = nixosConfig.config.disko.memSize;
-        virtualisation.useDefaultFilesystems = false;
-        virtualisation.diskImage = null;
-        virtualisation.qemu.drives = [ rootDisk ] ++ otherDisks;
-        boot.zfs.devNodes = "/dev/disk/by-uuid"; # needed because /dev/disk/by-id is empty in qemu-vms
-        boot.zfs.forceImportAll = true;
-      }
-      {
-        # generated from disko config
-        virtualisation.fileSystems = cfg_.disko.devices._config.fileSystems;
-        boot = cfg_.disko.devices._config.boot or { };
-        swapDevices = cfg_.disko.devices._config.swapDevices or [ ];
-      }
-      nixosConfig.config.disko.tests.extraConfig
-    ];
-  }).config.system.build.vm;
+
+  diskoBasedConfiguration = {
+    # generated from disko config
+    virtualisation.fileSystems = cfg_.disko.devices._config.fileSystems;
+    boot = cfg_.disko.devices._config.boot or { };
+    swapDevices = cfg_.disko.devices._config.swapDevices or [ ];
+  };
 in
 {
-  pure = pkgs.writers.writeDashBin "disko-vm" ''
+  imports = [
+    (modulesPath + "/virtualisation/qemu-vm.nix")
+    diskoBasedConfiguration
+  ];
+
+  virtualisation.useEFIBoot = config.disko.tests.efi;
+  virtualisation.memorySize = config.disko.memSize;
+  virtualisation.useDefaultFilesystems = false;
+  virtualisation.diskImage = null;
+  virtualisation.qemu.drives = [ rootDisk ] ++ otherDisks;
+  boot.zfs.devNodes = "/dev/disk/by-uuid"; # needed because /dev/disk/by-id is empty in qemu-vms
+  boot.zfs.forceImportAll = true;
+
+  system.build.vmWithDisko = pkgs.writers.writeDashBin "disko-vm" ''
     set -efux
     export tmp=$(mktemp -d)
     trap 'rm -rf "$tmp"' EXIT
@@ -84,6 +76,6 @@ in
       -F qcow2 "$tmp"/${disk.name}.qcow2
     '') disks}
     set +f
-    ${vm}/bin/run-*-vm
+    ${config.system.build.vm}/bin/run-*-vm
   '';
 }
diff --git a/module.nix b/module.nix
index c95261c..4583fbd 100644
--- a/module.nix
+++ b/module.nix
@@ -7,6 +7,14 @@ let
     eval-config = import (pkgs.path + "/nixos/lib/eval-config.nix");
   };
   cfg = config.disko;
+
+  vmVariantWithDisko = extendModules {
+    modules = [
+      ./lib/interactive-vm.nix
+      { _module.args = { inherit diskoLib; }; }
+      config.disko.tests.extraConfig
+    ];
+  };
 in
 {
   options.disko = {
@@ -137,6 +145,16 @@ in
       };
     };
   };
+
+  options.virtualisation.vmVariantWithDisko = lib.mkOption {
+    description = ''
+      Machine configuration to be added for the vm script available at `.system.build.vmWithDisko`.
+    '';
+    inherit (vmVariantWithDisko) type;
+    default = {};
+    visible = "shallow";
+  };
+
   config = lib.mkIf (cfg.devices.disk != { }) {
     system.build = (cfg.devices._scripts { inherit pkgs; checked = cfg.checkScripts; }) // {
 
@@ -161,10 +179,7 @@ in
         extraTestScript = cfg.tests.extraChecks;
       };
 
-      vmWithDisko = diskoLib.makeVMRunner {
-        inherit pkgs;
-        nixosConfig = args;
-      };
+      vmWithDisko = lib.mkDefault config.virtualisation.vmVariantWithDisko.system.build.vmWithDisko;
     };