docker

Dockerfile

@@ -0,0 +1,52 @@
+FROM nvcr.io/nvidia/tensorrt:23.12-py3
+
+# Rest of the Dockerfile remains the same...
+# Set working directory
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    git \
+    ffmpeg \
+    libsndfile1 \
+    build-essential \
+    python3-dev \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python packages with --no-deps for torch to avoid CUDA dependencies
+RUN pip install --no-cache-dir \
+    torch --no-deps \
+    && pip install --no-cache-dir \
+    cuda-python \
+    librosa \
+    tqdm \
+    filetype \
+    imageio \
+    opencv-python-headless \
+    scikit-image \
+    cython \
+    imageio-ffmpeg \
+    colored \
+    numpy==2.0.1
+
+# Clone the repository (without checkpoints)
+RUN git clone https://github.com/antgroup/ditto-talkinghead .
+
+# Build Cython extensions
+RUN cd core/utils/blend && \
+    cython blend.pyx && \
+    gcc -shared -pthread -fPIC -fwrapv -O2 -Wall -fno-strict-aliasing \
+        -I$(python3 -c "import sysconfig; print(sysconfig.get_paths()['include'])") \
+        -I$(python3 -c "import numpy; print(numpy.get_include())") \
+        blend.c -o blend_impl.so
+
+# Set environment variables
+ENV PYTHONPATH=/app
+
+# Command to run inference (can be overridden)
+CMD ["python", "inference.py", \
+     "--data_root", "./checkpoints/ditto_trt_Ampere_Plus", \
+     "--cfg_pkl", "./checkpoints/ditto_cfg/v0.4_hubert_cfg_trt.pkl", \
+     "--audio_path", "./example/audio.wav", \
+     "--source_path", "./example/image.png", \
+     "--output_path", "./output/result.mp4"]

README.md

@@ -166,6 +166,61 @@ python script/cvt_onnx_to_trt.py --onnx_dir "./checkpoints/ditto_onnx" --trt_dir
 Then run `inference.py` with `--data_root=./checkpoints/ditto_trt_custom`.
 
 
+
+
+## Docker + nvidia runtime container
+# Warning - (ubuntu docker + gpu will NOT WORK with Docker Desktop)
+https://docs.docker.com/desktop/features/gpu/
+
+
+Build the container:
+```shell
+./build.sh
+```
+
+
+Clone the checkpoints as above to the host
+
+
+Run the container with GPU support:
+
+```shell
+./run.sh
+```
+Or to run with custom input files:
+
+```shell
+docker run --gpus all \
+  -v $(pwd)/input:/app/input \
+  -v $(pwd)/output:/app/output \
+  ditto-talkinghead \
+  python inference.py \
+    --data_root "./checkpoints/ditto_trt_Ampere_Plus" \
+    --cfg_pkl "./checkpoints/ditto_cfg/v0.4_hubert_cfg_trt.pkl" \
+    --audio_path "/app/input/your_audio.wav" \
+    --source_path "/app/input/your_image.png" \
+    --output_path "/app/output/result.mp4"
+```
+
+To run the container you need nvidia runtime container
+
+# Setup the package repository and GPG key
+curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
+
+curl -fsSL https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
+  sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
+  sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
+
+# Update package listing and install
+sudo apt-get update
+sudo apt-get install -y nvidia-container-toolkit
+
+# Configure the Docker daemon to recognize NVIDIA runtime
+sudo nvidia-ctk runtime configure --runtime=docker
+
+# Restart Docker daemon
+sudo systemctl restart docker
+
 ## 📧 Acknowledgement
 Our implementation is based on [S2G-MDDiffusion](https://github.com/thuhcsi/S2G-MDDiffusion) and [LivePortrait](https://github.com/KwaiVGI/LivePortrait). Thanks for their remarkable contribution and released code! If we missed any open-source projects or related articles, we would like to complement the acknowledgement of this specific work immediately.
 

build.sh

@@ -0,0 +1 @@
+docker build -t ditto-talkinghead .

run.sh

@@ -0,0 +1,10 @@
+docker run --gpus all \
+  -v $(pwd)/input:/app/input \
+  -v $(pwd)/output:/app/output \
+  -v $(pwd)/checkpoints:/app/checkpoints ditto-talkinghead \
+  python inference.py \
+    --data_root "./checkpoints/ditto_trt_Ampere_Plus" \
+    --cfg_pkl "./checkpoints/ditto_cfg/v0.4_hubert_cfg_trt.pkl" \
+    --audio_path "./example/audio.wav" \
+    --source_path "./example/image.png" \
+    --output_path "/app/output/result.mp4"