Update Dockerfile

Dockerfile

@@ -1,26 +1,39 @@
-# Hugging Face Spaces (Docker) — GPU-compatible base (works on CPU too)
-FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
+# GPU-ready base with PyTorch + CUDA 12.1
+FROM pytorch/pytorch:2.3.1-cuda12.1-cudnn8-runtime
 
 ENV DEBIAN_FRONTEND=noninteractive \
     PIP_NO_CACHE_DIR=1 \
     HF_HUB_ENABLE_HF_TRANSFER=1 \
     UVICORN_WORKERS=1 \
-    PORT=7860
+    PORT=7860 \
+    HOME=/app \
+    XDG_CACHE_HOME=/app/.cache \
+    HF_HOME=/app/.cache/huggingface \
+    HF_HUB_CACHE=/app/.cache/huggingface/hub \
+    TRANSFORMERS_CACHE=/app/.cache/huggingface/transformers \
+    HF_DATASETS_CACHE=/app/.cache/huggingface/datasets \
+    TORCH_HOME=/app/.cache/torch \
+    PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
 
-# System deps
 RUN apt-get update && apt-get install -y --no-install-recommends \
     python3-pip python3-dev git build-essential && \
     ln -s /usr/bin/python3 /usr/bin/python && \
     pip install --upgrade pip && \
     rm -rf /var/lib/apt/lists/*
 
-# Copy files
+# Make caches writable
+RUN mkdir -p /app/.cache/huggingface/hub /app/.cache/huggingface/transformers \
+    /app/.cache/huggingface/datasets /app/.cache/torch && \
+    chmod -R 777 /app/.cache
+
 WORKDIR /app
 COPY requirements.txt /app/requirements.txt
-RUN pip install -r /app/requirements.txt
 
-COPY app.py /app/app.py
+# IMPORTANT: torch is already installed in the base image — don't overwrite it.
+# Remove "torch" from requirements.txt to avoid pulling a CPU wheel by mistake.
+RUN sed -i '/^torch==/d' /app/requirements.txt && pip install -r /app/requirements.txt
 
-# Expose & run
+COPY app.py /app/app.py
 EXPOSE 7860
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+