FROM nvcr.io/nvidia/pytorch:24.10-py3

RUN apt update
RUN DEBIAN_FRONTEND="noninteractive" TZ="US/Pacific" apt install -y zsh git wget bind9-host telnet iputils-ping
RUN DEBIAN_FRONTEND="noninteractive" TZ="US/Pacific" apt install -y ffmpeg libsm6 libxext6 libfreetype6-dev
RUN pip install setproctitle

# required for conversion scripts
RUN pip install omegaconf
RUN pip install safetensors

# download the huggingface models cache
# RUN mkdir -p /mldata/huggingface/hub/
# RUN chmod -R 777 /mldata
# COPY ./mldata_cache /mldata/huggingface
# We do this as separate COPY because we want seperate layers to increase the download speed
# You need one copy for each model which you can view with:
# for i in `ls mldata_cache/hub/` ; do echo COPY ./mldata_cache/huggingface/hub/$i /mldata/huggingface/hub/$i ; done
COPY ./mldata_cache/huggingface/hub/models--depth-anything--Depth-Anything-V2-Base-hf /mldata/huggingface/hub/models--depth-anything--Depth-Anything-V2-Base-hf
COPY ./mldata_cache/huggingface/hub/models--diffusers--controlnet-canny-sdxl-1.0-mid /mldata/huggingface/hub/models--diffusers--controlnet-canny-sdxl-1.0-mid
COPY ./mldata_cache/huggingface/hub/models--diffusers--controlnet-depth-sdxl-1.0-mid /mldata/huggingface/hub/models--diffusers--controlnet-depth-sdxl-1.0-mid
COPY ./mldata_cache/huggingface/hub/models--facebook--sam2.1-hiera-large /mldata/huggingface/hub/models--facebook--sam2.1-hiera-large
COPY ./mldata_cache/huggingface/hub/models--h94--IP-Adapter /mldata/huggingface/hub/models--h94--IP-Adapter
COPY ./mldata_cache/huggingface/hub/models--Hius--DreamFul-V2 /mldata/huggingface/hub/models--Hius--DreamFul-V2
COPY ./mldata_cache/huggingface/hub/models--IDEA-Research--grounding-dino-base /mldata/huggingface/hub/models--IDEA-Research--grounding-dino-base
COPY ./mldata_cache/huggingface/hub/models--lllyasviel--Annotators /mldata/huggingface/hub/models--lllyasviel--Annotators
COPY ./mldata_cache/huggingface/hub/models--lllyasviel--ControlNet /mldata/huggingface/hub/models--lllyasviel--ControlNet
COPY ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11f1e_sd15_tile /mldata/huggingface/hub/models--lllyasviel--control_v11f1e_sd15_tile
COPY ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11f1p_sd15_depth /mldata/huggingface/hub/models--lllyasviel--control_v11f1p_sd15_depth
COPY ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11p_sd15_canny /mldata/huggingface/hub/models--lllyasviel--control_v11p_sd15_canny
COPY ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11p_sd15_mlsd /mldata/huggingface/hub/models--lllyasviel--control_v11p_sd15_mlsd
COPY ./mldata_cache/huggingface/hub/models--openai--clip-vit-large-patch14 /mldata/huggingface/hub/models--openai--clip-vit-large-patch14
COPY ./mldata_cache/huggingface/hub/models--reimager--juggernaut-xl-inpainting /mldata/huggingface/hub/models--reimager--juggernaut-xl-inpainting
COPY ./mldata_cache/huggingface/hub/models--reimager--controlnet-canny-sdxl-white-cabinets-1.0 /mldata/huggingface/hub/models--reimager--controlnet-canny-sdxl-white-cabinets-1.0
COPY ./mldata_cache/huggingface/hub/models--reimager--controlnet-depth-sdxl-white-cabinets-1.0 /mldata/huggingface/hub/models--reimager--controlnet-depth-sdxl-white-cabinets-1.0
COPY ./mldata_cache/huggingface/hub/models--SG161222--RealVisXL_V4.0 /mldata/huggingface/hub/models--SG161222--RealVisXL_V4.0
COPY ./mldata_cache/huggingface/hub/models--stabilityai--stable-diffusion-xl-refiner-1.0 /mldata/huggingface/hub/models--stabilityai--stable-diffusion-xl-refiner-1.0
COPY ./mldata_cache/huggingface/hub/models--TencentARC--t2i-adapter-canny-sdxl-1.0 /mldata/huggingface/hub/models--TencentARC--t2i-adapter-canny-sdxl-1.0
COPY ./mldata_cache/huggingface/hub/models--TencentARC--t2i-adapter-depth-midas-sdxl-1.0 /mldata/huggingface/hub/models--TencentARC--t2i-adapter-depth-midas-sdxl-1.0
COPY ./mldata_cache/huggingface/hub/models--Zigeng--SlimSAM-uniform-77 /mldata/huggingface/hub/models--Zigeng--SlimSAM-uniform-77
COPY ./mldata_cache/huggingface/hub/version_diffusers_cache.txt /mldata/huggingface/hub/version_diffusers_cache.txt
COPY ./mldata_cache/huggingface/hub/version.txt /mldata/huggingface/hub/version.txt
COPY ./mldata_cache/nltk_data /mldata/nltk_data
COPY ./mldata_cache/torch /mldata/torch

# diffusers setup
RUN pip install pytorch_lightning
RUN pip install peft
RUN pip install --src /usr/local/src -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
RUN pip install --src /usr/local/src -e git+https://github.com/openai/CLIP.git@main#egg=clip

# controlnet deps
RUN pip install mediapipe
RUN pip install --src /usr/local/src -e git+https://github.com/patrickvonplaten/controlnet_aux#egg=controlnet_aux
RUN pip install git+https://github.com/huggingface/accelerate.git

# other libraries
RUN pip install --src /usr/local/src -e git+https://github.com/huggingface/transformers.git@v4.46.3#egg=transformers
RUN pip install compel ujson
RUN pip install nltk
RUN pip install fast-pytorch-kmeans
RUN pip install scikit-image
RUN pip install fast_pytorch_kmeans
RUN pip install simple-lama-inpainting

# diffusers
RUN pip install --src /usr/local/src -e git+https://github.com/huggingface/diffusers.git@v0.31.0#egg=diffusers

# SAM 2
RUN SAM2_BUILD_ALLOW_ERRORS=0 pip install --src /usr/local/src -e git+https://github.com/facebookresearch/sam2.git#egg=SAM-2

# this opencv-python fix # https://github.com/facebookresearch/nougat/issues/40
RUN sed -e '/DictValue/d' -i /usr/local/lib/python3.10/dist-packages/cv2/typing/__init__.py

RUN pip install vastai

# Install a specific opencv - remove all others
RUN pip uninstall --yes opencv opencv-python opencv-contrib-python opencv-python-headless opencv-contrib-python-headless
RUN rm -rf /usr/local/lib/python3.10/dist-packages/cv2/*
RUN pip install --no-deps opencv-contrib-python==4.10.0.84

# testing
RUN pip install pytest-cov
RUN pip install pytest-mock

# add the source
ADD ./src /usr/local/src/diffuser-models/src

ENV HOME=/tmp
ENTRYPOINT ["/bin/sh", "-c"]
