edm: expose the singularity image

❯ nom build .#pkgsCuda.some-pkgs-py.edm.image -L
❯ du -hs $(readlink ./result)
3.4G    /nix/store/axhpdc96qgzk720yciwfach93f3xrqby-singularity-image-edm.img
❯ du -hs /scratch/cs/graphics/singularity-images/images/edm.sif
7.6G    /scratch/cs/graphics/singularity-images/images/edm.sif
❯ rsync -LP ./result triton:
❯ ssh triton srun --mem=8G --time=0:05:00 --gres=gpu:a100:1 singularity exec -B /m:/m -B /scratch:/scratch -B /l:/l --nv ./result nixglhost -- python -m edm.example
...
Saving image grid to "imagenet-64x64.png"...
Done

flake.nix

@@ -3,9 +3,11 @@
 
   inputs.dream2nix.url = "github:nix-community/dream2nix";
   inputs.dream2nix.inputs.nixpkgs.follows = "nixpkgs";
+  inputs.nix-gl-host.url = "github:numtide/nix-gl-host";
+  inputs.nix-gl-host.inputs.nixpkgs.follows = "nixpkgs";
   inputs.nixpkgs.url = "github:NixOS/nixpkgs/master";
 
-  outputs = { self, dream2nix, nixpkgs }@inputs:
+  outputs = { self, dream2nix, nix-gl-host, nixpkgs }@inputs:
     let
       inherit (import ./lib/extend-lib.nix {
         inherit inputs;
@@ -64,7 +66,8 @@
           config = {
             allowUnfree = true;
             cudaSupport = true;
-            cudaCapabilities = [ "8.6" ];
+            # Support V100s and A100s on the Aalto's "Triton" and RTX3090 at the lab:
+            cudaCapabilities = [ "7.0" "8.0" "8.6" ];
             cudaEnableForwardCompat = false;
           };
           overlays = [ overlay ];

overlay.nix

@@ -65,6 +65,8 @@ in
 
   opencv4 = if final.config.cudaSupport then prev.opencv4.override { stdenv = final.cudaPackages.backendStdenv; } else prev.opencv4;
 
+  nixglhost = inputs.nix-gl-host.defaultPackage.${prev.system};
+
 } // lib'.optionalAttrs (lib'.versionOlder lib'.version "23.11") {
   # 2023-08-28: NUR still uses the 23.05 channel which doesen't handle pythonPackagesExtensions
   python3 =

python-packages/by-name/ed/edm/package.nix

@@ -14,6 +14,8 @@
 , click
 , pyspng
 , python
+, singularity-tools
+, nixglhost
 }:
 
 buildPythonPackage rec {
@@ -76,6 +78,15 @@ buildPythonPackage rec {
   ];
 
   passthru.pythonWith = python.withPackages (_: [ edm ]);
+  passthru.image = singularity-tools.buildImage {
+    name = "edm";
+    memSize = 4 * 1024; # MiB
+    diskSize = 20 * 1024; # MiB
+    contents = [
+      nixglhost
+      (python.withPackages (_: [ edm ]))
+    ];
+  };
 
   meta = with lib; {
     description = "Elucidating the Design Space of Diffusion-Based Generative Models (EDM)";