File size: 6,677 Bytes
1c7440e
 
08cc91a
d1afbc8
a628b37
d1afbc8
08cc91a
 
 
 
a628b37
08cc91a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a628b37
08cc91a
 
 
 
 
 
 
 
a628b37
1c7440e
08cc91a
a628b37
08cc91a
a628b37
08cc91a
a628b37
 
08cc91a
 
 
a628b37
08cc91a
 
a628b37
 
 
 
 
 
 
 
 
 
1c7440e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a628b37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1c7440e
 
a628b37
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1c7440e
 
a628b37
 
 
 
 
 
c8f33f0
561e7b8
1c7440e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c8f33f0
d1afbc8
 
1c7440e
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
149
150
151
152
153
154
155
156
157
158
# Dockerfile for HuggingFace Space deployment (x86_64)
# Based on working Modal Dockerfile.x86
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="cuda"

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
RUN uv pip install --system tf-nightly

# Install JAX with CUDA support (pinned for stability)
RUN uv pip install --system "jax[cuda12]==0.7.1" jaxlib

# -------------------------------------------------------------------
# HARD STABILIZER: prevent resampy->numba from selecting py<3.10-only wheels
# Also keep numpy < 2 to avoid later numba incompat headaches on py311.
# -------------------------------------------------------------------
RUN printf "numpy<2\nnumba>=0.57.0\nllvmlite>=0.40.0\n" > /tmp/constraints.txt && \
    echo "---- constraints.txt ----" && \
    cat /tmp/constraints.txt && \
    echo "------------------------"

# (Optional but often helpful) install pinned trio first so they land as wheels cleanly,
# then everything else resolves around them.
RUN uv pip install --system -c /tmp/constraints.txt "numpy<2" "numba>=0.57.0" "llvmlite>=0.40.0"

# Install base dependencies (constraints applied)
RUN uv pip install --system -c /tmp/constraints.txt \
    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 . && \
    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 (constraints applied)
RUN uv pip install --system -c /tmp/constraints.txt \
    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 (constraints applied)
RUN uv pip install --system -c /tmp/constraints.txt \
    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"

RUN uv pip install --system fastrtc

# Set working directory and create cache
WORKDIR /app
ENV MAGENTA_RT_CACHE_DIR=/root/.cache/magenta_rt
RUN mkdir -p $MAGENTA_RT_CACHE_DIR

# Copy application files for HuggingFace Space
COPY app.py /app/
COPY utils.py /app/
COPY jam_worker.py /app/
COPY one_shot_generation.py /app/
COPY model_management.py /app/
COPY documentation.html /app/
COPY lil_demo_540p.mp4 /app/
COPY magentaRT_rt_tester.html /app/
COPY magenta_prompts.js /app/
COPY docs/ /app/docs/
COPY fastrtc_magenta.py /app/

EXPOSE 7860
CMD ["python", "-m", "uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]