File size: 6,359 Bytes
a628b37 08cc91a d1afbc8 a628b37 d1afbc8 08cc91a a628b37 08cc91a a628b37 08cc91a a628b37 08cc91a a628b37 08cc91a a628b37 08cc91a a628b37 08cc91a a628b37 08cc91a a628b37 4130d86 a628b37 08cc91a a628b37 8f1aba9 08cc91a d1afbc8 08cc91a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
# thecollabagepatch/magenta:latest - Duplicable x86_64 build
FROM nvidia/cuda:12.6.2-cudnn-runtime-ubuntu22.04
# Ensure CUDA libraries are on loader path
RUN apt-get update && apt-get install -y --no-install-recommends \
cuda-libraries-12-4 && rm -rf /var/lib/apt/lists/*
ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda-12.4/lib64:/usr/local/cuda-12.4/compat:/usr/local/cuda/targets/x86_64-linux/lib:${LD_LIBRARY_PATH}
RUN ln -sf /usr/local/cuda/targets/x86_64-linux/lib /usr/local/cuda/lib64 || true
# Install cuDNN 9.8 for better compatibility
RUN set -eux; \
apt-get update && apt-get install -y --no-install-recommends gnupg ca-certificates curl; \
install -d -m 0755 /usr/share/keyrings; \
curl -fsSL https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub \
| gpg --batch --yes --dearmor -o /usr/share/keyrings/cuda-archive-keyring.gpg; \
apt-get update; \
apt-mark unhold libcudnn9-cuda-12 || true; \
apt-get install -y --no-install-recommends \
'libcudnn9-cuda-12=9.8.*' \
'libcudnn9-dev-cuda-12=9.8.*' \
--allow-downgrades --allow-change-held-packages; \
apt-mark hold libcudnn9-cuda-12 || true; \
ldconfig; \
rm -rf /var/lib/apt/lists/*
# Performance optimizations for L40S/Ada
ENV LD_PRELOAD=/usr/local/cuda/lib64/libcusparse.so.12:/usr/local/cuda/lib64/libcublas.so.12:/usr/local/cuda/lib64/libcublasLt.so.12:/usr/local/cuda/lib64/libcufft.so.11:/usr/local/cuda/lib64/libcusolver.so.11
ENV TF_GPU_ALLOCATOR=cuda_malloc_async
ENV TF_ENABLE_CUBLAS_TF32=1 NVIDIA_TF32_OVERRIDE=1
ENV DEBIAN_FRONTEND=noninteractive \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1 \
TF_FORCE_GPU_ALLOW_GROWTH=true \
XLA_PYTHON_CLIENT_PREALLOCATE=false \
JAX_PLATFORMS=""
SHELL ["/bin/bash", "-c"]
# Install system dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
software-properties-common curl ca-certificates \
build-essential pkg-config git \
libsndfile1 ffmpeg \
&& add-apt-repository ppa:deadsnakes/ppa -y \
&& apt-get update && apt-get install -y --no-install-recommends \
python3.11 python3.11-dev python3.11-venv python3-pip \
&& rm -rf /var/lib/apt/lists/*
# Make python3.11 default and install uv for faster package installs
RUN ln -sf /usr/bin/python3.11 /usr/bin/python && \
ln -sf /usr/bin/python3.11 /usr/bin/python3 && \
python -m pip install --upgrade pip && \
python -m pip install uv
# CRITICAL: Install TensorFlow FIRST to block tensorflow-cpu
# Using generic tf-nightly (no specific date) for duplicability
RUN uv pip install --system tf-nightly
# Install JAX with CUDA support (pinned for stability)
RUN uv pip install --system "jax[cuda12]" jaxlib
# Install base dependencies
RUN uv pip install --system \
absl-py chex gin-config numpy requests tqdm typing-extensions \
google-cloud-storage librosa resampy soundfile sentencepiece
# Clone and install t5x WITHOUT dependencies (avoid tensorflow-cpu)
RUN git clone https://github.com/google-research/t5x.git /t5x && \
cd /t5x && \
git checkout 92c5b467a5964d06c351c7eae4aa4bcd341c7ded && \
uv pip install --system --no-deps -e .
# Install flaxformer without deps
RUN git clone https://github.com/google/flaxformer.git /flaxformer && \
cd /flaxformer && \
git checkout 399ea3a && \
uv pip install --system --no-deps -e .
# Install seqio without deps and PATCH OUT tensorflow_text
RUN git clone https://github.com/google/seqio.git /seqio && \
cd /seqio && \
uv pip install --system --no-deps -e . && \
# CRITICAL FIX: Remove unused tensorflow_text import (not needed by Magenta RT)
sed -i '/import tensorflow_text as tf_text/d' /seqio/seqio/vocabularies.py
# Install airio (t5x dependency) without deps
RUN git clone https://github.com/google/airio.git /airio && \
cd /airio && \
uv pip install --system --no-deps -e .
# Install clu without deps
RUN git clone https://github.com/google/CommonLoopUtils.git /clu && \
cd /clu && \
uv pip install --system --no-deps -e .
# Now install all remaining dependencies these packages need
RUN uv pip install --system \
flax optax orbax-checkpoint \
fiddle cached_property tf2jax \
aqtp etils jestimator \
tensorflow-datasets tfds-nightly \
apache-beam pyyaml rouge-score sacrebleu scipy \
grain-nightly editdistance pyglove
# Patch jestimator for newer JAX - PartitionSpec moved to jax.sharding
RUN sed -i 's|from jax.experimental.pjit import PartitionSpec|from jax.sharding import PartitionSpec|g' \
/usr/local/lib/python3.11/dist-packages/jestimator/amos_helper.py || \
sed -i 's|from jax.experimental.pjit import PartitionSpec|from jax.sharding import PartitionSpec|g' \
/usr/lib/python3.11/dist-packages/jestimator/amos_helper.py || true
# Install magenta-realtime without deps
RUN git clone https://github.com/magenta/magenta-realtime.git /magenta-realtime-src && \
cd /magenta-realtime-src && \
uv pip install --system --no-deps -e .
# API and audio processing dependencies
RUN uv pip install --system \
fastapi uvicorn[standard] python-multipart \
pyloudnorm gradio soxr huggingface_hub
# Ensure compatible protobuf version
RUN uv pip install --system --force-reinstall "protobuf>=5.27.0"
# Create HuggingFace Space user
RUN useradd -m -u 1000 appuser
WORKDIR /home/appuser/app
# Set cache directory
ENV MAGENTA_RT_CACHE_DIR=/home/appuser/.cache/magenta_rt
RUN mkdir -p $MAGENTA_RT_CACHE_DIR && chown -R appuser:appuser /home/appuser/.cache
# Copy application files with proper ownership
COPY --chown=appuser:appuser app.py /home/appuser/app/
COPY --chown=appuser:appuser utils.py /home/appuser/app/
COPY --chown=appuser:appuser jam_worker.py /home/appuser/app/
COPY --chown=appuser:appuser one_shot_generation.py /home/appuser/app/
COPY --chown=appuser:appuser model_management.py /home/appuser/app/
COPY --chown=appuser:appuser documentation.html /home/appuser/app/
COPY --chown=appuser:appuser lil_demo_540p.mp4 /home/appuser/app/
COPY --chown=appuser:appuser magentaRT_rt_tester.html /home/appuser/app/
COPY --chown=appuser:appuser magenta_prompts.js /home/appuser/app/
COPY --chown=appuser:appuser docs/ /home/appuser/app/docs/
USER appuser
EXPOSE 7860
CMD ["bash", "-lc", "python -m uvicorn app:app --host 0.0.0.0 --port ${PORT:-7860}"] |