Update Dockerfile

Dockerfile

@@ -1,21 +1,19 @@
-# Use an official NVIDIA base image with PyTorch, CUDA, and cuDNN pre-installed
-# We choose a version with Python 3.11 to match your previous environment
+# 1. Start from the same great NVIDIA base image
 FROM nvcr.io/nvidia/pytorch:24.04-py3
 
-# Set the working directory
+# 2. Set the working directory
 WORKDIR /repository
 
-# Copy the requirements file into the container
-COPY requirements.txt .
+# 3. Add Hugging Face specific environment variables
+ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
+ENV PORT=80
 
-# Install your Python dependencies
-# Using --no-cache-dir is good practice for smaller image sizes
-RUN pip install --no-cache-dir -r requirements.txt
+# 4. Copy and install requirements AND the HF toolkit
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt huggingface_inference_toolkit
 
-# Copy the rest of your application code (e.g., your handler.py)
+# 5. Copy the rest of your application code
 COPY . .
 
-# (Optional) Set an environment variable if needed by your platform
-# ENV TRANSFORMERS_CACHE="/repository/cache"
-
-# Your inference platform will typically handle the CMD or ENTRYPOINT
+# 6. THE CRUCIAL PART: Set the command to run the inference server
+CMD ["python", "-m", "huggingface_inference_toolkit.main", "--model_dir", "."]