diff --git a/flake.nix b/flake.nix
index 4d2b84d..71b0501 100644
--- a/flake.nix
+++ b/flake.nix
@@ -15,53 +15,50 @@
     {
       self,
       nixpkgs,
-      home-manager,
-      flake-utils,
       chaotic,
       ...
     }@inputs:
-    flake-utils.lib.eachDefaultSystem (
-      system:
-      let
-        pkgs = import nixpkgs {
-          inherit system;
-          overlays = [ self.overlays.default ];
-        };
+    let
+      forAllSystems = nixpkgs.lib.genAttrs [
+        "x86_64-linux"
+        "aarch64-linux"
+      ];
+    in
+    {
+      lib = import ./lib { inherit (nixpkgs) lib; };
 
-        # Extend lib with play utilities
-        lib = nixpkgs.lib.extend (
-          final: prev: {
-            play = import ./lib { lib = final; };
-          }
-        );
+      nixosModules = {
+        play = import ./modules/nixos self;
+        default = self.nixosModules.play;
+      };
 
-      in
-      {
-        packages = {
-          proton-cachyos = pkgs.callPackage ./pkgs/proton-cachyos { };
-          default = self.packages.${system}.proton-cachyos;
-        };
+      homeManagerModules = {
+        play = import ./modules/home self;
+        default = self.homeManagerModules.play;
+      };
 
-        # Expose the extended lib
-        lib = lib;
-      }
-    )
-    // {
-      homeManagerModules.play = import ./modules/home inputs;
-      nixosModules.play = import ./modules/nixos inputs;
-
-      # Default module (home-manager)
-      homeManagerModules.default = self.homeManagerModules.play;
-
-      # Test configuration
       nixosConfigurations.test = nixpkgs.lib.nixosSystem {
         system = "x86_64-linux";
         modules = [
-          self.nixosModules.play
           ./test-config.nix
+          self.nixosModules.play
         ];
       };
 
+      packages = forAllSystems (
+        system:
+        let
+          pkgs = import nixpkgs {
+            inherit system;
+            overlays = [ self.overlays.default ];
+          };
+        in
+        {
+          proton-cachyos = pkgs.callPackage ./pkgs/proton-cachyos { };
+          default = self.packages.${system}.proton-cachyos;
+        }
+      );
+
       # Overlay for packages
       overlays.default = final: prev: {
         proton-cachyos = final.callPackage ./pkgs/proton-cachyos { };
diff --git a/modules/home/default.nix b/modules/home/default.nix
index b426bcc..d405d28 100644
--- a/modules/home/default.nix
+++ b/modules/home/default.nix
@@ -1,4 +1,4 @@
-inputs: {
+self: {
   imports = [
     ./monitors.nix
     ./gamescoperun.nix
@@ -6,5 +6,7 @@ inputs: {
   ];
 
   # Pass inputs to all modules via _module.args
-  _module.args = { inherit inputs; };
+  _module.args = {
+    inputs = self.inputs;
+  };
 }
diff --git a/modules/home/gamescoperun.nix b/modules/home/gamescoperun.nix
index 0f1cad3..8a4ac67 100644
--- a/modules/home/gamescoperun.nix
+++ b/modules/home/gamescoperun.nix
@@ -129,7 +129,6 @@ let
     # Execute gamescope with the final arguments and the command
     exec ${lib.getExe gamescopePackages.gamescope} $final_args -- $argv
   '';
-
 in
 {
   options.play.gamescoperun = {
diff --git a/modules/nixos/default.nix b/modules/nixos/default.nix
index 1ec7fa2..4f9ea1b 100644
--- a/modules/nixos/default.nix
+++ b/modules/nixos/default.nix
@@ -1,4 +1,4 @@
-inputs: {
+self: {
   imports = [
     ./amd.nix
     ./ananicy.nix
@@ -8,5 +8,7 @@ inputs: {
   ];
 
   # Pass inputs to all modules via _module.args
-  _module.args = { inherit inputs; };
+  _module.args = {
+    inputs = self.inputs;
+  };
 }