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

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

# 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/huggingface/hub/` ; do echo COPY --chmod=0777 ./mldata_cache/huggingface/hub/$i /mldata/huggingface/hub/$i ; done
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--depth-anything--Depth-Anything-V2-Base-hf /mldata/huggingface/hub/models--depth-anything--Depth-Anything-V2-Base-hf
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--diffusers--controlnet-depth-sdxl-1.0-mid /mldata/huggingface/hub/models--diffusers--controlnet-depth-sdxl-1.0-mid
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--diffusers--controlnet-canny-sdxl-1.0-mid /mldata/huggingface/hub/models--diffusers--controlnet-canny-sdxl-1.0-mid
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--h94--IP-Adapter /mldata/huggingface/hub/models--h94--IP-Adapter
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--IDEA-Research--grounding-dino-base /mldata/huggingface/hub/models--IDEA-Research--grounding-dino-base
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--IDEA-Research--grounding-dino-tiny /mldata/huggingface/hub/models--IDEA-Research--grounding-dino-tiny
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--lllyasviel--Annotators /mldata/huggingface/hub/models--lllyasviel--Annotators
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--lllyasviel--ControlNet /mldata/huggingface/hub/models--lllyasviel--ControlNet
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11f1e_sd15_tile /mldata/huggingface/hub/models--lllyasviel--control_v11f1e_sd15_tile
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11f1p_sd15_depth /mldata/huggingface/hub/models--lllyasviel--control_v11f1p_sd15_depth
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11p_sd15_canny /mldata/huggingface/hub/models--lllyasviel--control_v11p_sd15_canny
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--lllyasviel--control_v11p_sd15_mlsd /mldata/huggingface/hub/models--lllyasviel--control_v11p_sd15_mlsd
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--openai--clip-vit-large-patch14 /mldata/huggingface/hub/models--openai--clip-vit-large-patch14
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--reimager--HAT_SR /mldata/huggingface/hub/models--reimager--HAT_SR
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--reimager--juggernaut-xl-inpainting /mldata/huggingface/hub/models--reimager--juggernaut-xl-inpainting
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--SG161222--RealVisXL_V5.0 /mldata/huggingface/hub/models--SG161222--RealVisXL_V5.0
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--stabilityai--stable-diffusion-xl-refiner-1.0 /mldata/huggingface/hub/models--stabilityai--stable-diffusion-xl-refiner-1.0
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--TencentARC--t2i-adapter-canny-sdxl-1.0 /mldata/huggingface/hub/models--TencentARC--t2i-adapter-canny-sdxl-1.0
COPY --chmod=0777 ./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 --chmod=0777 ./mldata_cache/huggingface/hub/models--facebook--sam3 /mldata/huggingface/hub/models--facebook--sam3
COPY --chmod=0777 ./mldata_cache/huggingface/hub/models--Zigeng--SlimSAM-uniform-77 /mldata/huggingface/hub/models--Zigeng--SlimSAM-uniform-77
COPY --chmod=0777 ./mldata_cache/nltk_data /mldata/nltk_data
COPY --chmod=0777 ./mldata_cache/torch /mldata/torch

# upgrade base packages in container
RUN pip install torch transformers torchvision --upgrade --no-build-isolation
# RUN MAX_JOBS=4 pip install flash_attn --upgrade --no-build-isolation

# install libraries
RUN pip install setproctitle omegaconf safetensors pytorch_lightning peft compel ujson nltk fast-pytorch-kmeans scikit-image simple-lama-inpainting vastai pytest-cov pytest-mock
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
RUN pip install --src /usr/local/src -e git+https://github.com/patrickvonplaten/controlnet_aux#egg=controlnet_aux

# install huggingface & utils
# timm>1.0.0 required for transformers
RUN pip install 'timm>1.0.0'
# for faster huggingface downloads
RUN pip install hf_xet
RUN pip install --src /usr/local/src -e git+https://github.com/huggingface/transformers.git@v4.57.1#egg=transformers
RUN pip install --src /usr/local/src -e git+https://github.com/huggingface/diffusers.git@v0.35.2#egg=diffusers

# install BasicSR (upscaler)
RUN pip install --no-deps --src /usr/local/src -e git+https://git@github.com/XPixelGroup/BasicSR.git#egg=basicsr

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

# install sam3
RUN git clone https://github.com/facebookresearch/sam3.git /usr/local/src/sam3 && \
    pip install -e /usr/local/src/sam3 && \
    pip install -e "/usr/local/src/sam3[train,dev]"
RUN pip install decord

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

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

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