Dockerfile

FROM python:3.9

# Use an official TensorFlow GPU base image
FROM tensorflow/tensorflow:latest-gpu

# Set the working directory in the container
WORKDIR /code

# Install CUDA drivers (modify the version according to your requirements)
RUN apt-get update && apt-get install -y --no-install-recommends \
    cuda-drivers \
 && apt-get clean \
 && rm -rf /var/lib/apt/lists/*

# Copy requirements and install Python dependencies
COPY ./requirements.txt /code/requirements.txt
RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt

# Copy the rest of the application code
COPY . .

# Set the default command to run your application
CMD ["python", "caption_api.py", "--host", "0.0.0.0", "--port", "7860"]