Dockerfile.cuda

FROM python:3.10.4 as builder
WORKDIR /app

# Models download stage
RUN pip3 install --no-cache-dir tqdm==4.66.1 requests==2.31.0
RUN mkdir -p ./carvekit/utils/
RUN mkdir -p ./carvekit/ml/files
COPY ./carvekit/__init__.py ./carvekit/__init__.py
RUN touch ./carvekit/ml/__init__.py
RUN touch ./carvekit/utils/__init__.py
COPY ./carvekit/utils/download_models.py ./carvekit/utils/download_models.py
COPY ./carvekit/ml/files/__init__.py ./carvekit/ml/files/__init__.py
COPY ./carvekit/ml/files/models_loc.py ./carvekit/ml/files/models_loc.py
RUN python3 -c "from carvekit.ml.files.models_loc import download_all; download_all();"
RUN rm -rf ./carvekit

FROM pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime
WORKDIR /app

RUN apt-get update && apt-get -y install libgl1 libglib2.0-0  # Install cv2 dep.
COPY --from=builder /root/.cache/carvekit /root/.cache/carvekit

# Install requirements
COPY requirements.txt ./
RUN pip3 install --no-cache-dir -r requirements.txt
COPY requirements_test.txt ./
RUN pip3 install --no-cache-dir -r requirements_test.txt

# Update code
COPY ./ ./

# Install to site-packages to make possible run tests and process images by cli
RUN pip3 install -e ./

ENV CARVEKIT_PORT '5000'
ENV CARVEKIT_HOST '0.0.0.0'
ENV CARVEKIT_SEGMENTATION_NETWORK 'tracer_b7'
ENV CARVEKIT_PREPROCESSING_METHOD 'none'
ENV CARVEKIT_POSTPROCESSING_METHOD 'fba'
ENV CARVEKIT_DEVICE 'cuda'
ENV CARVEKIT_BATCH_SIZE_SEG '5'
ENV CARVEKIT_BATCH_SIZE_MATTING '1'
ENV CARVEKIT_SEG_MASK_SIZE '640'
ENV CARVEKIT_MATTING_MASK_SIZE '2048'
ENV CARVEKIT_AUTH_ENABLE '1'
ENV CARVEKIT_FP16 '0'
ENV CARVEKIT_TRIMAP_PROB_THRESHOLD=231
ENV CARVEKIT_TRIMAP_DILATION=30
ENV CARVEKIT_TRIMAP_EROSION=5

# Tokens will be generated automatically every time the container is restarted if ENV is not set.

# ENV CARVEKIT_ADMIN_TOKEN 'admin_token' # Do not use this env when creating an image as it is not safe.
# ENV CARVEKIT_ALLOWED_TOKENS 'test_token1,test_token2' # Do not use this env when creating an image as it is not safe.

EXPOSE 5000

CMD ["/bin/sh", "-c", "uvicorn carvekit.web.app:app --host $CARVEKIT_HOST --port $CARVEKIT_PORT"]