nixos/modules/hardware/nvidia.nix

{ config, lib, pkgs, ... }:

let inherit (lib) mkOption types;
in {
  options = {
    enableCuda = mkOption {
      type = types.bool;
      default = false;
      description = ''
        Enable CUDA support for Blender and other packages.
        Warning: This can cause very long build times.
      '';
    };
    nvidiaOpen = mkOption {
      type = types.bool;
      default = false;
    };
  };

  config = {
    hardware.graphics.enable = true;

    services.xserver.videoDrivers = [ "nvidia" ];

    hardware.nvidia = {
      modesetting.enable = true;
      powerManagement.enable = false;
      open = config.nvidiaOpen;
      nvidiaSettings = true;
      package = config.boot.kernelPackages.nvidiaPackages.stable;
    };

    # Apply CUDA-specific settings if enabled
    # Dont do all as this has ridiculous high build times
    #nixpkgs.config.cudaSupport = config.enableCuda;
    hardware.enableAllFirmware = config.enableCuda;

    environment.systemPackages = (if config.enableCuda then [
      (pkgs.blender.override { cudaSupport = true; })
      (pkgs.ollama.override { acceleration = "cuda"; })
    ] else
      [ ]);
  };
}