From 06073d33cb132e8a5b01b1943a1cd2d405a1b510 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Fri, 21 Mar 2025 14:37:46 +0100
Subject: [PATCH] Better logic to handle KS4 version and update to 4.0.30

---
 kilosort4/Dockerfile | 5 ++++-
 kilosort4/build.sh   | 6 +++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/kilosort4/Dockerfile b/kilosort4/Dockerfile
index 3647f63..10b8d11 100644
--- a/kilosort4/Dockerfile
+++ b/kilosort4/Dockerfile
@@ -2,6 +2,9 @@ FROM nvidia/cuda:12.0.0-base-ubuntu20.04
 
 LABEL maintainer="Alessio Buccino <alessiop.buccino@gmail.com>"
 
+ENV KILOSORT4_VERSION=4.0.30
+
+
 # USER root
 # Ubuntu package installs
 RUN apt update && \
@@ -35,4 +38,4 @@ RUN conda init bash
 RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
 
 # install kilosort4
-RUN pip install kilosort==4.0.18
+RUN pip install kilosort==$KILOSORT4_VERSION
diff --git a/kilosort4/build.sh b/kilosort4/build.sh
index 32763ba..cf9ff94 100755
--- a/kilosort4/build.sh
+++ b/kilosort4/build.sh
@@ -1,3 +1,7 @@
 #!/bin/bash
 
-docker build -t spikeinterface/kilosort4-base:latest -t spikeinterface/kilosort4-base:4.0.18_cuda-12.0.0 .
+# Extract the KILOSORT4_VERSION from the Dockerfile
+KILOSORT4_VERSION=$(grep -E '^ENV KILOSORT4_VERSION=' "Dockerfile" | awk -F= '{print $2}' | tr -d '[:space:]')
+echo "Building kilosort4-base:$KILOSORT4_VERSION"
+
+docker build -t spikeinterface/kilosort4-base:latest -t spikeinterface/kilosort4-base:${KILOSORT4_VERSION}_cuda-12.0.0 .