Add JAX CPU to L1 worker image, CUDA JAX to GPU image

CPU workers can now run GPU-queue rendering tasks via JAX on CPU.
GPU image overrides with jax[cuda12] for full CUDA acceleration.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

l1/Dockerfile.gpu

@@ -63,8 +63,8 @@ RUN python3 -m pip install --upgrade pip
 COPY l1/requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Install GPU-specific dependencies (CuPy for CUDA 12.x)
-RUN pip install --no-cache-dir cupy-cuda12x
+# Install GPU-specific dependencies (CuPy for CUDA 12.x, JAX with CUDA)
+RUN pip install --no-cache-dir cupy-cuda12x jax[cuda12]
 
 # Install PyNvVideoCodec for zero-copy GPU encoding
 RUN pip install --no-cache-dir PyNvVideoCodec

l1/requirements.txt

@@ -13,6 +13,7 @@ markdown>=3.5.0
 # Common effect dependencies (used by uploaded effects)
 numpy>=1.24.0
 opencv-python-headless>=4.8.0
+jax[cpu]>=0.4.20
 # core (artdag) and common (artdag_common) installed from local dirs in Dockerfile
 psycopg2-binary
 nest_asyncio