1 files changed, 6 insertions, 2 deletions
diff --git a/cli.nix b/cli.nix
index f99ad61..ed80160 100644
--- a/cli.nix
+++ b/cli.nix
@@ -1,10 +1,14 @@
 { pkgs ? import <nixpkgs> {}
 , mode ? "mount"
+, fromFlake ? null
 , diskoFile
 , ... }@args:
 let
-  disko = import ./. { inherit (pkgs) lib; };
-  diskFormat = import diskoFile;
+  disko = import ./. { };
+  diskFormat =
+    if fromFlake != null
+    then (builtins.getFlake fromFlake) + "/${diskoFile}"
+    else import diskoFile;
   diskoEval = if (mode == "create") then
     disko.createScript diskFormat pkgs
   else if (mode == "mount") then