FROM pytorch/torchserve:latest-gpu

# set user root
USER root

RUN pip install --upgrade pip

# Install dependencies
RUN pip install diffusers transformers accelerate invisible-watermark nvgpu google-cloud-storage tensorrt

# Copying model files
COPY ./config.properties /home/model-server/config.properties
COPY ./sketch-model-3.mar /home/model-server/sketch-model-3.mar

USER model-server

# Expose health and prediction listener ports from the image
EXPOSE 7080
EXPOSE 7081

# # Generate MAR file

CMD ["torchserve", \
     "--start", \
     "--ts-config=/home/model-server/config.properties", \
     "--models", \
     "sketch-model-3.mar", \
     "--model-store", \
     "/home/model-server"]