Update Dockerfile

Dockerfile

@@ -1,32 +1,27 @@
-# 1. Start from the official NVIDIA CUDA development image. This is the key change.
-FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
+# 1. Start from the PyTorch DEVEL image. This is the correct, full-featured base.
+FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-devel
 
-# 2. Set up the environment to be non-interactive and add Python paths
+# 2. Set environment variables. The LD_LIBRARY_PATH is good practice but may be redundant
+# in this well-configured image. It doesn't hurt to keep it.
 ENV DEBIAN_FRONTEND=noninteractive
-ENV PATH="/root/.local/bin:${PATH}"
+ENV HF_HOME="/data/huggingface"
+ENV UV_CACHE_DIR="/data/uv_cache"
+ENV LD_LIBRARY_PATH=/usr/local/nvidia/lib64:$LD_LIBRARY_PATH
 
-# 3. Install Python 3.11, pip, and other essential tools
-RUN apt-get update && apt-get install -y --no-install-recommends \
-    python3.11 \
-    python3-pip \
-    python-is-python3 \
-    git \
-    && rm -rf /var/lib/apt/lists/*
-
-# 4. Install uv for fast package installation
+# 3. Install git and uv. The base image has most tools, but we add these.
+RUN apt-get update && apt-get install -y --no-install-recommends git && rm -rf /var/lib/apt/lists/*
 RUN pip install uv
 
-# 5. Copy requirements and install them.
-# The base image is clean, so we need to install everything, including torch.
+# 4. Copy requirements and install them.
 COPY requirements.txt .
 RUN uv pip install --no-cache --system -r requirements.txt
 
-# 6. Copy the application code
+# 5. Copy the application code.
 COPY ./app.py /app/app.py
 WORKDIR /app
 
-# 7. Expose the port
+# 6. Expose the port.
 EXPOSE 7860
 
-# 8. Run the application with uvicorn
+# 7. Run the application with uvicorn.
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]