Update Dockerfile

Dockerfile

@@ -1,35 +1,46 @@
-# Use Python 3.9 slim base image
-FROM python:3.9-slim
-
-# Install system dependencies
-RUN apt-get update && apt-get install -y \
-    libgl1-mesa-glx \
+# Base image with CUDA 12.2
+FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04
+
+# System dependencies
+RUN apt-get update -q && \
+    apt-get install -y --no-install-recommends \
+    python3.9 \
+    python3-pip \
+    libgl1 \
     libglib2.0-0 \
     wget \
     git \
     && rm -rf /var/lib/apt/lists/*
 
-# Set working directory
+# CUDA/CuDNN setup
+RUN wget -qO /tmp/cuda-keyring.deb \
+    https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb && \
+    dpkg -i /tmp/cuda-keyring.deb && \
+    apt-get update -q && \
+    apt-get install -y --no-install-recommends \
+    libcudnn8 \
+    libcublas-12-2 && \
+    rm -rf /var/lib/apt/lists/*
+
 WORKDIR /app
 
-# Copy requirements and install dependencies
+# Python dependencies
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Create directories for models and temporary files with proper permissions
-RUN mkdir -p /app/weights/icon_detect /app/weights/icon_caption_florence /app/imgs \
-    && chown -R 1000:1000 /app
+# Create directories
+RUN mkdir -p /app/{weights,imgs} && \
+    chmod -R 777 /app
 
-# Copy application code
+# Application code
 COPY main.py utils.py download_models.py ./
 
-# Download models during build
+# Download models
 RUN python download_models.py
 
-# Set up user for HF Spaces
+# Set up user
 RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"
 
-# Run the application
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]