31 lines
1.3 KiB
Docker
31 lines
1.3 KiB
Docker
# Combined GPU worker: YOLO detect + Real-ESRGAN upscale + Whisper transcribe.
|
|
# CUDA 12.1 + torch cu121 (runs on RTX 3090/4090 driver-550 hosts). Build ctx = repo root.
|
|
FROM nvidia/cuda:12.1.1-cudnn8-runtime-ubuntu22.04
|
|
|
|
ENV DEBIAN_FRONTEND=noninteractive \
|
|
PYTHONUNBUFFERED=1 \
|
|
PIP_NO_CACHE_DIR=1
|
|
|
|
RUN apt-get update && apt-get install -y --no-install-recommends \
|
|
python3 python3-pip libgl1 libglib2.0-0 && \
|
|
ln -sf /usr/bin/python3 /usr/bin/python && \
|
|
rm -rf /var/lib/apt/lists/*
|
|
|
|
WORKDIR /app
|
|
|
|
RUN pip install --upgrade pip && \
|
|
pip install torch==2.1.2 torchvision==0.16.2 --index-url https://download.pytorch.org/whl/cu121
|
|
COPY workers/gpu-vision/requirements.txt .
|
|
RUN pip install -r requirements.txt
|
|
|
|
# Fix basicsr's removed torchvision import (the #1 Real-ESRGAN crash) — rewrite it.
|
|
RUN F=$(python3 -c "import basicsr, os; print(os.path.join(os.path.dirname(basicsr.__file__), 'data', 'degradations.py'))") && \
|
|
sed -i 's/from torchvision.transforms.functional_tensor import rgb_to_grayscale/from torchvision.transforms.functional import rgb_to_grayscale/' "$F"
|
|
|
|
# YOLO construction model + handler.
|
|
COPY workers/gpu-vision/best.pt .
|
|
ENV YOLO_MODEL=/app/best.pt
|
|
COPY workers/gpu-vision/handler.py .
|
|
|
|
CMD ["python3", "-u", "handler.py"]
|