diff --git a/justfile b/justfile
index 4fb3f7a..c0f3103 100644
--- a/justfile
+++ b/justfile
@@ -1,14 +1,12 @@
-alias nvidia := configure-nvidia
-
 # Configure the Nvidia driver
-configure-nvidia ACTION="prompt":
+_configure-nvidia ACTION="prompt":
     #!/usr/bin/bash
     source /usr/lib/ujust/ujust.sh
     OPTION={{ ACTION }}
     if [ "$OPTION" == "prompt" ]; then
       echo "${bold}Configuring Nvidia drivers${normal}"
       echo 'What would you like to do?'
-      OPTION=$(ugum choose "Set needed kernel arguments" "Remove kernel arguments" "Test CUDA support" "Enable Nvidia VAAPI in Firefox Flatpak")
+      OPTION=$(ugum choose "Set needed kernel arguments" "Remove kernel arguments")
     elif [ "$OPTION" == "help" ]; then
       echo "Usage: ujust configure-nvidia <option>"
       echo "  <option>: Specify the quick option - 'kargs', 'test-cuda' or 'firefox-vaapi'"
@@ -18,52 +16,17 @@ configure-nvidia ACTION="prompt":
       exit 0
     fi
     if [ "$OPTION" == "Set needed kernel arguments" ] || [ "${OPTION,,}" == "kargs" ]; then
-      if command -v nvidia-smi; then
         rpm-ostree kargs \
           --append-if-missing=rd.driver.blacklist=nouveau \
           --append-if-missing=modprobe.blacklist=nouveau \
           --append-if-missing=nvidia-drm.modeset=1 \
           --delete-if-present=nomodeset
-      else
-        echo 'You do not appear to be on a Nvidia image, please refer to the README for your uBlue-OS image.'
-      fi
     elif [ "$OPTION" == "Remove kernel arguments" ] || [ "${OPTION,,}" == "remove-kargs" ]; then
-      if command -v nvidia-smi; then
         rpm-ostree kargs \
           --delete-if-present=rd.driver.blacklist=nouveau \
           --delete-if-present=modprobe.blacklist=nouveau \
           --delete-if-present=nvidia-drm.modeset=1 \
           --append-if-missing=nomodeset
-      else
-        echo 'You do not appear to be on a Nvidia image, please refer to the README for your uBlue-OS image.'
-      fi
-    elif [ "$OPTION" == "Test CUDA support" ] || [ "${OPTION,,}" == "test-cuda" ]; then
-      if lsmod | grep -wq "nvidia"; then
-        podman run \
-          --user 1000:1000 \
-          --security-opt=no-new-privileges \
-          --cap-drop=ALL \
-          --security-opt label=type:nvidia_container_t  \
-          --device=nvidia.com/gpu=all \
-          docker.io/nvidia/samples:vectoradd-cuda11.2.1
-      else
-        echo 'The Nvidia kernel module is not loaded. You may be using secure boot without the needed signing key, lacking the needed kargs, or may not be on a Nvidia image. See "just enroll-secure-boot-key" and "just nvidia-set-kargs".'
-      fi
-    elif [ "$OPTION" == "Enable Nvidia VAAPI in Firefox Flatpak" ] || [ "${OPTION,,}" == "firefox-vaapi" ]; then
-      if lsmod | grep -wq "nvidia"; then
-        flatpak override \
-          --user \
-          --filesystem=host-os \
-          --env=LIBVA_DRIVER_NAME=nvidia \
-          --env=LIBVA_DRIVERS_PATH=/run/host/usr/lib64/dri \
-          --env=LIBVA_MESSAGING_LEVEL=1 \
-          --env=MOZ_DISABLE_RDD_SANDBOX=1 \
-          --env=NVD_BACKEND=direct \
-          --env=MOZ_ENABLE_WAYLAND=1 \
-          org.mozilla.firefox
-      else
-        echo 'The Nvidia kernel module is not loaded. You may be using secure boot without the needed signing key, lacking the needed kargs, or may not be on a Nvidia image. See "just enroll-secure-boot-key" and "just nvidia-set-kargs".'
-      fi
     fi
 
 # Switch between Nvidia image and NVK
@@ -72,12 +35,14 @@ toggle-nvk:
     CURRENT_IMAGE=$(rpm-ostree status -b --json | jq -r '.deployments[0]."container-image-reference"' | sed -E 's/^.+\/(.+:.+)$/\1/')
     CURRENT_URI=$(rpm-ostree status -b --json | jq -r '.deployments[0]."container-image-reference"' | sed -E 's/^(.+\/).+:.+$/\1/')
     NEW_IMAGE=$CURRENT_IMAGE
+    CONFIG_ACTION=""
     if [[ "$CURRENT_IMAGE" =~ "nvidia" ]]; then
       NEW_IMAGE=${CURRENT_IMAGE/-nvidia/}
-      just nvidia kargs
+      CONFIG_ACTION="kargs"
     else
       NEW_IMAGE=${CURRENT_IMAGE/:/-nvidia:}
-      just nvidia remove-kargs
+      CONFIG_ACTION="remove-kargs"
     fi
     echo "Rebasing to ${NEW_IMAGE}"
     rpm-ostree rebase ${CURRENT_URI}${NEW_IMAGE}
+    just _configure-nvidia ${CONFIG_ACTION}