flake.nix


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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97

{
  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";
        pythonPkgs = pkgs.python3.withPackages (
          ps:
            with ps; [
              mypy
              ruff
              torch
              accelerate
              transformers
              typing-extensions
              psutil
              ninja
              einops
              packaging
              sentence-transformers
              nix-ai-stuff.packages.${pkgs.system}.flash-attn
              # docling
              uv
            ]
        );

        py = with pkgs.python3Packages; [
          mypy
          ruff
          torch
          accelerate
          transformers
          typing-extensions
          psutil
          ninja
          einops
          packaging
          sentence-transformers
          nix-ai-stuff.packages.${pkgs.system}.flash-attn
          # docling
          uv
        ];
        buildInputs = [
          pyright
          pythonPkgs
          pkgs.sqlite
          pkgs.sqlite-vec
          pkgs.virtualenv
          pkgs.python3Packages.venvShellHook
        ];

        # LOLPATH = lib.makeSearchPath "" buildInputs;
        # LOLPATH = lib.makeSearchPath "" py;
        PYTORCH_CUDA_ALLOC_CONF = "max_split_size_mb:128";
        postVenvCreation = ''
          unset SOURCE_DATE_EPOCH
        '';
        shellHook = ''
          fish
        '';
      };
  };
}