chore: update optimum-neuron version to 0.4.1

Dockerfile

@@ -43,13 +43,16 @@ RUN apt-get update && apt-get install -y \
 
 ENV PATH =/opt/aws/neuron/bin:$PATH
 
-ARG OPTIMUM_NEURON_VERSION=0.4.0
+# Install cpu pytorch to avoid having to install cuda deps: it makes image leaner and builds faster
+RUN pip install torch==2.8.0 torchvision~=0.23 --index-url https://download.pytorch.org/whl/cpu
+# Install optimum-neuron
+ARG OPTIMUM_NEURON_VERSION=0.4.1
 RUN pip install --extra-index-url https://pip.repos.neuron.amazonaws.com --no-cache-dir optimum-neuron[neuronx]==$OPTIMUM_NEURON_VERSION
 # End Neuron setup
 RUN pip install --no-cache-dir pip -U && \
 	pip install --no-cache-dir \
 	datasets \
-	"huggingface-hub>=0.19" "hf_xet>=1.0.0,<2.0.0" "hf-transfer>=0.1.4" "protobuf<4" "click<8.1" "pydantic~=1.0" 
+	"huggingface-hub>=0.31.4" "protobuf<4" "click<8.1" "pydantic~=1.0" 
 RUN \
 	apt-get update && \
 	apt-get install -y curl && \
@@ -76,7 +79,7 @@ COPY --link --chown=1000 ./ /home/user/app
 COPY --from=pipfreeze --link --chown=1000 /tmp/freeze.txt /tmp/freeze.txt
 ENV PYTHONPATH=$HOME/app \
 	PYTHONUNBUFFERED=1 \
-	HF_HUB_ENABLE_HF_TRANSFER=1 \
+	HF_XET_HIGH_PERFORMANCE=1 \
 	GRADIO_ALLOW_FLAGGING=never \
 	GRADIO_FLAGGING_MODE=never \
 	GRADIO_CACHE_EXAMPLES=true \