FROM pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime

ENV DEBIAN_FRONTEND=noninteractive

# Set the MKL_THREADING_LAYER environment variable to GNU
ENV MKL_THREADING_LAYER=GNU

RUN apt-get update && apt-get install -y git libgl1-mesa-glx libglib2.0-0 build-essential cmake aria2

RUN useradd -m -u 1000 user

USER user

ENV HOME=/home/user \
	PATH=/home/user/.local/bin:$PATH \
    PYTHONPATH=$HOME/app \
	PYTHONUNBUFFERED=1 \
	GRADIO_ALLOW_FLAGGING=never \
	GRADIO_NUM_PORTS=1 \
	GRADIO_SERVER_NAME=0.0.0.0 \
	GRADIO_THEME=huggingface \
    GRADIO_SHARE=False \
	SYSTEM=spaces

# Set the working directory to the user's home directory
WORKDIR $HOME/app

RUN git clone -b dev https://github.com/camenduru/HandRefiner.git $HOME/app 
RUN git clone -b dev --recursive https://github.com/camenduru/MeshGraphormer $HOME/app/MeshGraphormer

# Install dependencies
RUN pip install -q pytorch-lightning omegaconf einops yacs trimesh mediapipe rtree boto3
RUN pip install -q git+https://gitlab.eecs.umich.edu/ngv-python-modules/opendr
RUN pip install -q https://download.pytorch.org/whl/cu121/xformers-0.0.22.post7-cp310-cp310-manylinux2014_x86_64.whl

RUN pip install --no-cache-dir gradio albumentations azureml chumpy einops matplotlib mediapipe numpy omegaconf opencv_contrib_python opencv_python opencv_python_headless Pillow pytorch_lightning pytorch_pretrained_bert safetensors scipy timm torchvision==0.15.1 tqdm transformers trimesh xformers yacs
    
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl -d $HOME/app/MeshGraphormer/src/modeling/data -o basicModel_neutral_lbs_10_207_0_v1.0.0.pkl
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/MANO_RIGHT.pkl -d $HOME/app/MeshGraphormer/src/modeling/data -o MANO_RIGHT.pkl
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/graphormer_release/graphormer_3dpw_state_dict.bin -d $HOME/app/MeshGraphormer/models/graphormer_release -o graphormer_3dpw_state_dict.bin
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/graphormer_release/graphormer_h36m_state_dict.bin -d $HOME/app/MeshGraphormer/models/graphormer_release -o graphormer_h36m_state_dict.bin
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/graphormer_release/graphormer_hand_state_dict.bin -d $HOME/app/MeshGraphormer/models/graphormer_release -o graphormer_hand_state_dict.bin
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/raw/main/hrnet/cls_hrnet_w40_sgd_lr5e-2_wd1e-4_bs32_x100.yaml -d $HOME/app/MeshGraphormer/models/hrnet -o cls_hrnet_w40_sgd_lr5e-2_wd1e-4_bs32_x100.yaml
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/raw/main/hrnet/cls_hrnet_w64_sgd_lr5e-2_wd1e-4_bs32_x100.yaml -d $HOME/app/MeshGraphormer/models/hrnet -o cls_hrnet_w64_sgd_lr5e-2_wd1e-4_bs32_x100.yaml
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/hrnet/hrnetv2_w40_imagenet_pretrained.pth -d $HOME/app/MeshGraphormer/models/hrnet -o hrnetv2_w40_imagenet_pretrained.pth
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/hrnet/hrnetv2_w64_imagenet_pretrained.pth -d $HOME/app/MeshGraphormer/models/hrnet -o hrnetv2_w64_imagenet_pretrained.pth

WORKDIR $HOME/app/MeshGraphormer
RUN pip install -e .
RUN pip install -e ./manopth/.

RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/hand_landmarker.task -d $HOME/app/preprocessor -o hand_landmarker.task
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/inpaint_depth_control.ckpt -d $HOME/app/models -o inpaint_depth_control.ckpt
RUN aria2c --console-log-level=error -c -x 16 -s 16 -k 1M https://huggingface.co/camenduru/HandRefiner/resolve/main/ckpt200-multisc-pred.zip -d $HOME/app/MeshGraphormer/predictions -o ckpt200-multisc-pred.zip

WORKDIR $HOME/app
RUN mkdir output
COPY app.py .

# Set the environment variable to specify the GPU device
ENV CUDA_DEVICE_ORDER=PCI_BUS_ID
ENV CUDA_VISIBLE_DEVICES=0

# Run your app.py script
CMD ["python", "app.py"]