fix: update Dockerfile to use CUDA base image and install Python dependencies, adjust segmentation model in configuration

Dockerfile

@@ -1,20 +1,26 @@
-FROM python:3.10-slim
+FROM nvidia/cuda:12.6.3-runtime-ubuntu22.04
 
 # Set environment variables
 ENV PYTHONUNBUFFERED=1 \
     PYTHONDONTWRITEBYTECODE=1 \
     HF_HOME=/home/user/.cache/huggingface \
     VISUALIZER_CONFIG=visualizer.hf.toml \
-    HOME=/home/user
+    HOME=/home/user \
+    DEBIAN_FRONTEND=noninteractive
 
-# Install system dependencies (git for compphoto/Intrinsic installation, ffmpeg, glib for OpenCV)
+# Install Python and system dependencies
 RUN apt-get update && \
     apt-get install -y --no-install-recommends \
+        python3.10 \
+        python3.10-venv \
+        python3-pip \
         git \
         ffmpeg \
         libglib2.0-0 \
         libgomp1 \
         build-essential && \
+    ln -sf /usr/bin/python3.10 /usr/bin/python3 && \
+    ln -sf /usr/bin/python3 /usr/bin/python && \
     rm -rf /var/lib/apt/lists/*
 
 # Set up a new user named "user" with UID 1000 for Hugging Face permissions
@@ -23,12 +29,11 @@ RUN useradd -m -u 1000 user
 WORKDIR /app
 
 # Copy requirements files first
-COPY requirements-base.txt ./
+COPY requirements-base.txt requirements-gpu-cu126.txt ./
 
-# Install CPU PyTorch/Torchvision first, then other base requirements
+# Install CUDA PyTorch/Torchvision first, then other base requirements
 RUN pip install --no-cache-dir --upgrade pip && \
-    pip install --no-cache-dir torch torchvision --extra-index-url https://download.pytorch.org/whl/cpu && \
-    pip install --no-cache-dir -r requirements-base.txt
+    pip install --no-cache-dir -r requirements-gpu-cu126.txt
 
 # Copy the rest of the application files
 COPY --chown=user:1000 . .

app.py

@@ -945,7 +945,7 @@ def build_segmentation_bundle(contents: bytes):
 
     t0 = time.perf_counter()
     img = Image.open(io.BytesIO(contents)).convert("RGB")
-    MAX_DIM = 640
+    MAX_DIM = 1280
     if max(img.width, img.height) > MAX_DIM:
         scale = MAX_DIM / max(img.width, img.height)
         img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS)

visualizer.hf.toml

@@ -3,8 +3,8 @@
 #   VISUALIZER_CONFIG=visualizer.hf.toml uvicorn app:app --host 0.0.0.0 --port 7860
 
 [models]
-segmentation_model = "segformer"
-segformer_model_name = "nvidia/segformer-b2-finetuned-ade-512-512"
+segmentation_model = "oneformer"
+oneformer_model_name = "shi-labs/oneformer_ade20k_swin_large"
 depth_model_name = "Intel/dpt-large"
 intrinsic_model_version = "v2"