wip split layers

Signed-off-by: Raphael Glon <[email protected]>

Author: oOraph

dockerfiles/pytorch/Dockerfile

@@ -33,13 +33,39 @@ RUN apt-get update && \
         python3.11 \
         libsndfile1-dev \
         ffmpeg \
+        gnupg2 \
+        wget \
     && apt-get clean autoremove --yes \
     && rm -rf /var/lib/{apt,cache,log}
 
 # Copying only necessary files as filtered by .dockerignore
 COPY . .
 
-RUN if [[ "$NEURONX" == "1" ]];then /bin/bash -c "./dockerfiles/pytorch/neuronx.sh";else pip install --no-cache-dir -U pip ".[torch, st, diffusers]";fi
+RUN if [[ "$NEURONX" != "1" ]];then pip install --no-cache-dir -U pip ".[torch, st, diffusers]";fi
+
+RUN if [[ "$NEURONX" == "1" ]];then . /etc/os-release && echo "deb https://apt.repos.neuron.amazonaws.com ${VERSION_CODENAME} main" > /etc/apt/sources.list.d/neuron.list;fi
+
+RUN if [[ "$NEURONX" == "1" ]];then wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB | apt-key add -;fi
+
+RUN if [[ "$NEURONX" == "1" ]];then apt-get update -y && \
+    apt-get install -y --no-install-recommends \
+        aws-neuronx-dkms=2.* \
+        aws-neuronx-collectives=2.* \
+        aws-neuronx-runtime-lib=2.* \
+        aws-neuronx-tools=2.* && \
+    apt-get clean autoremove --yes && \
+    rm -rf /var/lib/{apt,cache,log};fi
+
+RUN if [[ "$NEURONX" == "1" ]];then pip install -U pip && \
+    # Taken from optimum neuron, tgi dockerfile
+    pip3 install \
+        neuronx-cc==2.13.66.0 \
+        torch-neuronx==2.1.2.2.1.0 \
+        transformers-neuronx==0.10.0.21 \
+        --extra-index-url=https://pip.repos.neuron.amazonaws.com && \
+    pip3 install --extra-index-url=https://pip.repos.neuron.amazonaws.com optimum[neuronx,diffusers] && \
+    pip install ".[st,torch-neuronx]";fi
+
 
 # copy application
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit