Update Dockerfile

Dockerfile

@@ -1,20 +1,30 @@
-FROM python:3.10-slim
+FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
 
 WORKDIR /app
 
-RUN apt-get update && apt-get install -y \
+ENV PYTHONUNBUFFERED=1 \
+    HF_HUB_DISABLE_SYMLINKS_WARNING=1
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
     git \
     build-essential \
  && rm -rf /var/lib/apt/lists/*
 
-RUN pip install --no-cache-dir \
-    fastapi \
-    uvicorn \
-    torch \
-    transformers \
-    accelerate \
-    bitsandbytes
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir \
+      fastapi \
+      uvicorn[standard] \
+      "huggingface-hub>=0.23.2" \
+      "transformers==4.44.2" \
+      accelerate \
+      bitsandbytes \
+      torchvision \
+      pillow \
+      requests \
+      flash-attn --no-build-isolation
 
 COPY app.py /app/app.py
 
+EXPOSE 7860
+
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]