# Use NVIDIA's CUDA base image with Ubuntu 22.04
FROM pytorch/pytorch:2.4.1-cuda12.4-cudnn9-devel

ENV DEBIAN_FRONTEND=noninteractive

RUN useradd -m -u 1000 user

# Set home to the user's home directory
ENV HOME=/home/user \
    PATH=/home/user/.local/bin:$PATH

WORKDIR $HOME/app

RUN mkdir -p /home/user/.cache/huggingface/hub && \
    chown -R user:user /home/user/.cache/huggingface

RUN apt-get update && \
    apt-get install -y --no-install-recommends \
    wget \
    git \
    openssh-client \
    build-essential \
    ffmpeg \
    libsndfile1 \
    libffi-dev \
    python3 \
    python3-dev \
    python3-venv \
    python3-distutils \
    python3-pip && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/*

RUN python3 -m pip install --upgrade pip

RUN pip install uv
RUN python -m uv pip install packaging \
    wheel \
    accelerate \
    torch

RUN python -m uv pip install --no-build-isolation git+https://github.com/Zyphra/transformers_zamba2.git
#    git+https://github.com/Dao-AILab/causal-conv1d.git@v1.4.0 \
#    git+https://github.com/state-spaces/mamba@a07ff1b9ad2a4ac8b04eddf5eaaee5004f15aaf1 \

RUN python -m uv pip install gradio

COPY --chown=user main.py $HOME/app

CMD ["python3", "main.py"]