with import <nixos-unstable>
{
  config = { allowUnfree = true; cudaSupport = true; };
  overlays = [
    (
      final: prev: {
        final.python310 = prev.python310.override {
          enableOptimizations = true;
          reproducibleBuild = false;
          self = final.python310;
          buildInputs = [ final.ffmpeg-full ];
        };
        pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
          (
            python-final: python-prev: {
              pytorch-seed = python-final.callPackage ./pythonPackages/pytorch-seed { };
              audiolm-pytorch = python-final.callPackage ./pythonPackages/audiolm-pytorch { };
              vector-quantize-pytorch = python-final.callPackage ./pythonPackages/vector-quantize-pytorch { };
              local-attention = python-final.callPackage ./pythonPackages/local-attention { };
              ema-pytorch = python-final.callPackage ./pythonPackages/ema-pytorch { };

              openai-triton = python-prev.openai-triton-bin;
              torch = python-prev.torch-bin;
              torchaudio = python-prev.torchaudio-bin;

              bark-gui = python-final.callPackage ./pythonPackages/bark-gui.nix { };
            }
          )
        ];
      }
    )
  ];
};

pkgs.mkShell rec {
  name = "barkTestEnv";
  buildInputs = [
    (pkgs.python310.withPackages (ps: with ps; [ bark-gui ]))
  ];
  shellHook = ''
    echo ${python310Packages.bark-gui}
  '';
}