Dockerfile

# Use NVIDIA CUDA base image with Ubuntu (ARM64 compatible)
FROM nvidia/cuda:13.1.0-runtime-ubuntu24.04

# Avoid prompts from apt
ENV DEBIAN_FRONTEND=noninteractive

# ------------------------------------------------------------
# 2️⃣  Install basic dependencies
# ------------------------------------------------------------
RUN apt-get update && apt-get install -y \
    curl \
    git \
    python3-pip \
    python3-dev \
    build-essential \
    sudo \
    bash \
    ca-certificates \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# ------------------------------------------------------------
# 3️⃣  Install Python dependencies for z-image-turbo
# ------------------------------------------------------------
ENV PYTHONUNBUFFERED=1
RUN curl -LsSf https://astral.sh/uv/install.sh | sh
# Ensure uv is on PATH for this layer and all following ones
ENV PATH="/root/.local/bin:${PATH}"

# Create venv in /opt/venv
RUN uv venv /opt/venv
ENV PATH="/opt/venv/bin:${PATH}"
ENV VIRTUAL_ENV="/opt/venv"

# Copy and install webapp requirements
COPY webapp/requirements.txt /tmp/requirements.txt
# Install PyTorch with CUDA 12.1 support first
RUN uv pip install --no-cache-dir torch torchvision torchaudio --torch-backend cu130

# Install other requirements
RUN uv pip install --no-cache-dir -r /tmp/requirements.txt

# Install diffusers from git (as specified in pyproject.toml)
RUN uv pip install --no-cache-dir git+https://github.com/huggingface/diffusers.git

# ------------------------------------------------------------
# 4️⃣  GPU environment variables
# ------------------------------------------------------------
ENV NVIDIA_VISIBLE_DEVICES=all
ENV NVIDIA_DRIVER_CAPABILITIES=compute,utility

# ------------------------------------------------------------
# 5️⃣  Create a non-root user that will be used by VS Code
# ------------------------------------------------------------
ARG USERNAME=vscode
ARG USER_UID=1000
ARG USER_GID=$USER_UID

RUN groupadd --gid $USER_GID $USERNAME 2>/dev/null || groupmod -n $USERNAME $(getent group $USER_GID | cut -d: -f1) \
    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME 2>/dev/null || usermod -l $USERNAME -d /home/$USERNAME -m $(getent passwd $USER_UID | cut -d: -f1) \
    && echo "$USERNAME ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/$USERNAME \
    && chmod 0440 /etc/sudoers.d/$USERNAME

# ------------------------------------------------------------
# 6️⃣  Create output directory with proper permissions
# ------------------------------------------------------------
RUN mkdir -p /tim/data/z-image-turbo && \
    mkdir -p /home/vscode/.cache/huggingface && \
    chown -R $USERNAME:$USERNAME /tim/data && \
    chown -R $USERNAME:$USERNAME /home/vscode/.cache && \
    chown -R $USERNAME:$USERNAME /opt/venv

# ------------------------------------------------------------
# 7️⃣  Set working directory and copy webapp files
# ------------------------------------------------------------
WORKDIR /workspace
COPY --chown=$USERNAME:$USERNAME webapp /workspace/webapp
COPY --chown=$USERNAME:$USERNAME image_generation.py /workspace/

# ------------------------------------------------------------
# 8️⃣  Set default user
# ------------------------------------------------------------
USER $USERNAME

# ------------------------------------------------------------
# 9️⃣  Expose port for Flask webapp
# ------------------------------------------------------------
EXPOSE 3010

# ------------------------------------------------------------
# 🔟  Run the Flask webapp
# ----/opt/venv/bin/--------------------------------------------------------
CMD ["python3", "/workspace/webapp/app.py"]