From 4270885b11a8d4882e14600a1b2f1cc435199ea9 Mon Sep 17 00:00:00 2001
From: Jared Baur <jaredbaur@fastmail.com>
Date: Wed, 17 Apr 2024 14:24:11 -0700
Subject: [PATCH] Move `devicePkgs` to the `pkgs.nvidia-jetpack` package-set
 and clear up ambiguities with systems compatible for flash/fuse scripts

Instead of maintaining a separate package-set location (outside of
`pkgs`), we can just use regular overlays to apply our changes that are
device-specific (stored under the `pkgs.nvidia-jetpack` scope). An alias
is added to the old locations (`config.system.build.jetsonDevicePkgs`,
`config.hardware.nvidia-jetpack.devicePkgs`, etc.) with a warning to
indicate that users should just use `pkgs.nvidia-jetpack`.

Also included is clearing up the ambiguities of what systems are
compatible with flash/fuse scripts. NVIDIA makes the decision for us as
to what platforms we can run these tools on (x86_64-linux only), so we
shouldn't allow for any flash/fuse derivations to be built for
aarch64-linux.

The rundown:
- a jetson device nixos config's `hostPlatform` _must_ be `aarch64-linux`
- a flash/fuse script's `hostPlatform` _must_ be `x86_64-linux`

What we were doing before was mixing package-sets willy-nilly without
much control over which hostPlatform we were dealing with (leading to
lots of usage of hardcoded `pkgsAarch64` to force a package-set to
aarch64). This change makes a logical separation of what needs to be
built with an aarch64 hostPlatform package-set vs an x86_64 hostPlatform
package-set.
---
 UPGRADE_CHECKLIST.md               |   4 +-
 default.nix                        | 214 ---------------------
 device-pkgs/default.nix            | 223 ++++++----------------
 device-pkgs/flash-script.nix       |   7 +-
 flake.nix                          |  64 +++++--
 kernel/default.nix                 |  13 +-
 modules/default.nix                |  51 ++++-
 modules/flash-script.nix           | 292 +++++++++++++++--------------
 overlay-with-config.nix            | 149 +++++++++++++++
 overlay.nix                        | 136 +++++++++++++-
 pkgs/cuda-packages/default.nix     |  12 +-
 pkgs/flash-from-device/default.nix |  13 +-
 pkgs/samples/default.nix           |   6 +-
 13 files changed, 609 insertions(+), 575 deletions(-)
 delete mode 100644 default.nix
 create mode 100644 overlay-with-config.nix

diff --git a/UPGRADE_CHECKLIST.md b/UPGRADE_CHECKLIST.md
index 4c2c79b..83261e0 100644
--- a/UPGRADE_CHECKLIST.md
+++ b/UPGRADE_CHECKLIST.md
@@ -1,7 +1,7 @@
 ### Updating
-- [ ] Update `l4tVersion`, `jetpackVersion`, and `cudaVersion` in default.nix
+- [ ] Update `l4tVersion`, `jetpackVersion`, and `cudaVersion` in overlay.nix
 - [ ] Update branch/revision/sha256s in:
-    - [ ] `default.nix`
+    - [ ] `overlay.nix`
     - [ ] `kernel/default.nix`
     - [ ] `uefi-firmware.nix`
     - [ ] Grep for "sha256 = ", see if there is anything else not covered
diff --git a/default.nix b/default.nix
deleted file mode 100644
index bd88283..0000000
--- a/default.nix
+++ /dev/null
@@ -1,214 +0,0 @@
-{ callPackage
-, callPackages
-, stdenv
-, stdenvNoCC
-, lib
-, runCommand
-, fetchurl
-, fetchgit
-, bzip2_1_1
-, dpkg
-, pkgs
-, dtc
-, python3
-, runtimeShell
-, writeShellApplication
-}:
-
-let
-  # Grab this from nixpkgs cudaPackages
-  inherit (pkgs.cudaPackages) autoAddOpenGLRunpathHook;
-
-  pkgsAarch64 = if pkgs.stdenv.buildPlatform.isAarch64 then pkgs else pkgs.pkgsCross.aarch64-multiplatform;
-
-  jetpackVersion = "5.1.2";
-  l4tVersion = "35.4.1";
-  cudaVersion = "11.4";
-
-  # https://developer.nvidia.com/embedded/jetson-linux-archive
-  # https://repo.download.nvidia.com/jetson/
-
-  src = fetchurl {
-    url = with lib.versions; "https://developer.download.nvidia.com/embedded/L4T/r${major l4tVersion}_Release_v${minor l4tVersion}.${patch l4tVersion}/release/Jetson_Linux_R${l4tVersion}_aarch64.tbz2";
-    sha256 = "sha256-crdaDH+jv270GuBmNLtnw4qSaCFV0SBgJtvuSmuaAW8=";
-  };
-
-  sourceInfo = import ./sourceinfo { inherit lib fetchurl fetchgit l4tVersion; };
-  inherit (sourceInfo) debs gitRepos;
-
-  # we use a more recent version of bzip2 here because we hit this bug extracting nvidia's archives:
-  # https://bugs.launchpad.net/ubuntu/+source/bzip2/+bug/1834494
-  bspSrc = runCommand "l4t-unpacked" { nativeBuildInputs = [ bzip2_1_1 ]; } ''
-    bzip2 -d -c ${src} | tar xf -
-    mv Linux_for_Tegra $out
-  '';
-
-  # Here for convenience, to see what is in upstream Jetpack
-  unpackedDebs = pkgs.runCommand "unpackedDebs-${l4tVersion}" { nativeBuildInputs = [ dpkg ]; } ''
-    mkdir -p $out
-    ${lib.concatStringsSep "\n" (lib.mapAttrsToList (n: p: "echo Unpacking ${n}; dpkg -x ${p.src} $out/${n}") debs.common)}
-    ${lib.concatStringsSep "\n" (lib.mapAttrsToList (n: p: "echo Unpacking ${n}; dpkg -x ${p.src} $out/${n}") debs.t234)}
-  '';
-
-  # Also just for convenience,
-  unpackedDebsFilenames = pkgs.runCommand "unpackedDebsFilenames-${l4tVersion}" { nativeBuildInputs = [ dpkg ]; } ''
-    mkdir -p $out
-    ${lib.concatStringsSep "\n" (lib.mapAttrsToList (n: p: "echo Extracting file list from ${n}; dpkg --fsys-tarfile ${p.src} | tar --list > $out/${n}") debs.common)}
-    ${lib.concatStringsSep "\n" (lib.mapAttrsToList (n: p: "echo Extracting file list from ${n}; dpkg --fsys-tarfile ${p.src} | tar --list > $out/${n}") debs.t234)}
-  '';
-
-  unpackedGitRepos = pkgs.runCommand "unpackedGitRepos-${l4tVersion}" { } (
-    lib.mapAttrsToList
-      (relpath: repo: ''
-        mkdir -p $out/${relpath}
-        cp --no-preserve=all -r ${repo}/. $out/${relpath}
-      '')
-      gitRepos
-  );
-
-  inherit (pkgsAarch64.callPackages ./pkgs/uefi-firmware { inherit l4tVersion; })
-    edk2-jetson uefi-firmware;
-
-  inherit (pkgsAarch64.callPackages ./pkgs/optee {
-    # Nvidia's recommended toolchain is gcc9:
-    # https://nv-tegra.nvidia.com/r/gitweb?p=tegra/optee-src/nv-optee.git;a=blob;f=optee/atf_and_optee_README.txt;h=591edda3d4ec96997e054ebd21fc8326983d3464;hb=5ac2ab218ba9116f1df4a0bb5092b1f6d810e8f7#l33
-    stdenv = pkgsAarch64.gcc9Stdenv;
-    inherit bspSrc gitRepos l4tVersion;
-  }) buildTOS buildOpteeTaDevKit opteeClient;
-
-  flash-tools = callPackage ./pkgs/flash-tools {
-    inherit bspSrc l4tVersion;
-  };
-
-  board-automation = callPackage ./pkgs/board-automation {
-    inherit bspSrc l4tVersion;
-  };
-
-  python-jetson = python3.pkgs.callPackage ./pkgs/python-jetson { };
-
-  tegra-eeprom-tool = pkgsAarch64.callPackage ./pkgs/tegra-eeprom-tool { };
-  tegra-eeprom-tool-static = pkgsAarch64.pkgsStatic.callPackage ./pkgs/tegra-eeprom-tool { };
-
-  l4t = callPackages ./pkgs/l4t { inherit debs l4tVersion; };
-
-  cudaPackages = callPackages ./pkgs/cuda-packages { inherit debs cudaVersion autoAddOpenGLRunpathHook l4t; };
-
-  samples = callPackages ./pkgs/samples { inherit debs cudaVersion autoAddOpenGLRunpathHook l4t cudaPackages; };
-
-  kernel = callPackage ./kernel { inherit (l4t) l4t-xusb-firmware; kernelPatches = [ ]; };
-  kernelPackagesOverlay = self: super: {
-    nvidia-display-driver = self.callPackage ./kernel/display-driver.nix { inherit gitRepos l4tVersion; };
-  };
-  kernelPackages = (pkgs.linuxPackagesFor kernel).extend kernelPackagesOverlay;
-
-  rtkernel = callPackage ./kernel { inherit (l4t) l4t-xusb-firmware; kernelPatches = [ ]; realtime = true; };
-  rtkernelPackages = (pkgs.linuxPackagesFor rtkernel).extend kernelPackagesOverlay;
-
-  nxJetsonBenchmarks = pkgs.callPackage ./pkgs/jetson-benchmarks {
-    targetSom = "nx";
-    inherit cudaPackages;
-  };
-  xavierAgxJetsonBenchmarks = pkgs.callPackage ./pkgs/jetson-benchmarks {
-    targetSom = "xavier-agx";
-    inherit cudaPackages;
-  };
-  orinAgxJetsonBenchmarks = pkgs.callPackage ./pkgs/jetson-benchmarks {
-    targetSom = "orin-agx";
-    inherit cudaPackages;
-  };
-
-  supportedConfigurations = lib.listToAttrs (map
-    (c: {
-      name = "${c.som}-${c.carrierBoard}";
-      value = c;
-    }) [
-    { som = "orin-agx"; carrierBoard = "devkit"; }
-    { som = "orin-agx-industrial"; carrierBoard = "devkit"; }
-    { som = "orin-nx"; carrierBoard = "devkit"; }
-    { som = "orin-nano"; carrierBoard = "devkit"; }
-    { som = "xavier-agx"; carrierBoard = "devkit"; }
-    { som = "xavier-nx"; carrierBoard = "devkit"; }
-    { som = "xavier-nx-emmc"; carrierBoard = "devkit"; }
-  ]);
-
-  supportedNixOSConfigurations = lib.mapAttrs
-    (n: c: {
-      imports = [ ./modules/default.nix ];
-      hardware.nvidia-jetpack = { enable = true; } // c;
-      networking.hostName = "${c.som}-${c.carrierBoard}"; # Just so it sets the flash binary name.
-    })
-    supportedConfigurations;
-
-  flashFromDevice = callPackage ./pkgs/flash-from-device {
-    inherit pkgsAarch64 tegra-eeprom-tool-static;
-  };
-
-  # Packages whose contents are parameterized by NixOS configuration
-  devicePkgsFromNixosConfig = callPackage ./device-pkgs {
-    inherit l4tVersion pkgsAarch64 flash-tools flashFromDevice edk2-jetson uefi-firmware buildTOS buildOpteeTaDevKit;
-  };
-
-  otaUtils = callPackage ./pkgs/ota-utils {
-    inherit tegra-eeprom-tool l4tVersion;
-  };
-in
-rec {
-  inherit jetpackVersion l4tVersion cudaVersion;
-
-  # Just for convenience
-  inherit bspSrc debs gitRepos;
-  inherit unpackedDebs unpackedDebsFilenames unpackedGitRepos;
-
-  inherit cudaPackages samples;
-  inherit flash-tools;
-  inherit board-automation; # Allows automation of Orin AGX devkit
-  inherit python-jetson; # Allows automation of Xavier AGX devkit
-  inherit tegra-eeprom-tool;
-
-  inherit kernel kernelPackages;
-  inherit rtkernel rtkernelPackages;
-
-  inherit nxJetsonBenchmarks xavierAgxJetsonBenchmarks orinAgxJetsonBenchmarks;
-
-  inherit edk2-jetson uefi-firmware;
-  inherit otaUtils;
-
-  inherit opteeClient;
-
-  # TODO: Source packages. source_sync.sh from bspSrc
-  # GST plugins
-
-  inherit flashFromDevice;
-
-  inherit devicePkgsFromNixosConfig;
-
-  devicePkgs = lib.mapAttrs (n: c: devicePkgsFromNixosConfig (pkgs.nixos c).config) supportedNixOSConfigurations;
-
-  flash-generic = writeShellApplication {
-    name = "flash-generic";
-    text = callPackage ./device-pkgs/flash-script.nix {
-      inherit flash-tools uefi-firmware;
-      flashCommands = ''
-        ${runtimeShell}
-      '';
-      # Use cross-compiled machine here so we don't have to depend on aarch64 builders
-      # TODO: Do a smaller cross-compiled version from old jetpack dir
-      dtbsDir = (pkgsAarch64.nixos {
-        imports = [ ./modules/default.nix ];
-        hardware.nvidia-jetpack.enable = true;
-      }).config.hardware.deviceTree.package;
-    };
-  };
-
-  l4tCsv = callPackage ./pkgs/containers/l4t-csv.nix { inherit bspSrc; };
-  genL4tJson = runCommand "l4t.json" { nativeBuildInputs = [ python3 ]; } ''
-    python3 ${./pkgs/containers/gen_l4t_json.py} ${l4tCsv} ${unpackedDebsFilenames} > $out
-  '';
-  containerDeps = callPackage ./pkgs/containers/deps.nix { inherit debs; };
-  nvidia-ctk = callPackage ./pkgs/containers/nvidia-ctk.nix { };
-
-  flashScripts = lib.mapAttrs' (n: c: lib.nameValuePair "flash-${n}" c.flashScript) devicePkgs;
-  initrdFlashScripts = lib.mapAttrs' (n: c: lib.nameValuePair "initrd-flash-${n}" c.initrdFlashScript) devicePkgs;
-  uefiCapsuleUpdates = lib.mapAttrs' (n: c: lib.nameValuePair "uefi-capsule-update-${n}" c.uefiCapsuleUpdate) devicePkgs;
-}
-  // l4t
diff --git a/device-pkgs/default.nix b/device-pkgs/default.nix
index 5f60785..1bc2468 100644
--- a/device-pkgs/default.nix
+++ b/device-pkgs/default.nix
@@ -1,78 +1,35 @@
+# These come from the device's nixos module arguments, so `pkgs` is actually an
+# aarch64 hostPlatform packaget-set.
+{ config, pkgs, ... }:
+
+# These must be filled in by a `callPackage` from an x86_64 hostPlatform
+# package-set to satisfy being able to run nvidia's prebuilt binaries on an
+# x86-compatible platform.
 { lib
-, callPackage
+, dtc
+, gcc
+, makeInitrd
+, nvidia-jetpack
+, openssl
+, python3
 , runCommand
 , writeScript
 , writeShellApplication
-, makeInitrd
-, makeModulesClosure
-, flashFromDevice
-, edk2-jetson
-, uefi-firmware
-, flash-tools
-, buildTOS
-, buildOpteeTaDevKit
-, python3
-, openssl
-, gcc
-, dtc
-, l4tVersion
-, pkgsAarch64
-,
 }:
 
-config:
-
 let
   cfg = config.hardware.nvidia-jetpack;
-  hostName = config.networking.hostName;
-
-  socType =
-    if cfg.som == null then null
-    else if lib.hasPrefix "orin-" cfg.som then "t234"
-    else if lib.hasPrefix "xavier-" cfg.som then "t194"
-    else throw "Unknown SoC type";
-
-  inherit (cfg.flashScriptOverrides)
-    flashArgs fuseArgs partitionTemplate;
-
-  flash-tools-patched = flash-tools.overrideAttrs ({ patches ? [ ], postPatch ? "", ... }: {
-    patches = patches ++ cfg.flashScriptOverrides.patches;
-    postPatch = postPatch + cfg.flashScriptOverrides.postPatch;
-  });
-
-  tosArgs = {
-    inherit socType;
-    inherit (cfg.firmware.optee) taPublicKeyFile;
-    opteePatches = cfg.firmware.optee.patches;
-    extraMakeFlags = cfg.firmware.optee.extraMakeFlags;
-  };
-  tosImage = buildTOS tosArgs;
-  taDevKit = buildOpteeTaDevKit tosArgs;
-
-  mkFlashScript = args: import ./flash-script.nix ({
-    inherit lib flashArgs partitionTemplate;
-
-    inherit (cfg.flashScriptOverrides) additionalDtbOverlays;
+  inherit (config.networking) hostName;
 
-    flash-tools = flash-tools-patched;
+  # We need to grab some packages from the device's aarch64 package set.
+  inherit (pkgs.nvidia-jetpack)
+    chipId
+    flashInitrd
+    l4tVersion
+    mkFlashScript
+    ;
 
-    uefi-firmware = uefi-firmware.override ({
-      bootLogo = cfg.firmware.uefi.logo;
-      debugMode = cfg.firmware.uefi.debugMode;
-      errorLevelInfo = cfg.firmware.uefi.errorLevelInfo;
-      edk2NvidiaPatches = cfg.firmware.uefi.edk2NvidiaPatches;
-      edk2UefiPatches = cfg.firmware.uefi.edk2UefiPatches;
-    } // lib.optionalAttrs cfg.firmware.uefi.capsuleAuthentication.enable {
-      inherit (cfg.firmware.uefi.capsuleAuthentication) trustedPublicCertPemFile;
-    });
-
-    inherit socType;
-
-    inherit tosImage;
-    eksFile = cfg.firmware.eksFile;
-
-    dtbsDir = config.hardware.deviceTree.package;
-  } // args);
+  inherit (cfg.flashScriptOverrides) flashArgs fuseArgs partitionTemplate;
 
   # This produces a script where we have already called the ./flash.sh script
   # with `--no-flash` and produced a file under bootloader/flashcmd.txt.
@@ -103,7 +60,7 @@ let
 
         ${cfg.firmware.secureBoot.preSignCommands}
 
-        ${mkFlashScript (args // { flashArgs = [ "--no-root-check" "--no-flash" ] ++ (args.flashArgs or flashArgs); }) }
+        ${mkFlashScript nvidia-jetpack.flash-tools (args // { flashArgs = [ "--no-root-check" "--no-flash" ] ++ (args.flashArgs or flashArgs); }) }
 
         cp -r ./ $out
       '';
@@ -120,19 +77,20 @@ let
   # With either produce a standard flash script, which does variant detection,
   # or if there is only a single variant, will produce a script specialized to
   # that particular variant.
-  mkFlashScriptAuto = if builtins.length cfg.firmware.variants == 1 then mkFlashCmdScript else mkFlashScript;
+  mkFlashScriptAuto = if builtins.length cfg.firmware.variants == 1 then mkFlashCmdScript else (mkFlashScript nvidia-jetpack.flash-tools);
 
   # Generate a flash script using the built configuration options set in a NixOS configuration
   flashScript = writeShellApplication {
     name = "flash-${hostName}";
     text = (mkFlashScriptAuto { });
+    meta.platforms = [ "x86_64-linux" ];
   };
 
   # Produces a script that boots a given kernel, initrd, and cmdline using the RCM boot method
   mkRcmBootScript = { kernelPath, initrdPath, kernelCmdline }: mkFlashScriptAuto {
     preFlashCommands = ''
       cp ${kernelPath} kernel/Image
-      cp ${initrdPath}/initrd bootloader/l4t_initrd.img
+      cp ${initrdPath} bootloader/l4t_initrd.img
 
       export CMDLINE="${builtins.toString kernelCmdline}"
       export INITRD_IN_BOOTIMG="yes"
@@ -147,70 +105,47 @@ let
     name = "rcmboot-nixos";
     text = mkRcmBootScript {
       # See nixpkgs nixos/modules/system/activatation/top-level.nix for standard usage of these paths
-      kernelPath = "${config.boot.kernelPackages.kernel}/${config.system.boot.loader.kernelFile}";
+      kernelPath = "${config.system.build.kernel}/${config.system.boot.loader.kernelFile}";
       initrdPath = "${config.system.build.initialRamdisk}/${config.system.boot.loader.initrdFile}";
       kernelCmdline = "init=${config.system.build.toplevel}/init initrd=initrd ${toString config.boot.kernelParams}";
     };
+    meta.platforms = [ "x86_64-linux" ];
   };
 
   # TODO: The flash script should not have the kernel output in its runtime closure
   initrdFlashScript =
-    let
-      modules = [ "qspi_mtd" "spi_tegra210_qspi" "at24" "spi_nor" ];
-      modulesClosure = makeModulesClosure {
-        rootModules = modules;
-        kernel = config.system.modulesTree;
-        firmware = config.hardware.firmware;
-        allowMissing = false;
-      };
-      jetpack-init = writeScript "init" ''
-        #!${pkgsAarch64.pkgsStatic.busybox}/bin/sh
-        export PATH=${pkgsAarch64.pkgsStatic.busybox}/bin
-        mkdir -p /proc /dev /sys
-        mount -t proc proc -o nosuid,nodev,noexec /proc
-        mount -t devtmpfs none -o nosuid /dev
-        mount -t sysfs sysfs -o nosuid,nodev,noexec /sys
-
-        for mod in ${builtins.toString modules}; do
-          modprobe -v $mod
-        done
-
-        if ${flashFromDevice}/bin/${flashFromDevice.name} ${signedFirmware}; then
-          echo "Flashing platform firmware successful. Rebooting now."
-          sync
-          reboot -f
-        else
-          echo "Flashing platform firmware unsuccessful. Entering console"
-          exec ${pkgsAarch64.pkgsStatic.busybox}/bin/sh
-        fi
-      '';
-      initrd = makeInitrd {
-        contents = [
-          { object = jetpack-init; symlink = "/init"; }
-          { object = "${modulesClosure}/lib/modules"; symlink = "/lib/modules"; }
-          { object = "${modulesClosure}/lib/firmware"; symlink = "/lib/firmware"; }
-        ];
-      };
-    in
     writeShellApplication {
       name = "initrd-flash-${hostName}";
       text = ''
         ${mkRcmBootScript {
-          kernelPath = "${config.boot.kernelPackages.kernel}/Image";
-          initrdPath = initrd;
+          kernelPath = "${config.system.build.kernel}/${config.system.boot.loader.kernelFile}";
+          initrdPath =
+            let
+              signedFirmwareInitrd = makeInitrd {
+                contents = [{ object = signedFirmware; symlink = "/signed-firmware"; }];
+              };
+            in
+            # The linux kernel supports concatenated initrds where each initrd
+            # can be optionally compressed with any compression algorithm
+            # supported by the kernel (initrds don't need to match in
+            # compression algorithm).
+            runCommand "combined-initrd" { } ''
+              cat ${flashInitrd}/initrd ${signedFirmwareInitrd}/initrd > $out
+            '';
           kernelCmdline = "initrd=initrd console=ttyTCU0,115200";
         }}
         echo
         echo "Jetson device should now be flashing and will reboot when complete."
         echo "You may watch the progress of this on the device's serial port"
       '';
+      meta.platforms = [ "x86_64-linux" ];
     };
 
   signedFirmware = runCommand "signed-${hostName}-${l4tVersion}"
     {
       inherit (cfg.firmware.secureBoot) requiredSystemFeatures;
     }
-    (mkFlashScript {
+    (mkFlashScript nvidia-jetpack.flash-tools {
       flashCommands = cfg.firmware.secureBoot.preSignCommands + lib.concatMapStringsSep "\n"
         (v: with v; ''
           BOARDID=${boardid} BOARDSKU=${boardsku} FAB=${fab} BOARDREV=${boardrev} FUSELEVEL=${fuselevel} CHIPREV=${chiprev} ./flash.sh ${lib.optionalString (partitionTemplate != null) "-c flash.xml"} --no-root-check --no-flash --sign ${builtins.toString flashArgs}
@@ -239,71 +174,31 @@ let
         cfg.firmware.variants;
     });
 
-  # Bootloader Update Package (BUP)
-  # TODO: Maybe generate this ourselves from signedFirmware so we dont have multiple scripts using the same keys to sign the same artifacts
-  bup = runCommand "bup-${hostName}-${l4tVersion}"
-    {
-      inherit (cfg.firmware.secureBoot) requiredSystemFeatures;
-    }
-    ((mkFlashScript {
-      # TODO: Remove preSignCommands when we switch to using signedFirmware directly
-      flashCommands = cfg.firmware.secureBoot.preSignCommands + lib.concatMapStringsSep "\n"
-        (v: with v;
-        "BOARDID=${boardid} BOARDSKU=${boardsku} FAB=${fab} BOARDREV=${boardrev} FUSELEVEL=${fuselevel} CHIPREV=${chiprev} ./flash.sh ${lib.optionalString (partitionTemplate != null) "-c flash.xml"} --no-flash --bup --multi-spec ${builtins.toString flashArgs}"
-        )
-        cfg.firmware.variants;
-    }) + ''
-      mkdir -p $out
-      cp -r bootloader/payloads_*/* $out/
-    '');
-
-  # See l4t_generate_soc_bup.sh
-  # python ${edk2-jetson}/BaseTools/BinWrappers/PosixLike/GenerateCapsule -v --encode --monotonic-count 1
-  # NOTE: providing null public certs here will use the test certs in the EDK2 repo
-  uefiCapsuleUpdate = runCommand "uefi-${hostName}-${l4tVersion}.Cap"
-    {
-      nativeBuildInputs = [ python3 openssl ];
-      inherit (cfg.firmware.uefi.capsuleAuthentication) requiredSystemFeatures;
-    }
-    (''
-      ${cfg.firmware.uefi.capsuleAuthentication.preSignCommands}
-      bash ${flash-tools-patched}/generate_capsule/l4t_generate_soc_capsule.sh \
-    '' + (lib.optionalString cfg.firmware.uefi.capsuleAuthentication.enable ''
-      --trusted-public-cert ${cfg.firmware.uefi.capsuleAuthentication.trustedPublicCertPemFile} \
-      --other-public-cert ${cfg.firmware.uefi.capsuleAuthentication.otherPublicCertPemFile} \
-      --signer-private-cert ${cfg.firmware.uefi.capsuleAuthentication.signerPrivateCertPemFile} \
-    '') + ''
-      -i ${bup}/bl_only_payload \
-      -o $out \
-      ${socType}
-    '');
-
   fuseScript = writeShellApplication {
     name = "fuse-${hostName}";
     text = import ./flash-script.nix {
       inherit lib;
-      flash-tools = flash-tools-patched;
-      flashCommands =
-        let
-          chipId =
-            if cfg.som == null then null
-            else if lib.hasPrefix "orin-" cfg.som then "0x23"
-            else if lib.hasPrefix "xavier-" cfg.som then "0x19"
-            else throw "Unknown SoC type";
-        in
-        ''
-          ./odmfuse.sh -i ${chipId} "$@" ${builtins.toString fuseArgs}
-        '';
+      inherit (nvidia-jetpack) flash-tools;
+      flashCommands = ''
+        ./odmfuse.sh -i ${chipId} "$@" ${builtins.toString fuseArgs}
+      '';
 
       # Fuse script needs device tree files, which aren't already present for
       # non-devkit boards, so we need to get our built version of them
       dtbsDir = config.hardware.deviceTree.package;
     };
+    meta.platforms = [ "x86_64-linux" ];
   };
 in
 {
-  inherit (tosImage) nvLuksSrv hwKeyAgent;
-  inherit mkFlashScript mkFlashCmdScript mkFlashScriptAuto;
-  inherit flashScript initrdFlashScript tosImage taDevKit signedFirmware bup fuseScript uefiCapsuleUpdate;
-  inherit mkRcmBootScript rcmBoot;
+  inherit
+    flashScript
+    fuseScript
+    initrdFlashScript
+    mkFlashCmdScript
+    mkFlashScriptAuto
+    mkRcmBootScript
+    rcmBoot
+    signedFirmware
+    ;
 }
diff --git a/device-pkgs/flash-script.nix b/device-pkgs/flash-script.nix
index 4b2abdc..4a29fb4 100644
--- a/device-pkgs/flash-script.nix
+++ b/device-pkgs/flash-script.nix
@@ -1,5 +1,5 @@
+flash-tools:
 { lib
-, flash-tools
 , preFlashCommands ? ""
 , flashCommands ? ""
 , postFlashCommands ? ""
@@ -17,9 +17,8 @@
   eksFile ? null
 , # Additional DTB overlays to use during device flashing
   additionalDtbOverlays ? [ ]
-,
 }:
-''
+(''
   set -euo pipefail
 
   if [[ -z ''${WORKDIR-} ]]; then
@@ -72,4 +71,4 @@
   ./flash.sh ${lib.optionalString (partitionTemplate != null) "-c flash.xml"} "$@" ${builtins.toString flashArgs}
 '') + ''
   ${postFlashCommands}
-''
+'')
diff --git a/flake.nix b/flake.nix
index 9e56326..b9f4e08 100644
--- a/flake.nix
+++ b/flake.nix
@@ -3,22 +3,20 @@
     nixpkgs.url = "github:nixos/nixpkgs/nixos-23.11";
   };
 
-  outputs = { self, nixpkgs, ... }@inputs:
+  outputs = { self, nixpkgs, ... }:
     let
       inherit (nixpkgs) lib;
 
       installer_minimal_config = {
         imports = [
           "${nixpkgs}/nixos/modules/installer/cd-dvd/installation-cd-minimal.nix"
-          ./modules/default.nix
+          self.nixosModules.default
         ];
         # Avoids a bunch of extra modules we don't have in the tegra_defconfig, like "ata_piix",
         disabledModules = [ "profiles/all-hardware.nix" ];
 
         hardware.nvidia-jetpack.enable = true;
       };
-
-      x86_packages = nixpkgs.legacyPackages.x86_64-linux.callPackage ./default.nix { };
     in
     {
       nixosConfigurations = {
@@ -31,19 +29,47 @@
       overlays.default = import ./overlay.nix;
 
       packages = {
-        x86_64-linux = {
-          # TODO: Untested
-          iso_minimal = self.nixosConfigurations.installer_minimal_cross.config.system.build.isoImage;
-
-          inherit (x86_packages)
-            board-automation python-jetson;
-          inherit (x86_packages.cudaPackages)
-            nsight_systems_host nsight_compute_host;
-        }
-        # Flashing and board automation scripts _only_ work on x86_64-linux
-        // x86_packages.flashScripts
-        // x86_packages.initrdFlashScripts
-        // x86_packages.uefiCapsuleUpdates;
+        x86_64-linux =
+          let
+            supportedConfigurations = lib.listToAttrs (map
+              (c: {
+                name = "${c.som}-${c.carrierBoard}";
+                value = c;
+              }) [
+              { som = "orin-agx"; carrierBoard = "devkit"; }
+              { som = "orin-agx-industrial"; carrierBoard = "devkit"; }
+              { som = "orin-nx"; carrierBoard = "devkit"; }
+              { som = "orin-nano"; carrierBoard = "devkit"; }
+              { som = "xavier-agx"; carrierBoard = "devkit"; }
+              { som = "xavier-nx"; carrierBoard = "devkit"; }
+              { som = "xavier-nx-emmc"; carrierBoard = "devkit"; }
+            ]);
+
+            supportedNixOSConfigurations = lib.mapAttrs
+              (n: c: (nixpkgs.legacyPackages.x86_64-linux.pkgsCross.aarch64-multiplatform.nixos {
+                imports = [ self.nixosModules.default ];
+                hardware.nvidia-jetpack = { enable = true; } // c;
+                networking.hostName = "${c.som}-${c.carrierBoard}"; # Just so it sets the flash binary name.
+              }).config)
+              supportedConfigurations;
+
+            flashScripts = lib.mapAttrs' (n: c: lib.nameValuePair "flash-${n}" c.system.build.flashScript) supportedNixOSConfigurations;
+            initrdFlashScripts = lib.mapAttrs' (n: c: lib.nameValuePair "initrd-flash-${n}" c.system.build.initrdFlashScript) supportedNixOSConfigurations;
+            uefiCapsuleUpdates = lib.mapAttrs' (n: c: lib.nameValuePair "uefi-capsule-update-${n}" c.system.build.uefiCapsuleUpdate) supportedNixOSConfigurations;
+          in
+          {
+            # TODO: Untested
+            iso_minimal = self.nixosConfigurations.installer_minimal_cross.config.system.build.isoImage;
+
+            inherit (self.legacyPackages.x86_64-linux)
+              board-automation python-jetson;
+            inherit (self.legacyPackages.x86_64-linux.cudaPackages)
+              nsight_systems_host nsight_compute_host;
+          }
+          # Flashing and board automation scripts _only_ work on x86_64-linux
+          // flashScripts
+          // initrdFlashScripts
+          // uefiCapsuleUpdates;
 
         aarch64-linux = {
           iso_minimal = self.nixosConfigurations.installer_minimal.config.system.build.isoImage;
@@ -57,7 +83,7 @@
         self.legacyPackages;
 
       # Not everything here should be cross-compiled to aarch64-linux
-      legacyPackages.x86_64-linux = nixpkgs.legacyPackages.x86_64-linux.callPackage ./default.nix { };
-      legacyPackages.aarch64-linux = nixpkgs.legacyPackages.aarch64-linux.callPackage ./default.nix { };
+      legacyPackages.x86_64-linux = (import nixpkgs { system = "x86_64-linux"; overlays = [ self.overlays.default ]; }).nvidia-jetpack;
+      legacyPackages.aarch64-linux = (import nixpkgs { system = "aarch64-linux"; overlays = [ self.overlays.default ]; }).nvidia-jetpack;
     };
 }
diff --git a/kernel/default.nix b/kernel/default.nix
index 8adf253..8981e78 100644
--- a/kernel/default.nix
+++ b/kernel/default.nix
@@ -1,20 +1,15 @@
-{ pkgs
+{ applyPatches
 , lib
 , fetchFromGitHub
-, fetchpatch
 , l4t-xusb-firmware
 , realtime ? false
 , kernelPatches ? [ ]
 , structuredExtraConfig ? { }
-, extraMeta ? { }
 , argsOverride ? { }
+, buildLinux
 , ...
 }@args:
-let
-  isNative = pkgs.stdenv.isAarch64;
-  pkgsAarch64 = if isNative then pkgs else pkgs.pkgsCross.aarch64-multiplatform;
-in
-pkgsAarch64.buildLinux (args // {
+buildLinux (args // {
   version = "5.10.120" + lib.optionalString realtime "-rt70";
   extraMeta.branch = "5.10";
 
@@ -22,7 +17,7 @@ pkgsAarch64.buildLinux (args // {
 
   # Using applyPatches here since it's not obvious how to append an extra
   # postPatch. This is not very efficient.
-  src = pkgs.applyPatches {
+  src = applyPatches {
     src = fetchFromGitHub {
       owner = "OE4T";
       repo = "linux-tegra-5.10";
diff --git a/modules/default.nix b/modules/default.nix
index 617f243..042ec67 100644
--- a/modules/default.nix
+++ b/modules/default.nix
@@ -1,4 +1,4 @@
-{ config, lib, pkgs, ... }:
+{ options, config, lib, pkgs, ... }:
 
 let
   inherit (lib)
@@ -82,17 +82,58 @@ in
         description = "Whether to mount the ESP partition on eMMC under /opt/nvidia/esp on Xavier AGX platforms. Needed for capsule updates";
         internal = true;
       };
+
+      hostPkgs = mkOption {
+        type = options.nixpkgs.pkgs.type;
+        default = import pkgs.path {
+          system = "x86_64-linux";
+          inherit (pkgs) config;
+        };
+        apply = p: p.appendOverlays pkgs.overlays;
+        defaultText = ''
+          import pkgs.path {
+            system = "x86_64-linux";
+            inherit (pkgs) overlays config;
+          }
+        '';
+        description = ''
+          The package set that is used to build packages that run on an
+          external host for purposes of flashing/fusing a Jetson device. This
+          defaults to a package set that can run NVIDIA's pre-built x86
+          binaries needed for flashing/fusing Jetson SOMs.
+        '';
+      };
     };
   };
 
   config = mkIf cfg.enable {
-    assertions = [{
-      assertion = (config.virtualisation.docker.enable && config.virtualisation.docker.enableNvidia) -> lib.versionAtLeast config.virtualisation.docker.package.version "25";
-      message = "Docker version < 25 does not support CDI";
-    }];
+    assertions = [
+      {
+        # NixOS provides two main ways to feed a package set into a config:
+        # 1. The options nixpkgs.hostPlatform/nixpkgs.buildPlatform, which are
+        #    used to construct an import of nixpkgs.
+        # 2. The option nixpkgs.pkgs (set by default if you use the pkgs.nixos
+        #    function), which is a pre-configured import of nixpkgs.
+        #
+        # Regardless of how the package set is setup, it _must_ have its
+        # hostPlatform compatible with aarch64 in order to run on the Jetson
+        # platform.
+        assertion = pkgs.stdenv.hostPlatform.isAarch64;
+        message = ''
+          NixOS config has an invalid package set for the Jetson platform. Try
+          setting nixpkgs.hostPlatform to "aarch64-linux" or otherwise using an
+          aarch64-linux compatible package set.
+        '';
+      }
+      {
+        assertion = (config.virtualisation.docker.enable && config.virtualisation.docker.enableNvidia) -> lib.versionAtLeast config.virtualisation.docker.package.version "25";
+        message = "Docker version < 25 does not support CDI";
+      }
+    ];
 
     nixpkgs.overlays = [
       (import ../overlay.nix)
+      (import ../overlay-with-config.nix config)
     ];
 
     boot.kernelPackages =
diff --git a/modules/flash-script.nix b/modules/flash-script.nix
index 6a26b80..7427756 100644
--- a/modules/flash-script.nix
+++ b/modules/flash-script.nix
@@ -9,6 +9,8 @@ let
     types;
 
   cfg = config.hardware.nvidia-jetpack;
+
+  flashTools = cfg.hostPkgs.callPackages (import ../device-pkgs { inherit config pkgs; }) { };
 in
 {
   imports = with lib; [
@@ -323,152 +325,154 @@ in
     };
   };
 
-  config =
-    let
-      devicePkgs = pkgs.nvidia-jetpack.devicePkgsFromNixosConfig config;
-    in
-    {
-      hardware.nvidia-jetpack.flashScript = devicePkgs.flashScript; # Left for backwards-compatibility
-      hardware.nvidia-jetpack.devicePkgs = devicePkgs; # Left for backwards-compatibility
-      system.build.jetsonDevicePkgs = devicePkgs;
-
-      hardware.nvidia-jetpack.flashScriptOverrides.flashArgs = lib.mkAfter (
-        lib.optional (cfg.firmware.secureBoot.pkcFile != null) "-u ${cfg.firmware.secureBoot.pkcFile}" ++
-        lib.optional (cfg.firmware.secureBoot.sbkFile != null) "-v ${cfg.firmware.secureBoot.sbkFile}" ++
-        [ cfg.flashScriptOverrides.configFileName "mmcblk0p1" ]
-      );
-
-      hardware.nvidia-jetpack.flashScriptOverrides.additionalDtbOverlays =
-        let
-          bootOrder = pkgs.runCommand "DefaultBootOrder.dtbo" { nativeBuildInputs = with pkgs.buildPackages; [ dtc ]; } ''
-            export bootOrder=${lib.concatStringsSep "," cfg.firmware.bootOrder}
-            substituteAll ${./uefi-boot-order.dts} keys.dts
-            dtc -I dts -O dtb keys.dts -o $out
-          '';
-          uefiDefaultKeysDtbo = pkgs.runCommand "UefiDefaultSecurityKeys.dtbo" { nativeBuildInputs = with pkgs.buildPackages; [ dtc ]; } ''
-            export pkDefault=$(od -t x1 -An "${cfg.firmware.uefi.secureBoot.defaultPkEslFile}")
-            export kekDefault=$(od -t x1 -An "${cfg.firmware.uefi.secureBoot.defaultKekEslFile}")
-            export dbDefault=$(od -t x1 -An "${cfg.firmware.uefi.secureBoot.defaultDbEslFile}")
-            substituteAll ${./uefi-default-keys.dts} keys.dts
-            dtc -I dts -O dtb keys.dts -o $out
-          '';
-        in
-        (lib.optional (cfg.firmware.bootOrder != null) bootOrder) ++
-        (lib.optional cfg.firmware.uefi.secureBoot.enrollDefaultKeys uefiDefaultKeysDtbo);
-
-      hardware.nvidia-jetpack.flashScriptOverrides.fuseArgs = lib.mkAfter [ cfg.flashScriptOverrides.configFileName ];
-
-      # These are from l4t_generate_soc_bup.sh, plus some additional ones found in the wild.
-      hardware.nvidia-jetpack.firmware.variants =
-        if (cfg.som != null) then
-          (lib.mkOptionDefault (
-            {
-              xavier-agx = [
-                { boardid = "2888"; boardsku = "0001"; fab = "400"; boardrev = "D.0"; fuselevel = "fuselevel_production"; chiprev = "2"; }
-                { boardid = "2888"; boardsku = "0001"; fab = "400"; boardrev = "E.0"; fuselevel = "fuselevel_production"; chiprev = "2"; } # 16GB
-                { boardid = "2888"; boardsku = "0004"; fab = "400"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; } # 32GB
-                { boardid = "2888"; boardsku = "0005"; fab = "402"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; } # 64GB
-              ];
-              xavier-nx = [
-                # Dev variant
-                { boardid = "3668"; boardsku = "0000"; fab = "100"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
-                { boardid = "3668"; boardsku = "0000"; fab = "301"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
-              ];
-              xavier-nx-emmc = [
-                # Prod variant
-                { boardid = "3668"; boardsku = "0001"; fab = "100"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
-                { boardid = "3668"; boardsku = "0003"; fab = "301"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
-              ];
-
-              orin-agx = [
-                { boardid = "3701"; boardsku = "0000"; fab = "300"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; }
-                { boardid = "3701"; boardsku = "0004"; fab = "300"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # 32GB
-                { boardid = "3701"; boardsku = "0005"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # 64GB
-              ];
-
-              orin-agx-industrial = [
-                { boardid = "3701"; boardsku = "0008"; fab = "300"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; }
-              ];
-
-              orin-nx = [
-                { boardid = "3767"; boardsku = "0000"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin NX 16GB
-                { boardid = "3767"; boardsku = "0001"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin NX 8GB
-              ];
-
-              orin-nano = [
-                { boardid = "3767"; boardsku = "0003"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin Nano 8GB
-                { boardid = "3767"; boardsku = "0004"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin Nano 4GB
-                { boardid = "3767"; boardsku = "0005"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin Nano devkit module
-              ];
-            }.${cfg.som}
-          )) else lib.mkOptionDefault [ ];
-
-      systemd.services.setup-jetson-efi-variables = lib.mkIf (cfg.flashScriptOverrides.targetBoard != null) {
-        description = "Setup Jetson OTA UEFI variables";
-        wantedBy = [ "multi-user.target" ];
-        after = [ "opt-nvidia-esp.mount" ];
-        serviceConfig.Type = "oneshot";
-        serviceConfig.ExecStart = "${pkgs.nvidia-jetpack.otaUtils}/bin/ota-setup-efivars ${cfg.flashScriptOverrides.targetBoard}";
-      };
+  config = {
+    hardware.nvidia-jetpack.flashScript = lib.warn "hardware.nvidia-jetpack.flashScript is deprecated, use config.system.build.flashScript" config.system.build.flashScript;
+    hardware.nvidia-jetpack.devicePkgs = lib.warn "hardware.nvidia-jetpack.devicePkgs is deprecated, use pkgs.nvidia-jetpack" pkgs.nvidia-jetpack;
 
-      systemd.services.firmware-update = lib.mkIf (cfg.firmware.autoUpdate && cfg.som != null && cfg.flashScriptOverrides.targetBoard != null) {
-        wantedBy = [ "multi-user.target" ];
-        path = [ pkgs.nvidia-jetpack.otaUtils ];
-        after = [
-          "${utils.escapeSystemdPath config.boot.loader.efi.efiSysMountPoint}.mount"
-          "opt-nvidia-esp.mount"
-        ];
-        unitConfig = {
-          ConditionPathExists = "/sys/devices/virtual/dmi/id/bios_version";
-          # This directory is populated by ota-apply-capsule-update, don't run
-          # if we already have a capsule update present on the ESP.
-          ConditionDirectoryNotEmpty = "!${config.boot.loader.efi.efiSysMountPoint}/EFI/UpdateCapsule";
-        };
-        script =
-          # NOTE: Our intention is to not apply any capsule update if the
-          # user's intention is to "test" a new nixos config without having it
-          # persist across reboots. "nixos-rebuild test" does not append a new
-          # generation to /nix/var/nix/profiles for the system profile, so we
-          # can compare that symlink to /run/current-system to see if our
-          # current active config has been persisted as a generation. Note that
-          # this check _may_ break down if not using nixos-rebuild and using
-          # switch-to-configuration directly, however it is well-documented
-          # that a user would need to self-manage their system profile's
-          # generations if switching a system in that manner.
-          lib.optionalString config.system.switch.enable ''
-            if [[ -L /nix/var/nix/profiles/system ]]; then
-              latest_generation=$(readlink -f /nix/var/nix/profiles/system)
-              current_system=$(readlink -f /run/current-system)
-              if [[ $latest_generation == /nix/store* ]] && [[ $latest_generation != "$current_system" ]]; then
-                echo "Skipping capsule update, current active system not persisted to /nix/var/nix/profiles/system"
-                exit 0
-              fi
-            fi
-          '' + ''
-            CUR_VER=$(cat /sys/devices/virtual/dmi/id/bios_version)
-            NEW_VER=${pkgs.nvidia-jetpack.l4tVersion}
-
-            if [[ "$CUR_VER" != "$NEW_VER" ]]; then
-              echo "Current Jetson firmware version is: $CUR_VER"
-              echo "New Jetson firmware version is: $NEW_VER"
-              echo
-
-              # Set efi vars here as well as in systemd service, in case we're
-              # upgrading from an older nixos generation that doesn't have the
-              # systemd service. Plus, this ota-setup-efivars will be from the
-              # generation we're switching to, which can contain additional
-              # fixes/improvements.
-              ota-setup-efivars ${cfg.flashScriptOverrides.targetBoard}
-
-              ota-apply-capsule-update ${config.system.build.jetsonDevicePkgs.uefiCapsuleUpdate}
-            fi
-          '';
-      };
+    system.build = {
+      jetsonDevicePkgs = lib.warn "system.build.jetsonDevicePkgs is deprecated, use pkgs.nvidia-jetpack" pkgs.nvidia-jetpack;
+
+      inherit (pkgs.nvidia-jetpack) uefiCapsuleUpdate;
+      inherit (flashTools) flashScript initrdFlashScript fuseScript signedFirmware;
+    };
+
+    hardware.nvidia-jetpack.flashScriptOverrides.flashArgs = lib.mkAfter (
+      lib.optional (cfg.firmware.secureBoot.pkcFile != null) "-u ${cfg.firmware.secureBoot.pkcFile}" ++
+      lib.optional (cfg.firmware.secureBoot.sbkFile != null) "-v ${cfg.firmware.secureBoot.sbkFile}" ++
+      [ cfg.flashScriptOverrides.configFileName "mmcblk0p1" ]
+    );
+
+    hardware.nvidia-jetpack.flashScriptOverrides.additionalDtbOverlays =
+      let
+        bootOrder = pkgs.runCommand "DefaultBootOrder.dtbo" { nativeBuildInputs = with pkgs.buildPackages; [ dtc ]; } ''
+          export bootOrder=${lib.concatStringsSep "," cfg.firmware.bootOrder}
+          substituteAll ${./uefi-boot-order.dts} keys.dts
+          dtc -I dts -O dtb keys.dts -o $out
+        '';
+        uefiDefaultKeysDtbo = pkgs.runCommand "UefiDefaultSecurityKeys.dtbo" { nativeBuildInputs = with pkgs.buildPackages; [ dtc ]; } ''
+          export pkDefault=$(od -t x1 -An "${cfg.firmware.uefi.secureBoot.defaultPkEslFile}")
+          export kekDefault=$(od -t x1 -An "${cfg.firmware.uefi.secureBoot.defaultKekEslFile}")
+          export dbDefault=$(od -t x1 -An "${cfg.firmware.uefi.secureBoot.defaultDbEslFile}")
+          substituteAll ${./uefi-default-keys.dts} keys.dts
+          dtc -I dts -O dtb keys.dts -o $out
+        '';
+      in
+      (lib.optional (cfg.firmware.bootOrder != null) bootOrder) ++
+      (lib.optional cfg.firmware.uefi.secureBoot.enrollDefaultKeys uefiDefaultKeysDtbo);
+
+    hardware.nvidia-jetpack.flashScriptOverrides.fuseArgs = lib.mkAfter [ cfg.flashScriptOverrides.configFileName ];
+
+    # These are from l4t_generate_soc_bup.sh, plus some additional ones found in the wild.
+    hardware.nvidia-jetpack.firmware.variants =
+      if (cfg.som != null) then
+        (lib.mkOptionDefault (
+          {
+            xavier-agx = [
+              { boardid = "2888"; boardsku = "0001"; fab = "400"; boardrev = "D.0"; fuselevel = "fuselevel_production"; chiprev = "2"; }
+              { boardid = "2888"; boardsku = "0001"; fab = "400"; boardrev = "E.0"; fuselevel = "fuselevel_production"; chiprev = "2"; } # 16GB
+              { boardid = "2888"; boardsku = "0004"; fab = "400"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; } # 32GB
+              { boardid = "2888"; boardsku = "0005"; fab = "402"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; } # 64GB
+            ];
+            xavier-nx = [
+              # Dev variant
+              { boardid = "3668"; boardsku = "0000"; fab = "100"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
+              { boardid = "3668"; boardsku = "0000"; fab = "301"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
+            ];
+            xavier-nx-emmc = [
+              # Prod variant
+              { boardid = "3668"; boardsku = "0001"; fab = "100"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
+              { boardid = "3668"; boardsku = "0003"; fab = "301"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = "2"; }
+            ];
+
+            orin-agx = [
+              { boardid = "3701"; boardsku = "0000"; fab = "300"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; }
+              { boardid = "3701"; boardsku = "0004"; fab = "300"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # 32GB
+              { boardid = "3701"; boardsku = "0005"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # 64GB
+            ];
+
+            orin-agx-industrial = [
+              { boardid = "3701"; boardsku = "0008"; fab = "300"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; }
+            ];
+
+            orin-nx = [
+              { boardid = "3767"; boardsku = "0000"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin NX 16GB
+              { boardid = "3767"; boardsku = "0001"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin NX 8GB
+            ];
+
+            orin-nano = [
+              { boardid = "3767"; boardsku = "0003"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin Nano 8GB
+              { boardid = "3767"; boardsku = "0004"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin Nano 4GB
+              { boardid = "3767"; boardsku = "0005"; fab = "000"; boardrev = ""; fuselevel = "fuselevel_production"; chiprev = ""; } # Orin Nano devkit module
+            ];
+          }.${cfg.som}
+        )) else lib.mkOptionDefault [ ];
+
+    systemd.services.setup-jetson-efi-variables = lib.mkIf (cfg.flashScriptOverrides.targetBoard != null) {
+      description = "Setup Jetson OTA UEFI variables";
+      wantedBy = [ "multi-user.target" ];
+      after = [ "opt-nvidia-esp.mount" ];
+      serviceConfig.Type = "oneshot";
+      serviceConfig.ExecStart = "${pkgs.nvidia-jetpack.otaUtils}/bin/ota-setup-efivars ${cfg.flashScriptOverrides.targetBoard}";
+    };
 
-      environment.systemPackages = lib.mkIf (cfg.firmware.autoUpdate && cfg.som != null && cfg.flashScriptOverrides.targetBoard != null) [
-        (pkgs.writeShellScriptBin "ota-apply-capsule-update-included" ''
-          ${pkgs.nvidia-jetpack.otaUtils}/bin/ota-apply-capsule-update ${config.system.build.jetsonDevicePkgs.uefiCapsuleUpdate}
-        '')
+    systemd.services.firmware-update = lib.mkIf (cfg.firmware.autoUpdate && cfg.som != null && cfg.flashScriptOverrides.targetBoard != null) {
+      wantedBy = [ "multi-user.target" ];
+      path = [ pkgs.nvidia-jetpack.otaUtils ];
+      after = [
+        "${utils.escapeSystemdPath config.boot.loader.efi.efiSysMountPoint}.mount"
+        "opt-nvidia-esp.mount"
       ];
+      unitConfig = {
+        ConditionPathExists = "/sys/devices/virtual/dmi/id/bios_version";
+        # This directory is populated by ota-apply-capsule-update, don't run
+        # if we already have a capsule update present on the ESP.
+        ConditionDirectoryNotEmpty = "!${config.boot.loader.efi.efiSysMountPoint}/EFI/UpdateCapsule";
+      };
+      script =
+        # NOTE: Our intention is to not apply any capsule update if the
+        # user's intention is to "test" a new nixos config without having it
+        # persist across reboots. "nixos-rebuild test" does not append a new
+        # generation to /nix/var/nix/profiles for the system profile, so we
+        # can compare that symlink to /run/current-system to see if our
+        # current active config has been persisted as a generation. Note that
+        # this check _may_ break down if not using nixos-rebuild and using
+        # switch-to-configuration directly, however it is well-documented
+        # that a user would need to self-manage their system profile's
+        # generations if switching a system in that manner.
+        lib.optionalString config.system.switch.enable ''
+          if [[ -L /nix/var/nix/profiles/system ]]; then
+            latest_generation=$(readlink -f /nix/var/nix/profiles/system)
+            current_system=$(readlink -f /run/current-system)
+            if [[ $latest_generation == /nix/store* ]] && [[ $latest_generation != "$current_system" ]]; then
+              echo "Skipping capsule update, current active system not persisted to /nix/var/nix/profiles/system"
+              exit 0
+            fi
+          fi
+        '' + ''
+          CUR_VER=$(cat /sys/devices/virtual/dmi/id/bios_version)
+          NEW_VER=${pkgs.nvidia-jetpack.l4tVersion}
+
+          if [[ "$CUR_VER" != "$NEW_VER" ]]; then
+            echo "Current Jetson firmware version is: $CUR_VER"
+            echo "New Jetson firmware version is: $NEW_VER"
+            echo
+
+            # Set efi vars here as well as in systemd service, in case we're
+            # upgrading from an older nixos generation that doesn't have the
+            # systemd service. Plus, this ota-setup-efivars will be from the
+            # generation we're switching to, which can contain additional
+            # fixes/improvements.
+            ota-setup-efivars ${cfg.flashScriptOverrides.targetBoard}
+
+            ota-apply-capsule-update ${pkgs.nvidia-jetpack.uefiCapsuleUpdate}
+          fi
+        '';
     };
+
+    environment.systemPackages = lib.mkIf (cfg.firmware.autoUpdate && cfg.som != null && cfg.flashScriptOverrides.targetBoard != null) [
+      (pkgs.writeShellScriptBin "ota-apply-capsule-update-included" ''
+        ${pkgs.nvidia-jetpack.otaUtils}/bin/ota-apply-capsule-update ${pkgs.nvidia-jetpack.uefiCapsuleUpdate}
+      '')
+    ];
+  };
 }
diff --git a/overlay-with-config.nix b/overlay-with-config.nix
new file mode 100644
index 0000000..edcb4ba
--- /dev/null
+++ b/overlay-with-config.nix
@@ -0,0 +1,149 @@
+# device-specific packages that are influenced by the nixos config
+config:
+final: prev: (
+  let
+    cfg = config.hardware.nvidia-jetpack;
+
+    inherit (prev) lib;
+
+    tosArgs = {
+      inherit (final.nvidia-jetpack) socType;
+      inherit (cfg.firmware.optee) taPublicKeyFile;
+      opteePatches = cfg.firmware.optee.patches;
+      extraMakeFlags = cfg.firmware.optee.extraMakeFlags;
+    };
+  in
+  {
+    nvidia-jetpack = prev.nvidia-jetpack.overrideScope (finalJetpack: prevJetpack: {
+      socType =
+        if cfg.som == null then null
+        else if lib.hasPrefix "orin-" cfg.som then "t234"
+        else if lib.hasPrefix "xavier-" cfg.som then "t194"
+        else throw "Unknown SoC type";
+
+      chipId =
+        if cfg.som == null then null
+        else if lib.hasPrefix "orin-" cfg.som then "0x23"
+        else if lib.hasPrefix "xavier-" cfg.som then "0x19"
+        else throw "Unknown SoC type";
+
+      uefi-firmware = prevJetpack.uefi-firmware.override ({
+        bootLogo = cfg.firmware.uefi.logo;
+        debugMode = cfg.firmware.uefi.debugMode;
+        errorLevelInfo = cfg.firmware.uefi.errorLevelInfo;
+        edk2NvidiaPatches = cfg.firmware.uefi.edk2NvidiaPatches;
+        edk2UefiPatches = cfg.firmware.uefi.edk2UefiPatches;
+      } // lib.optionalAttrs cfg.firmware.uefi.capsuleAuthentication.enable {
+        inherit (cfg.firmware.uefi.capsuleAuthentication) trustedPublicCertPemFile;
+      });
+
+      flash-tools = prevJetpack.flash-tools.overrideAttrs ({ patches ? [ ], postPatch ? "", ... }: {
+        patches = patches ++ cfg.flashScriptOverrides.patches;
+        postPatch = postPatch + cfg.flashScriptOverrides.postPatch;
+      });
+
+      tosImage = finalJetpack.buildTOS tosArgs;
+      taDevKit = finalJetpack.buildOpteeTaDevKit tosArgs;
+      inherit (finalJetpack.tosImage) nvLuksSrv hwKeyAgent;
+
+      flashInitrd =
+        let
+          modules = [ "qspi_mtd" "spi_tegra210_qspi" "at24" "spi_nor" ];
+          modulesClosure = prev.makeModulesClosure {
+            rootModules = modules;
+            kernel = config.system.modulesTree;
+            firmware = config.hardware.firmware;
+            allowMissing = false;
+          };
+          jetpack-init = prev.writeScript "init" ''
+            #!${prev.pkgsStatic.busybox}/bin/sh
+            export PATH=${prev.pkgsStatic.busybox}/bin
+            mkdir -p /proc /dev /sys
+            mount -t proc proc -o nosuid,nodev,noexec /proc
+            mount -t devtmpfs none -o nosuid /dev
+            mount -t sysfs sysfs -o nosuid,nodev,noexec /sys
+
+            for mod in ${builtins.toString modules}; do
+              modprobe -v $mod
+            done
+
+            # `signedFirmware` must be built on x86_64, so we make a
+            # concatenated initrd that places `signedFirmware` at a well
+            # known path so that the final initrd can be constructed from
+            # outside the context of this nixos config (which has an
+            # aarch64-linux package-set).
+            if ${lib.getExe finalJetpack.flashFromDevice} /signed-firmware ; then
+              echo "Flashing platform firmware successful. Rebooting now."
+              sync
+              reboot -f
+            else
+              echo "Flashing platform firmware unsuccessful. Entering console"
+              exec ${prev.pkgsStatic.busybox}/bin/sh
+            fi
+          '';
+        in
+        prev.makeInitrd {
+          contents = [
+            { object = jetpack-init; symlink = "/init"; }
+            { object = "${modulesClosure}/lib"; symlink = "/lib"; }
+          ];
+        };
+
+      # mkFlashScript is declared here due to its dependence on values from
+      # `config`, but it is not inherently tied to any one particular
+      # hostPlatform or buildPlatform (for example, it can be used to build
+      # the "bup" entirely on an aarch64 build machine). mkFlashScript's
+      # hostPlatform and buildPlatform is determined by which flash-tools
+      # you give it, so if your flash-tools is for an x86_64-linux
+      # hostPlatform, then mkFlashScript will generate script commands that
+      # will need to be ran on x86_64-linux.
+      mkFlashScript = flash-tools: args: (import ./device-pkgs/flash-script.nix flash-tools) ({
+        inherit lib;
+        inherit (cfg.firmware) eksFile;
+        inherit (cfg.flashScriptOverrides) additionalDtbOverlays flashArgs partitionTemplate;
+        inherit (finalJetpack) tosImage socType uefi-firmware;
+
+        dtbsDir = config.hardware.deviceTree.package;
+      } // args);
+
+      bup = prev.runCommand "bup-${config.networking.hostName}-${finalJetpack.l4tVersion}"
+        {
+          inherit (cfg.firmware.secureBoot) requiredSystemFeatures;
+        }
+        ((finalJetpack.mkFlashScript
+          final.pkgsBuildBuild.nvidia-jetpack.flash-tools # we need flash-tools for the buildPlatform
+          {
+            # TODO: Remove preSignCommands when we switch to using signedFirmware directly
+            flashCommands = cfg.firmware.secureBoot.preSignCommands + lib.concatMapStringsSep "\n"
+              (v: with v;
+              "BOARDID=${boardid} BOARDSKU=${boardsku} FAB=${fab} BOARDREV=${boardrev} FUSELEVEL=${fuselevel} CHIPREV=${chiprev} ./flash.sh ${lib.optionalString (cfg.flashScriptOverrides.partitionTemplate != null) "-c flash.xml"} --no-flash --bup --multi-spec ${builtins.toString cfg.flashScriptOverrides.flashArgs}"
+              )
+              cfg.firmware.variants;
+          }) + ''
+          mkdir -p $out
+          cp -r bootloader/payloads_*/* $out/
+        '');
+
+      # See l4t_generate_soc_bup.sh
+      # python ${edk2-jetson}/BaseTools/BinWrappers/PosixLike/GenerateCapsule -v --encode --monotonic-count 1
+      # NOTE: providing null public certs here will use the test certs in the EDK2 repo
+      uefiCapsuleUpdate = prev.runCommand "uefi-${config.networking.hostName}-${finalJetpack.l4tVersion}.Cap"
+        {
+          nativeBuildInputs = [ prev.buildPackages.python3 prev.buildPackages.openssl ];
+          inherit (cfg.firmware.uefi.capsuleAuthentication) requiredSystemFeatures;
+        }
+        (''
+          ${cfg.firmware.uefi.capsuleAuthentication.preSignCommands}
+          bash ${finalJetpack.flash-tools}/generate_capsule/l4t_generate_soc_capsule.sh \
+        '' + (lib.optionalString cfg.firmware.uefi.capsuleAuthentication.enable ''
+          --trusted-public-cert ${cfg.firmware.uefi.capsuleAuthentication.trustedPublicCertPemFile} \
+          --other-public-cert ${cfg.firmware.uefi.capsuleAuthentication.otherPublicCertPemFile} \
+          --signer-private-cert ${cfg.firmware.uefi.capsuleAuthentication.signerPrivateCertPemFile} \
+        '') + ''
+          -i ${finalJetpack.bup}/bl_only_payload \
+          -o $out \
+          ${finalJetpack.socType}
+        '');
+    });
+  }
+)
diff --git a/overlay.nix b/overlay.nix
index 4706fd1..88d1d4f 100644
--- a/overlay.nix
+++ b/overlay.nix
@@ -1,3 +1,135 @@
-final: prev: {
-  nvidia-jetpack = final.callPackage ./default.nix { };
+# TODO(jared): Get rid of usage of `callPackages` where possible so we can take
+# advantage of scope's `self.callPackage` (callPackages does not exist under
+# `self`).
+
+final: prev:
+let
+  jetpackVersion = "5.1.2";
+  l4tVersion = "35.4.1";
+  cudaVersion = "11.4";
+
+  sourceInfo = import ./sourceinfo {
+    inherit l4tVersion;
+    inherit (prev) lib fetchurl fetchgit;
+  };
+in
+{
+  nvidia-jetpack = prev.lib.makeScope prev.newScope (self: ({
+    inherit jetpackVersion l4tVersion cudaVersion;
+
+    inherit (sourceInfo) debs gitRepos;
+
+    bspSrc = prev.runCommand "l4t-unpacked"
+      {
+        # https://developer.nvidia.com/embedded/jetson-linux-archive
+        # https://repo.download.nvidia.com/jetson/
+        src = prev.fetchurl {
+          url = with prev.lib.versions; "https://developer.download.nvidia.com/embedded/L4T/r${major l4tVersion}_Release_v${minor l4tVersion}.${patch l4tVersion}/release/Jetson_Linux_R${l4tVersion}_aarch64.tbz2";
+          hash = "sha256-crdaDH+jv270GuBmNLtnw4qSaCFV0SBgJtvuSmuaAW8=";
+        };
+        # We use a more recent version of bzip2 here because we hit this bug
+        # extracting nvidia's archives:
+        # https://bugs.launchpad.net/ubuntu/+source/bzip2/+bug/1834494
+        nativeBuildInputs = [ prev.buildPackages.bzip2_1_1 ];
+      } ''
+      bzip2 -d -c $src | tar xf -
+      mv Linux_for_Tegra $out
+    '';
+
+    # Here for convenience, to see what is in upstream Jetpack
+    unpackedDebs = prev.runCommand "unpackedDebs-${l4tVersion}" { nativeBuildInputs = [ prev.buildPackages.dpkg ]; } ''
+      mkdir -p $out
+      ${prev.lib.concatStringsSep "\n" (prev.lib.mapAttrsToList (n: p: "echo Unpacking ${n}; dpkg -x ${p.src} $out/${n}") self.debs.common)}
+      ${prev.lib.concatStringsSep "\n" (prev.lib.mapAttrsToList (n: p: "echo Unpacking ${n}; dpkg -x ${p.src} $out/${n}") self.debs.t234)}
+    '';
+
+    # Also just for convenience,
+    unpackedDebsFilenames = prev.runCommand "unpackedDebsFilenames-${l4tVersion}" { nativeBuildInputs = [ prev.buildPackages.dpkg ]; } ''
+      mkdir -p $out
+      ${prev.lib.concatStringsSep "\n" (prev.lib.mapAttrsToList (n: p: "echo Extracting file list from ${n}; dpkg --fsys-tarfile ${p.src} | tar --list > $out/${n}") self.debs.common)}
+      ${prev.lib.concatStringsSep "\n" (prev.lib.mapAttrsToList (n: p: "echo Extracting file list from ${n}; dpkg --fsys-tarfile ${p.src} | tar --list > $out/${n}") self.debs.t234)}
+    '';
+
+    unpackedGitRepos = prev.runCommand "unpackedGitRepos-${l4tVersion}" { } (
+      prev.lib.mapAttrsToList
+        (relpath: repo: ''
+          mkdir -p $out/${relpath}
+          cp --no-preserve=all -r ${repo}/. $out/${relpath}
+        '')
+        self.gitRepos
+    );
+
+    inherit (prev.callPackages ./pkgs/uefi-firmware { inherit (self) l4tVersion; })
+      edk2-jetson uefi-firmware;
+
+    inherit (prev.callPackages ./pkgs/optee {
+      # Nvidia's recommended toolchain is gcc9:
+      # https://nv-tegra.nvidia.com/r/gitweb?p=tegra/optee-src/nv-optee.git;a=blob;f=optee/atf_and_optee_README.txt;h=591edda3d4ec96997e054ebd21fc8326983d3464;hb=5ac2ab218ba9116f1df4a0bb5092b1f6d810e8f7#l33
+      stdenv = prev.gcc9Stdenv;
+      inherit (self) bspSrc gitRepos l4tVersion;
+    }) buildTOS buildOpteeTaDevKit opteeClient;
+
+    flash-tools = self.callPackage ./pkgs/flash-tools { };
+
+    # Allows automation of Orin AGX devkit
+    board-automation = self.callPackage ./pkgs/board-automation { };
+
+    # Allows automation of Xavier AGX devkit
+    python-jetson = prev.python3.pkgs.callPackage ./pkgs/python-jetson { };
+
+    tegra-eeprom-tool = prev.callPackage ./pkgs/tegra-eeprom-tool { };
+    tegra-eeprom-tool-static = prev.pkgsStatic.callPackage ./pkgs/tegra-eeprom-tool { };
+
+    cudaPackages = prev.callPackages ./pkgs/cuda-packages {
+      inherit (self) debs cudaVersion
+        l4t-3d-core
+        l4t-core
+        l4t-cuda
+        l4t-cupva
+        l4t-multimedia;
+      inherit (prev.cudaPackages) autoAddOpenGLRunpathHook;
+    };
+
+    samples = prev.callPackages ./pkgs/samples {
+      inherit (self) debs cudaVersion cudaPackages l4t-cuda l4t-multimedia l4t-camera;
+      inherit (prev.cudaPackages) autoAddOpenGLRunpathHook;
+    };
+
+    kernelPackagesOverlay = final: prev: {
+      nvidia-display-driver = self.callPackage ./kernel/display-driver.nix { };
+    };
+
+    kernel = self.callPackage ./kernel { kernelPatches = [ ]; };
+    kernelPackages = (prev.linuxPackagesFor self.kernel).extend self.kernelPackagesOverlay;
+
+    rtkernel = self.callPackage ./kernel { kernelPatches = [ ]; realtime = true; };
+    rtkernelPackages = (prev.linuxPackagesFor self.rtkernel).extend self.kernelPackagesOverlay;
+
+    nxJetsonBenchmarks = self.callPackage ./pkgs/jetson-benchmarks {
+      targetSom = "nx";
+    };
+    xavierAgxJetsonBenchmarks = self.callPackage ./pkgs/jetson-benchmarks {
+      targetSom = "xavier-agx";
+    };
+    orinAgxJetsonBenchmarks = self.callPackage ./pkgs/jetson-benchmarks {
+      targetSom = "orin-agx";
+    };
+
+    flashFromDevice = self.callPackage ./pkgs/flash-from-device { };
+
+    otaUtils = self.callPackage ./pkgs/ota-utils { };
+
+    l4tCsv = self.callPackage ./pkgs/containers/l4t-csv.nix { };
+    genL4tJson = prev.runCommand "l4t.json" { nativeBuildInputs = [ prev.buildPackages.python3 ]; } ''
+      python3 ${./pkgs/containers/gen_l4t_json.py} ${self.l4tCsv} ${self.unpackedDebsFilenames} > $out
+    '';
+    containerDeps = self.callPackage ./pkgs/containers/deps.nix { };
+    nvidia-ctk = self.callPackage ./pkgs/containers/nvidia-ctk.nix { };
+
+    # TODO(jared): deprecate this
+    devicePkgsFromNixosConfig = config: config.system.build.jetsonDevicePkgs;
+  } // (prev.callPackages ./pkgs/l4t {
+    inherit l4tVersion;
+    inherit (sourceInfo) debs;
+  })));
 }
diff --git a/pkgs/cuda-packages/default.nix b/pkgs/cuda-packages/default.nix
index 3c41b34..2a04059 100644
--- a/pkgs/cuda-packages/default.nix
+++ b/pkgs/cuda-packages/default.nix
@@ -15,7 +15,11 @@ in
 , expat
 , pkg-config
 , substituteAll
-, l4t
+, l4t-3d-core
+, l4t-core
+, l4t-cuda
+, l4t-cupva
+, l4t-multimedia
 , zlib
 , qt6
 , xorg
@@ -238,7 +242,7 @@ let
     libcusolver = buildFromSourcePackage { name = "libcusolver"; buildInputs = [ cudaPackages.libcublas ]; };
     libcusparse = buildFromSourcePackage { name = "libcusparse"; };
     libnpp = buildFromSourcePackage { name = "libnpp"; };
-    libcudla = buildFromSourcePackage { name = "libcudla"; buildInputs = [ l4t.l4t-cuda ]; };
+    libcudla = buildFromSourcePackage { name = "libcudla"; buildInputs = [ l4t-cuda ]; };
     nsight_compute_target = buildFromDebs {
       name = "nsight-compute-target";
       version = nsight_compute_version;
@@ -452,7 +456,7 @@ let
         version = (lib.head tensorrtDebs).version;
         srcs = builtins.map (deb: deb.src) tensorrtDebs;
 
-        buildInputs = (with cudaPackages; [ cuda_cudart libcublas libcudla cudnn ]) ++ (with l4t; [ l4t-core l4t-multimedia ]);
+        buildInputs = (with cudaPackages; [ cuda_cudart libcublas libcudla cudnn ]) ++ ([ l4t-core l4t-multimedia ]);
         # Remove unnecessary (and large) static libs
         postPatch = ''
           rm -rf lib/*.a
@@ -474,7 +478,7 @@ let
       version = debs.common.vpi2-dev.version;
       srcs = [ debs.common.libnvvpi2.src debs.common.vpi2-dev.src ];
       sourceRoot = "source/opt/nvidia/vpi2";
-      buildInputs = (with l4t; [ l4t-core l4t-3d-core l4t-multimedia l4t-cupva ])
+      buildInputs = ([ l4t-core l4t-3d-core l4t-multimedia l4t-cupva ])
         ++ (with cudaPackages; [ libcufft libnpp ]);
       patches = [ ./vpi2.patch ];
       postPatch = ''
diff --git a/pkgs/flash-from-device/default.nix b/pkgs/flash-from-device/default.nix
index 7d21f77..53c56b8 100644
--- a/pkgs/flash-from-device/default.nix
+++ b/pkgs/flash-from-device/default.nix
@@ -1,20 +1,21 @@
-{ pkgsAarch64, lib, writeScriptBin, runCommand, tegra-eeprom-tool-static }:
+{ lib, pkgsStatic, runCommand, tegra-eeprom-tool-static }:
 
 let
   # Make the package smaller so it doesn't blow up the initrd size
   staticDeps = runCommand "static-deps" { } ''
     mkdir -p $out/bin
-    cp ${pkgsAarch64.pkgsStatic.mtdutils}/bin/mtd_debug $out/bin
-    cp ${pkgsAarch64.pkgsStatic.mtdutils}/bin/flash_erase $out/bin
+    cp ${pkgsStatic.mtdutils}/bin/mtd_debug $out/bin
+    cp ${pkgsStatic.mtdutils}/bin/flash_erase $out/bin
     cp ${tegra-eeprom-tool-static}/bin/tegra-boardspec $out/bin
   '';
+  name = "flash-from-device";
 in
-runCommand "flash-from-device" { } ''
+runCommand name { meta.mainProgram = name; } ''
   mkdir -p $out/bin
 
   cat > $out/bin/flash-from-device <<EOF
-  #!${pkgsAarch64.pkgsStatic.busybox}/bin/sh
-  export PATH="${lib.makeBinPath [ pkgsAarch64.pkgsStatic.busybox staticDeps ]}:$PATH"
+  #!${pkgsStatic.busybox}/bin/sh
+  export PATH="${lib.makeBinPath [ pkgsStatic.busybox staticDeps ]}:$PATH"
   EOF
   cat ${./flash-from-device.sh} >> $out/bin/flash-from-device
   substituteInPlace $out/bin/flash-from-device \
diff --git a/pkgs/samples/default.nix b/pkgs/samples/default.nix
index c6abfc5..fa63834 100644
--- a/pkgs/samples/default.nix
+++ b/pkgs/samples/default.nix
@@ -22,7 +22,9 @@
 , vulkan-headers
 , vulkan-loader
 , writeShellApplication
-, l4t
+, l4t-cuda
+, l4t-multimedia
+, l4t-camera
 , cudaPackages
 , cudaVersion
 , debs
@@ -232,7 +234,7 @@ let
 
     nativeBuildInputs = [ dpkg python3 ];
     buildInputs = [ libX11 libdrm libglvnd opencv2 vulkan-headers vulkan-loader ]
-      ++ (with l4t; [ l4t-cuda l4t-multimedia l4t-camera ])
+      ++ ([ l4t-cuda l4t-multimedia l4t-camera ])
       ++ (with cudaPackages; [ cudatoolkit tensorrt ]);
 
     # Usually provided by pkg-config, but the samples don't use it.