I've cooked something up which does what I want:

let
 vmVariantModule =
    {
      config,
      pkgs,
      lib,
      ...
    }:
    let

      ovmf = pkgs.OVMF.override { secureBoot = false; };
    in
    {

      config.system.build.vmWithDiskoSecrets = (
        pkgs.writers.writeBashBin "disko-vm"
          {
            makeWrapperArgs = [
              "--prefix"
              "PATH"
              ":"
              "${lib.makeBinPath [
                pkgs.qemu_kvm
                pkgs.coreutils
              ]}"
            ];
          }
          ''
            set -efux

            # ---------- argument handling ----------
            usage() {
              echo "Usage: $0 --keyfile <path> [--boot-location <dest>] [--mem 4G] [--smp 2]"
              echo ""
              echo "  --keyfile        Path to the keyfile to inject (required)"
              echo "  --boot-location  Destination path inside the image (default: /run/secret)"
              echo "  --mem            QEMU RAM (default: 4G)"
              echo "  --smp            QEMU vCPUs (default: 2)"
              exit 1
            }

            KEYFILE=""
            BOOT_LOCATION="/run/secret"
            MEM="4G"
            SMP="2"

            while [[ $# -gt 0 ]]; do
              case "$1" in
                --keyfile)        KEYFILE="$2";        shift 2 ;;
                --boot-location)  BOOT_LOCATION="$2";  shift 2 ;;
                --mem)            MEM="$2";             shift 2 ;;
                --smp)            SMP="$2";             shift 2 ;;
                *)                usage ;;
              esac
            done

            if [[ -z "$KEYFILE" ]]; then
              echo "Error: --keyfile is required."
              usage
            fi

            # ---------- build image ----------
            WORKDIR=$(mktemp -d /tmp/vm-XXXXXX)
            trap 'rm -rf "$WORKDIR"' EXIT

            echo "[vm-disko] Building disko image in $WORKDIR …"
            # diskoImagesScript writes <hostName>.raw (or similar) into the cwd.
            pushd "$WORKDIR" > /dev/null
            ${config.system.build.diskoImagesScript} --post-format-files "$KEYFILE" "$BOOT_LOCATION"
            popd > /dev/null

            # Find the raw image produced (disko names it after the disk label,
            # so we just grab the first .raw file in the workdir).
            IMAGE=$(find "$WORKDIR" -maxdepth 1 -name '*.raw' | head -n1)
            if [[ -z "$IMAGE" ]]; then
              echo "Error: no .raw image found after diskoImagesScript."
              exit 1
            fi

            echo "[vm-disko] Booting $IMAGE …"

            # ---------- boot in QEMU ----------
            exec qemu-system-x86_64 \
              -enable-kvm \
              -m "$MEM" \
              -cpu max \
              -smp "$SMP" \
              -netdev user,id=net0,hostfwd=tcp::2222-:22 \
              -device virtio-net-pci,netdev=net0 \
              -drive if=pflash,format=raw,readonly=on,file=${ovmf.firmware} \
              -drive if=pflash,format=raw,readonly=on,file=${ovmf.variables} \
              -drive "if=virtio,format=raw,file=$IMAGE" \
              "$@"
          ''
      );
    };
    vmVariantWithDiskoSecrets = extendModules {
          modules = [
            vmVariantModule
          ];
        };
      in
      {
        options.virtualisation.vmVariantWithDiskoSecrets = lib.mkOption {
          description = ''
            Machine configuration to be added for the vm script available at `.system.build.vmWithDiskoSecrets`.
          '';
          inherit (vmVariantWithDiskoSecrets) type;
          default = { };
          visible = "shallow";
        };
        config.system.build.vmWithDiskoSecrets = lib.mkDefault config.virtualisation.vmVariantWithDiskoSecrets.system.build.vmWithDiskoSecrets;
      };