Dockerfile refactored, conda env added

Dockerfile

@@ -6,12 +6,17 @@ ARG CUDNN_VER=8
 # See possible types: https://hub.docker.com/r/nvidia/cuda/tags?page=1&ordering=last_updated
 ARG IMAGE_TYPE=runtime
 
-FROM nvidia/cuda:${CUDA_VERSION}-${IMAGE_TYPE}-ubuntu20.04 AS builder
+FROM nvidia/cuda:${CUDA_VERSION}-cudnn${CUDNN_VER}-${IMAGE_TYPE}-ubuntu20.04 AS builder
 
 ENV DEBIAN_FRONTEND=noninteractive
 
 RUN apt update && \
-    apt install -y wget \
+    apt install -y software-properties-common && \
+    # the latest version of Git
+    add-apt-repository ppa:git-core/ppa && \
+    apt update && \
+    apt install -y --no-install-recommends \
+    git ca-certificates wget unzip libgl1-mesa-glx \
     && \
     apt autoremove -y && \
     apt clean -y
@@ -32,7 +37,14 @@ WORKDIR /home
 COPY ./environment.yml ./environment.yml 
 
 RUN conda env update -n base --prune --file ./environment.yml && \
-    conda clean -ay && rm ./environment.yml
+    conda clean -ayf && rm ./environment.yml
+
+RUN git clone --single-branch https://github.com/neuralchen/SimSwap.git && \
+    cd ./SimSwap && \
+    rm -rf .git && \
+    ./download-weights.sh && \
+    wget -P ./parsing_model/checkpoint https://github.com/neuralchen/SimSwap/releases/download/1.0/79999_iter.pth
+
 
 FROM nvidia/cuda:${CUDA_VERSION}-cudnn${CUDNN_VER}-${IMAGE_TYPE}-ubuntu20.04
 
@@ -46,30 +58,10 @@ ENV PATH=$CONDA_DIR/bin:$PATH
 ENV DEBIAN_FRONTEND=noninteractive
 
 RUN apt update && \
-    apt install -y software-properties-common && \
-    # the latest version of Git
-    add-apt-repository ppa:git-core/ppa && \
-    apt update && \
-    apt install -y --no-install-recommends \
-    git ca-certificates wget unzip \
-    libgl1-mesa-glx \
-    && \
+    apt install -y --no-install-recommends libgl1-mesa-glx && \
     apt autoremove -y && \
     apt clean -y
 
-WORKDIR /home
-
-RUN git clone --single-branch https://github.com/neuralchen/SimSwap.git && \
-    cd ./SimSwap && \
-    rm -rf .git && \
-    ./download-weights.sh && \
-    python test_video_swapsingle.py \
-        --isTrain false --use_mask --name people \
-        --Arc_path arcface_model/arcface_checkpoint.tar \
-        --pic_a_path ./demo_file/Iron_man.jpg \
-        --video_path ./demo_file/multi_people_1080p.mp4 \
-        --output_path ./output/multi_test_swapsingle.mp4 \
-        --temp_path ./temp_results
+COPY --from=builder /home/SimSwap /home/SimSwap
 
 WORKDIR /home/SimSwap
-

environment.yml

@@ -0,0 +1,16 @@
+channels:
+  - defaults
+  - conda-forge
+  - pytorch
+dependencies:
+  - numpy=1.20.2
+  - opencv=4.5.0
+  - pillow=8.2.0
+  - moviepy
+  - pip=21.2.2
+  - pip:
+    - insightface==0.2.1
+    - onnxruntime-gpu==1.8.1
+    - -f https://download.pytorch.org/whl/torch_stable.html
+    - torch==1.9.0+cu111
+    - torchvision==0.10.0+cu111