blob: a24ebffbf915a0abf41d9cc0e6fe1a7d77f23e69 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
|
{
description = "Torch cuda flake using nix-community cachix";
nixConfig = {
extra-substituters = [
"https://nix-community.cachix.org"
"https://nix-ai-stuff.cachix.org"
"https://ai.cachix.org"
"https://cuda-maintainers.cachix.org"
];
extra-trusted-public-keys = [
"nix-community.cachix.org-1:mB9FSh9qf2dCimDSUo8Zy7bkq5CX+/rkCWyvRCYg3Fs="
"cuda-maintainers.cachix.org-1:0dq3bujKpuEPMCX6U4WylrUDZ9JyUG0VpVZa7CNfq5E="
"ai.cachix.org-1:N9dzRK+alWwoKXQlnn0H6aUx0lU/mspIoz8hMvGvbbc="
"nix-ai-stuff.cachix.org-1:WlUGeVCs26w9xF0/rjyg32PujDqbVMlSHufpj1fqix8="
];
};
inputs = {
nixpkgs.url = "github:nixos/nixpkgs/nixos-unstable";
nix-ai-stuff = {
url = "github:BatteredBunny/nix-ai-stuff";
inputs.nixpkgs.follows = "nixpkgs";
};
};
outputs = {
self,
nixpkgs,
nix-ai-stuff,
...
}: let
pkgs = import nixpkgs {
system = "x86_64-linux";
config.allowUnfree = true;
config.cudaSupport = true;
};
in {
devShell.x86_64-linux = with pkgs;
mkShell rec {
venvDir = "./.venv";
buildInputs = [
(pkgs.python3.withPackages (
ps:
with ps; [
torch
accelerate
transformers
typing-extensions
psutil
ninja
einops
packaging
sentence-transformers
nix-ai-stuff.packages.${pkgs.system}.flash-attn
]
))
pkgs.virtualenv
pkgs.python3Packages.venvShellHook
];
HENLO = "${pkgs.lib.makeLibraryPath buildInputs}";
postVenvCreation = ''
unset SOURCE_DATE_EPOCH
'';
shellHook = ''
fish
'';
};
};
}
|
