diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..f65c83b --- /dev/null +++ b/.dockerignore @@ -0,0 +1,22 @@ +.git +.github +.vscode +.idea + +__pycache__ +*.pyc +*.pyo +*.log + +.venv +venv +env + +deploy/deps +deploy/tmp +deploy/recordings + +*.pth +*.pt +*.safetensors +*.onnx \ No newline at end of file diff --git a/.github/workflows/build-runpod-h200.yml b/.github/workflows/build-runpod-h200.yml new file mode 100644 index 0000000..faf94fd --- /dev/null +++ b/.github/workflows/build-runpod-h200.yml @@ -0,0 +1,76 @@ +name: Build RunPod H200 container image + +on: + workflow_dispatch: + push: + branches: + - main + - "agent/**" + paths: + - .dockerignore + - .github/workflows/build-runpod-h200.yml + - Dockerfile.h200 + - deploy/** + - runpod/** + +permissions: + contents: read + packages: write + +concurrency: + group: runpod-h200-image-${{ github.ref }} + cancel-in-progress: true + +env: + REGISTRY: ghcr.io + IMAGE_NAME: samuellucky2424-afk/joyai-video-edit + +jobs: + build-and-push: + name: Build and publish H200 image + runs-on: ubuntu-24.04 + timeout-minutes: 330 + + steps: + - name: Free disk space for the CUDA image + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be + with: + tool-cache: false + android: true + dotnet: true + haskell: true + large-packages: true + docker-images: true + swap-storage: false + + - name: Check out the repository + uses: actions/checkout@v6 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@65b78e6e13532edd9afa3aa52ac7964289d1a9c1 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and publish the H200 container image + id: build + uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4 + with: + context: . + file: ./Dockerfile.h200 + platforms: linux/amd64 + push: true + provenance: false + build-args: | + MAX_JOBS=2 + tags: | + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:runpod-h200 + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:h200-sha-${{ github.sha }} + labels: | + org.opencontainers.image.source=${{ github.server_url }}/${{ github.repository }} + org.opencontainers.image.revision=${{ github.sha }} + org.opencontainers.image.description=JoyAI Video Edit for RunPod H200 + + - name: Show immutable deployment reference + run: echo "${REGISTRY}/${IMAGE_NAME}@${{ steps.build.outputs.digest }}" diff --git a/.github/workflows/build-runpod-image.yml b/.github/workflows/build-runpod-image.yml new file mode 100644 index 0000000..f2ab5b1 --- /dev/null +++ b/.github/workflows/build-runpod-image.yml @@ -0,0 +1,65 @@ +name: Build RunPod container image + +on: + workflow_dispatch: + +permissions: + contents: read + packages: write + +concurrency: + group: runpod-image-${{ github.ref }} + cancel-in-progress: false + +env: + REGISTRY: ghcr.io + IMAGE_NAME: samuellucky2424-afk/joyai-video-edit + +jobs: + build-and-push: + name: Build and publish RTX PRO 6000 image + runs-on: ubuntu-24.04 + timeout-minutes: 330 + + steps: + - name: Free disk space for the CUDA image + uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be + with: + tool-cache: false + android: true + dotnet: true + haskell: true + large-packages: true + docker-images: true + swap-storage: false + + - name: Check out the RunPod branch + uses: actions/checkout@v6 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@65b78e6e13532edd9afa3aa52ac7964289d1a9c1 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and publish the container image + id: build + uses: docker/build-push-action@f2a1d5e99d037542a71f64918e516c093c6f3fc4 + with: + context: . + file: ./Dockerfile + platforms: linux/amd64 + push: true + build-args: | + MAX_JOBS=2 + tags: | + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:runpod-rtx-pro-6000 + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:sha-${{ github.sha }} + labels: | + org.opencontainers.image.source=${{ github.server_url }}/${{ github.repository }} + org.opencontainers.image.revision=${{ github.sha }} + org.opencontainers.image.description=JoyAI Video Edit for RunPod RTX PRO 6000 + + - name: Show immutable deployment reference + run: echo "${REGISTRY}/${IMAGE_NAME}@${{ steps.build.outputs.digest }}" diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..14620fe --- /dev/null +++ b/Dockerfile @@ -0,0 +1,80 @@ +# syntax=docker/dockerfile:1.7 + +FROM nvidia/cuda:12.8.1-devel-ubuntu22.04 + +ARG DEBIAN_FRONTEND=noninteractive +ARG MAX_JOBS=8 + +ENV PYTHONUNBUFFERED=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_NO_CACHE_DIR=1 \ + CUDA_HOME=/usr/local/cuda \ + TORCH_CUDA_ARCH_LIST=12.0 \ + JOYOMNI_OPS_CUDA_ARCHS=120a \ + MAX_JOBS=${MAX_JOBS} + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + ffmpeg \ + git \ + libglib2.0-0 \ + libgl1 \ + ninja-build \ + python3 \ + python3-dev \ + python3-pip \ + python3-venv \ + python-is-python3 \ + && rm -rf /var/lib/apt/lists/* \ + && python3 -m pip install --upgrade pip setuptools wheel + +WORKDIR /opt/joyai + +COPY deploy/requirements.txt /tmp/requirements.txt + +RUN python3 -m pip install -r /tmp/requirements.txt + +# Install patched SageAttention for RTX PRO 6000 Blackwell. +COPY deploy/sageattention-cudagraph-stream.patch /tmp/sageattention.patch + +RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \ + && git -C /tmp/SageAttention checkout d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5 \ + && git -C /tmp/SageAttention apply /tmp/sageattention.patch \ + && cd /tmp/SageAttention \ + && EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" python3 setup.py install \ + && rm -rf /tmp/SageAttention + +# Build the repository's original FP8 CUDA operations. +COPY deploy/joyomni_ops /opt/joyai/deploy/joyomni_ops + +RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ + && git -C /tmp/cutlass checkout dcf215af \ + && JOYOMNI_OPS_CUTLASS_DIR=/tmp/cutlass \ + python3 -m pip install --no-build-isolation /opt/joyai/deploy/joyomni_ops \ + && rm -rf /tmp/cutlass /root/.cache/pip + +COPY . /opt/joyai + +RUN chmod +x /opt/joyai/deploy/run_server.sh \ + && mkdir -p /runpod-volume/joyai \ + && mkdir -p /tmp/joyomni-recordings + +ENV JOYOMNI_DEVICE=cuda:0 \ + JOYOMNI_HOST=0.0.0.0 \ + JOYOMNI_PORT=8080 \ + JOYOMNI_CKPT_ROOT=/runpod-volume/joyai/checkpoints \ + JOYOMNI_WIDTH=840 \ + JOYOMNI_HEIGHT=480 \ + JOYOMNI_FPS=24 \ + JOYOMNI_FP8_IMG=1 \ + JOYOMNI_FP8_TXT=1 \ + JOYOMNI_CUDA_GRAPH=1 \ + JOYOMNI_SAGE_ATTN=1 \ + JOYOMNI_TXT_PARALLEL=1 \ + JOYOMNI_RECORD_DIR=/tmp/joyomni-recordings + +EXPOSE 8080 8081 + +CMD ["python3", "/opt/joyai/runpod/start.py"] \ No newline at end of file diff --git a/Dockerfile.h200 b/Dockerfile.h200 new file mode 100644 index 0000000..81b09f3 --- /dev/null +++ b/Dockerfile.h200 @@ -0,0 +1,119 @@ +# syntax=docker/dockerfile:1.7 + +ARG PYTORCH_IMAGE_VERSION=2.9.1-cuda12.8-cudnn9 + +# Compile the Hopper-specific CUDA extensions in a development image. Nothing +# from this stage is shipped except the two finished Python wheels. +FROM pytorch/pytorch:${PYTORCH_IMAGE_VERSION}-devel AS extensions + +ARG DEBIAN_FRONTEND=noninteractive +ARG MAX_JOBS=8 + +ENV PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_NO_CACHE_DIR=1 \ + CUDA_HOME=/usr/local/cuda \ + TORCH_CUDA_ARCH_LIST=9.0 \ + JOYOMNI_OPS_CUDA_ARCHS=90a \ + MAX_JOBS=${MAX_JOBS} + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + git \ + ninja-build \ + && rm -rf /var/lib/apt/lists/* \ + && python -m pip install --upgrade pip setuptools wheel + +WORKDIR /opt/joyai +RUN mkdir -p /wheels + +# SageAttention supports Hopper GPUs. Compile its CUDA extension for H200 +# (sm_90) and keep only the resulting wheel. +COPY deploy/sageattention-cudagraph-stream.patch /tmp/sageattention.patch + +RUN git clone https://github.com/thu-ml/SageAttention.git /tmp/SageAttention \ + && git -C /tmp/SageAttention checkout d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5 \ + && git -C /tmp/SageAttention apply /tmp/sageattention.patch \ + && cd /tmp/SageAttention \ + && EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" \ + python setup.py bdist_wheel --dist-dir /wheels \ + && rm -rf /tmp/SageAttention + +# Build JoyAI's FP8 CUDA operations for Hopper (sm_90a) as a wheel. +# CUTLASS WGMMA instructions are architecture-accelerated and abort at runtime +# when the extension is compiled only for the generic sm_90 target. +COPY deploy/joyomni_ops /opt/joyai/deploy/joyomni_ops + +RUN git clone https://github.com/NVIDIA/cutlass.git /tmp/cutlass \ + && git -C /tmp/cutlass checkout dcf215af \ + && JOYOMNI_OPS_CUTLASS_DIR=/tmp/cutlass \ + python -m pip wheel --no-build-isolation --no-deps \ + --wheel-dir /wheels /opt/joyai/deploy/joyomni_ops \ + && rm -rf /tmp/cutlass /root/.cache/pip + + +# The final image already contains CUDA 12.8, cuDNN 9, PyTorch 2.9.1, and +# torchvision. Keep the small host compiler toolchain as an escape hatch for +# deployments that explicitly re-enable Torch Inductor VAE compilation. +# CUDA extensions and the full CUDA development toolchain stay in the builder. +FROM pytorch/pytorch:${PYTORCH_IMAGE_VERSION}-runtime + +ARG DEBIAN_FRONTEND=noninteractive + +ENV PYTHONUNBUFFERED=1 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + PIP_NO_CACHE_DIR=1 \ + CC=gcc \ + CXX=g++ + +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + ffmpeg \ + libglib2.0-0 \ + libgl1 \ + libgomp1 \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /opt/joyai + +COPY deploy/requirements-h200-runtime.txt /tmp/requirements.txt +RUN python -m pip install -r /tmp/requirements.txt \ + && rm -f /tmp/requirements.txt + +COPY --from=extensions /wheels /tmp/wheels +RUN python -m pip install --no-deps /tmp/wheels/*.whl \ + && rm -rf /tmp/wheels + +COPY . /opt/joyai + +RUN chmod +x /opt/joyai/deploy/run_server.sh \ + && mkdir -p /runpod-volume/joyai + +ENV JOYOMNI_DEVICE=cuda:0 \ + JOYOMNI_HOST=0.0.0.0 \ + JOYOMNI_PORT=8080 \ + JOYOMNI_CKPT_ROOT=/runpod-volume/joyai/checkpoints \ + JOYOMNI_CACHE_ROOT=/runpod-volume/joyai/cache/h200-torch291-cu128 \ + JOYOMNI_CACHE_READY_MARKER=/runpod-volume/joyai/cache/h200-torch291-cu128/ready.json \ + JOYOMNI_PRELOAD=1 \ + JOYOMNI_WIDTH=840 \ + JOYOMNI_HEIGHT=480 \ + JOYOMNI_FPS=20 \ + JOYOMNI_FP8_IMG=1 \ + JOYOMNI_FP8_TXT=1 \ + JOYOMNI_CUDA_GRAPH=1 \ + JOYOMNI_SAGE_ATTN=1 \ + JOYOMNI_TXT_PARALLEL=1 \ + JOYOMNI_VAE_COMPILE=1 \ + JOYOMNI_VAE_COMPILE_STRICT=1 \ + JOYOMNI_LOAD_WARMUP_STRICT=1 \ + JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS=300 \ + JOYOMNI_WARMUP_BOTH_ORIENTATIONS=0 \ + JOYOMNI_WARMUP_REFERENCE_BUCKETS=0 \ + JOYOMNI_RECORD_ENABLED=0 \ + JOYOMNI_ONLINE_GATE_ENABLED=0 + +EXPOSE 8080 8081 + +CMD ["python", "/opt/joyai/runpod/start.py"] diff --git a/deploy/requirements-h200-runtime.txt b/deploy/requirements-h200-runtime.txt new file mode 100644 index 0000000..50955f6 --- /dev/null +++ b/deploy/requirements-h200-runtime.txt @@ -0,0 +1,16 @@ +transformers>=4.57.1,<4.58 +accelerate==1.10.1 +diffusers==0.36.0 +einops==0.8.2 +numpy==2.2.6 +pillow==12.2.0 +opencv-python-headless==4.13.0.92 +av==13.1.0 +imageio-ffmpeg==0.6.0 +fastapi==0.117.1 +uvicorn==0.37.0 +uvloop==0.22.1 +ninja==1.13.0 +websockets==16.0 +openai==2.41.0 +loguru==0.7.3 diff --git a/deploy/run_server.sh b/deploy/run_server.sh index bb47f8f..9c26c96 100755 --- a/deploy/run_server.sh +++ b/deploy/run_server.sh @@ -26,13 +26,17 @@ fi cd "$HERE" -export TORCHINDUCTOR_CACHE_DIR="$HERE/deps/cache/torchinductor" -export TRITON_CACHE_DIR="$HERE/deps/cache/triton" -export CUDA_CACHE_PATH="$HERE/deps/cache/nv_compute" -export TORCHINDUCTOR_FX_GRAPH_CACHE=1 +CACHE_ROOT="${JOYOMNI_CACHE_ROOT:-$HERE/deps/cache}" +export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-$CACHE_ROOT/torchinductor}" +export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-$CACHE_ROOT/triton}" +export CUDA_CACHE_PATH="${CUDA_CACHE_PATH:-$CACHE_ROOT/nv_compute}" +export TORCHINDUCTOR_FX_GRAPH_CACHE="${TORCHINDUCTOR_FX_GRAPH_CACHE:-1}" export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}" mkdir -p "$TORCHINDUCTOR_CACHE_DIR" "$TRITON_CACHE_DIR" "$CUDA_CACHE_PATH" +echo "JoyAI compile cache root: $CACHE_ROOT" +echo "JoyAI VAE compile: ${JOYOMNI_VAE_COMPILE:-1} (strict=${JOYOMNI_VAE_COMPILE_STRICT:-0})" + export PYTHONUNBUFFERED=1 export PYTHONPATH="$HERE" @@ -43,6 +47,33 @@ export JOYOMNI_SAGE_ATTN="${JOYOMNI_SAGE_ATTN:-1}" export JOYOMNI_TXT_PARALLEL="${JOYOMNI_TXT_PARALLEL:-1}" RECORD_DIR="${JOYOMNI_RECORD_DIR:-$HERE/recordings}" +RECORD_ENABLED="${JOYOMNI_RECORD_ENABLED:-1}" +ONLINE_GATE_ENABLED="${JOYOMNI_ONLINE_GATE_ENABLED:-1}" +EXTRA_ARGS=() + +case "${RECORD_ENABLED,,}" in + 1|true|yes|on) + EXTRA_ARGS+=(--record-dir "$RECORD_DIR") + ;; + 0|false|no|off) + ;; + *) + echo "JOYOMNI_RECORD_ENABLED must be one of: 1, 0, true, false, yes, no, on, off" >&2 + exit 2 + ;; +esac + +case "${ONLINE_GATE_ENABLED,,}" in + 1|true|yes|on) + ;; + 0|false|no|off) + EXTRA_ARGS+=(--no-online-gate) + ;; + *) + echo "JOYOMNI_ONLINE_GATE_ENABLED must be one of: 1, 0, true, false, yes, no, on, off" >&2 + exit 2 + ;; +esac CKPT_ROOT="${JOYOMNI_CKPT_ROOT:-$HERE/deps/checkpoints}" DIT_CKPT="${JOYOMNI_DIT_CKPT:-$CKPT_ROOT/JoyAI-Video-Edit/dit/joyai_video_edit_dit_0811.pth}" @@ -61,7 +92,6 @@ python xvideo/serving/serve_joyomni_streaming.py \ --text-encoder-ckpt "$TE_CKPT" \ --face-detector-onnx "$FACE_ONNX" \ --person-detector-onnx "$PERSON_ONNX" \ - --record-dir "$RECORD_DIR" \ --device "$DEVICE" \ --vae-encode-device "$DEVICE" \ --vae-decode-device "$DEVICE" \ @@ -70,4 +100,5 @@ python xvideo/serving/serve_joyomni_streaming.py \ --width "${JOYOMNI_WIDTH:-840}" --height "${JOYOMNI_HEIGHT:-480}" \ --fps "${JOYOMNI_FPS:-24}" \ --host "$HOST" --port "$PORT" \ + "${EXTRA_ARGS[@]}" \ "$@" diff --git a/deploy/static/index.html b/deploy/static/index.html index b7d5ff2..75f3391 100644 --- a/deploy/static/index.html +++ b/deploy/static/index.html @@ -162,8 +162,6 @@ - -
@@ -171,7 +169,6 @@ - @@ -192,21 +189,6 @@
编辑输出 - - -
-
-
等待编辑指令
-
- - -
@@ -241,12 +222,12 @@ 上传清晰度 - + - + - +
-
-
+
+
-
+
-
-
+
+
@@ -608,6 +589,11 @@ try { new ResizeObserver(() => updateCropMask()).observe(camera); } catch (e) {} } let ws = null; +let streamingWanted = false; +let reconnectTimer = null; +let reconnectAttempt = 0; +const RECONNECT_BASE_DELAY_MS = 750; +const RECONNECT_MAX_DELAY_MS = 5000; let stream = null; let videoFileUrl = null; let usingVideoFile = false; @@ -636,8 +622,8 @@ let pendingPeCacheKey = null; let outputPlaybackTimer = null; let outputPlaybackIntervalMs = null; -const DEFAULT_TARGET_OUTPUT_QUEUE_DELAY_MS = 200; -const DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS = 300; +const DEFAULT_TARGET_OUTPUT_QUEUE_DELAY_MS = 100; +const DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS = 200; function targetOutputQueueDelayMs() { const el = document.getElementById("targetQueueDelayMs"); const raw = el && el.value !== "" ? el.value : null; @@ -651,7 +637,7 @@ const v = raw === null ? NaN : Number(raw); return Number.isFinite(v) && v >= 0 ? v : DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS; } -const MAX_BACKEND_PENDING_FRAMES = 32; +const MAX_BACKEND_PENDING_FRAMES = 16; let sendRealtimeSkips = 0; let sentFrames = 0; let backendAckedFrames = 0; @@ -807,6 +793,25 @@ }, 1000); } +function cancelReconnect() { + if (reconnectTimer) clearTimeout(reconnectTimer); + reconnectTimer = null; +} + +function scheduleReconnect() { + if (!streamingWanted || reconnectTimer) return; + const delay = Math.min( + RECONNECT_MAX_DELAY_MS, + RECONNECT_BASE_DELAY_MS * Math.pow(2, Math.min(reconnectAttempt, 3)) + ); + reconnectAttempt += 1; + reconnectTimer = setTimeout(async () => { + reconnectTimer = null; + if (!streamingWanted || ws) return; + await start(); + }, delay); +} + function readRefImageDataUrl() { const el = document.getElementById("refImage"); const file = el && el.files && el.files[0]; @@ -1156,10 +1161,20 @@ } let h264DecodeOk = false; +let h264DecodeReady = Promise.resolve(false); +const CODEC_PROBE_TIMEOUT_MS = 1200; +function codecProbeWithTimeout(probe, fallback = false) { + return Promise.race([ + Promise.resolve(probe).catch(() => fallback), + new Promise((resolve) => setTimeout(() => resolve(fallback), CODEC_PROBE_TIMEOUT_MS)), + ]); +} if (typeof VideoDecoder !== "undefined") { - VideoDecoder.isConfigSupported({ codec: "avc1.640028", optimizeForLatency: true }) - .then((s) => { h264DecodeOk = !!(s && s.supported); }) - .catch(() => {}); + h264DecodeReady = codecProbeWithTimeout( + VideoDecoder.isConfigSupported({ codec: "avc1.640028", optimizeForLatency: true }) + ) + .then((s) => { h264DecodeOk = !!(s && s.supported); return h264DecodeOk; }) + .catch(() => false); } let outCodecH264 = false; let outDecoder = null; @@ -1168,7 +1183,7 @@ let upCodecH264 = false; let upEncoder = null; -const UPLINK_KEYFRAME_INTERVAL = 8; +const UPLINK_KEYFRAME_INTERVAL = 20; let upSeq = 0; const upPendingByTs = new Map(); @@ -1608,6 +1623,7 @@ async function beginSession() { if (!ws || ws.readyState !== WebSocket.OPEN) return false; + h264DecodeOk = await h264DecodeReady; resetRunMetrics(); pePausedSend = false; clearResultVideo(); @@ -1646,7 +1662,10 @@ upCodecH264 = false; if (!usingVideoFile && typeof VideoEncoder !== "undefined") { try { - const s = await VideoEncoder.isConfigSupported(upEncoderConfig(width, height)); + const s = await codecProbeWithTimeout( + VideoEncoder.isConfigSupported(upEncoderConfig(width, height)), + null + ); upCodecH264 = !!(s && s.supported); } catch (e) {} } @@ -1714,10 +1733,16 @@ sessionGranted = true; setSendBusy(true, t("busy_starting")); showOutputStart(true, t("busy_starting")); - beginSession(); + try { + const started = await beginSession(); + if (!started) throw new Error("session start was not sent"); + } catch (err) { + stopActiveRun(); + } return; } if (msg.type === "started") { + reconnectAttempt = 0; if (Number(msg.width) > 0 && Number(msg.height) > 0) { document.getElementById("width").value = msg.width; document.getElementById("height").value = msg.height; @@ -1959,6 +1984,7 @@ setSendBusy(false, t("busy_disconnected")); showOutputIdle(true); updateMetrics(); + scheduleReconnect(); }; socket.onerror = () => { if (ws === socket) { @@ -1969,6 +1995,8 @@ } function stopActiveRun() { + streamingWanted = false; + cancelReconnect(); if (timer) clearInterval(timer); timer = null; stopVideoPull(); @@ -1992,6 +2020,9 @@ currentWs.onclose = null; try { currentWs.close(); } catch (err) {} } + startingRun = false; + sessionGranted = false; + setSendBusy(false, ""); } async function send() { @@ -2012,6 +2043,7 @@ sessionGranted = false; } stopActiveRun(); + streamingWanted = true; await start(); } @@ -2031,8 +2063,9 @@ }, { key: "ref_tryon", - label: "参考图换装", label_en: "Try-On from Reference", + label: "参考图编辑", label_en: "Reference Image", cases: [ + { title: "我的参考人物", title_en: "My Reference Person", desc: "用上传图片中的人物替换主角", desc_en: "Replace the subject with your uploaded person", text: "将视频中的主角替换为上传的参考图中的人物。保持参考人物的身份、面部、发型和服装,同时保留原视频的姿势、动作、镜头和背景。", text_en: "Replace the main subject in the video with the person shown in the uploaded reference image. Preserve the reference person's identity, face, hairstyle, and clothing while keeping the source pose, motion, camera, and background unchanged." }, { title: "NY 棒球帽", title_en: "NY Cap", desc: "戴上参考图里的黑色 NY 棒球帽", desc_en: "Black NY cap from the reference image", ref: "hat", text: "将参考图中的黑色 NY 棒球帽戴到主角头上,帽子正面朝前。", text_en: "Put the black NY baseball cap from the reference image on the subject's head, facing forward." }, { title: "红色围巾", title_en: "Red Scarf", desc: "围上参考图里的红色流苏围巾", desc_en: "Red tasseled scarf from the reference image", ref: "scarf", text: "将参考图中的红色流苏围巾围到主角的颈部。", text_en: "Wrap the red tasseled scarf from the reference image around the subject's neck." }, { title: "粉色T恤", title_en: "Pink Tee", desc: "换上参考图里的粉色短袖 T 恤", desc_en: "Pink short-sleeve tee from the reference image", ref: "pink_tee", text: "将主角的衣服换成参考图里的粉色短袖 T 恤。", text_en: "Replace the subject's clothing with the pink short-sleeve T-shirt from the reference image." }, @@ -2558,7 +2591,7 @@ const Q_TIERS = [20, 40, 60]; const downQualityUi = document.getElementById("downQualityUi"); const qtierEl = document.getElementById("qtier"); -let autoQTier = 2; +let autoQTier = 0; let autoQLowStreak = 0; let autoQHoldTicks = 0; function sendOutputQuality(v) { @@ -2590,7 +2623,7 @@ b.addEventListener("click", () => { const d = b.getAttribute("data-q"); if (d === "auto") { - autoQuality = true; autoQTier = 2; autoQLowStreak = 0; autoQHoldTicks = 0; + autoQuality = true; autoQTier = 0; autoQLowStreak = 0; autoQHoldTicks = 0; highlightQTier("auto"); setDownQuality(Q_TIERS[autoQTier], { live: true }); } else { @@ -2705,6 +2738,7 @@ showCameraOverlay(true); }); }); +window.addEventListener("pagehide", () => stopActiveRun()); diff --git a/deploy/xvideo/checkpoint_status.py b/deploy/xvideo/checkpoint_status.py new file mode 100644 index 0000000..a40d052 --- /dev/null +++ b/deploy/xvideo/checkpoint_status.py @@ -0,0 +1,93 @@ +from __future__ import annotations + +import hashlib +from pathlib import Path +from typing import Any + + +JOYAI_DIT_FILENAME = "dit/joyai_video_edit_dit_0811.pth" +JOYAI_DIT_RELEASE_COMMIT = "eda14f342ef99c52485bbb8dc271c29b42298089" +JOYAI_DIT_XET_HASH = "86a577acfe936e9b56ae7c89f04d2db61d5b69f97f6fa6496da8f7cfcc47305f" +JOYAI_DIT_SHA256 = "b3904b6fda53d13b230918bb616f322d12cfb2337b0e8d9dc203cdabc36605ba" + + +def _metadata_path(checkpoint_path: Path) -> Path: + """Return huggingface_hub's local_dir metadata path for the DiT file.""" + local_dir = checkpoint_path.parent.parent + return local_dir / ".cache" / "huggingface" / "download" / "dit" / ( + checkpoint_path.name + ".metadata" + ) + + +def _read_metadata(checkpoint_path: Path) -> tuple[str | None, str | None, str | None]: + metadata_path = _metadata_path(checkpoint_path) + if not metadata_path.is_file(): + return None, None, None + + try: + with metadata_path.open(encoding="utf-8") as metadata: + revision = metadata.readline().strip() or None + etag = metadata.readline().strip().strip('"') or None + metadata.readline() # timestamp; validate only by successful reading + except (OSError, UnicodeError) as error: + return None, None, f"{type(error).__name__}: {error}" + + return revision, etag, None + + +def checkpoint_status( + checkpoint: str | Path, + *, + full_hash: bool = False, +) -> dict[str, Any]: + """Inspect the RV2V DiT without reading 32.5 GB unless explicitly requested.""" + checkpoint_path = Path(checkpoint) + report: dict[str, Any] = { + "status": "missing", + "path": str(checkpoint_path), + "expected_release_commit": JOYAI_DIT_RELEASE_COMMIT, + "expected_xet_hash": JOYAI_DIT_XET_HASH, + "expected_sha256": JOYAI_DIT_SHA256, + "metadata_revision": None, + "metadata_etag": None, + "sha256": None, + "size_bytes": None, + "verification": "none", + } + + if not checkpoint_path.is_file(): + return report + + try: + report["size_bytes"] = checkpoint_path.stat().st_size + except OSError: + pass + + revision, etag, metadata_error = _read_metadata(checkpoint_path) + report["metadata_revision"] = revision + report["metadata_etag"] = etag + if metadata_error is not None: + report["metadata_error"] = metadata_error + + if etag in {JOYAI_DIT_XET_HASH, JOYAI_DIT_SHA256}: + report["status"] = "current" + report["verification"] = "huggingface_metadata" + elif etag: + report["status"] = "stale" + report["verification"] = "huggingface_metadata" + else: + report["status"] = "unknown" + + if full_hash: + digest = hashlib.sha256() + with checkpoint_path.open("rb") as checkpoint_file: + for block in iter(lambda: checkpoint_file.read(16 * 1024 * 1024), b""): + digest.update(block) + actual_sha256 = digest.hexdigest() + report["sha256"] = actual_sha256 + report["verification"] = "sha256" + report["status"] = ( + "current" if actual_sha256 == JOYAI_DIT_SHA256 else "stale" + ) + + return report diff --git a/deploy/xvideo/models/pipeline.py b/deploy/xvideo/models/pipeline.py index d9da262..3b3759f 100644 --- a/deploy/xvideo/models/pipeline.py +++ b/deploy/xvideo/models/pipeline.py @@ -305,7 +305,8 @@ def _encode_vae_single( from xvideo.models.vae import vae_compile as _vc inputs = _vc.prep_input(inputs) - latents = self.vae.encode(inputs).latent_dist.sample() + with _vc.call_guard(): + latents = self.vae.encode(inputs).latent_dist.sample() if enable_denormalization: latents = self.normalize_latents(latents) return latents.to(original_device) diff --git a/deploy/xvideo/models/vae/vae.py b/deploy/xvideo/models/vae/vae.py index 4bc0d9c..88b8e3d 100644 --- a/deploy/xvideo/models/vae/vae.py +++ b/deploy/xvideo/models/vae/vae.py @@ -6,6 +6,7 @@ import torch from torch import nn, Tensor import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel from diffusers.utils.torch_utils import randn_tensor from diffusers.configuration_utils import ConfigMixin, register_to_config @@ -127,11 +128,15 @@ def attention(self, x: Tensor) -> Tensor: k = self.k(x) v = self.v(x) - q = rearrange(q, "b c t h w -> (b t) 1 (h w) c") - k = rearrange(k, "b c t h w -> (b t) 1 (h w) c") - v = rearrange(v, "b c t h w -> (b t) 1 (h w) c") + q = rearrange(q, "b c t h w -> (b t) 1 (h w) c").contiguous() + k = rearrange(k, "b c t h w -> (b t) 1 (h w) c").contiguous() + v = rearrange(v, "b c t h w -> (b t) 1 (h w) c").contiguous() - x = F.scaled_dot_product_attention(q, k, v) + # This VAE uses a 1024-wide attention head. CUDA flash attention only + # supports head dimensions up to 256, so fused-only SDPA aborts on H200. + # Keep this fallback local to the VAE instead of changing DiT attention. + with sdpa_kernel(SDPBackend.MATH): + x = F.scaled_dot_product_attention(q, k, v) x = rearrange(x, "(b t) 1 (h w) c -> b c t h w", b=b, t=t, h=h, w=w) x = self.proj_out(x) @@ -528,7 +533,7 @@ class XVAEChunkCausal(ModelMixin, ConfigMixin): """For more technical details on high-resolution causal VAE decoding, see: https://github.com/xin1u/UltraFlash """ - + @register_to_config def __init__( self, diff --git a/deploy/xvideo/models/vae/vae_compile.py b/deploy/xvideo/models/vae/vae_compile.py index ff0401a..8d89d83 100644 --- a/deploy/xvideo/models/vae/vae_compile.py +++ b/deploy/xvideo/models/vae/vae_compile.py @@ -1,61 +1,162 @@ from __future__ import annotations +import os +import threading +from contextlib import contextmanager + import torch import torch.nn as nn from xvideo.inductor_autotune_fix import install as _install_autotune_fix +def compile_enabled() -> bool: + """Return whether Torch Inductor VAE compilation is enabled. + + Compilation remains enabled by default. Serverless deployments should put + the Inductor/Triton caches on persistent storage so later workers reuse the + first worker's autotuning artifacts. + """ + value = os.getenv("JOYOMNI_VAE_COMPILE", "1").strip().lower() + return value not in {"0", "false", "no", "off"} + + +def strict_enabled() -> bool: + """Fail startup when an explicitly required compiled path cannot warm.""" + value = os.getenv("JOYOMNI_VAE_COMPILE_STRICT", "0").strip().lower() + return value in {"1", "true", "yes", "on"} + + # Must run before any compiled function executes, so warm restarts reuse the # on-disk autotune results instead of re-running coordinate descent. -_install_autotune_fix() +if compile_enabled(): + _install_autotune_fix() _configured: set[int] = set() _configured_encode: set[int] = set() _configured_encode_dynamic: set[int] = set() +_skip_notices: set[str] = set() +_compile_failures: list[str] = [] +_call_lock = threading.RLock() + + +@contextmanager +def call_guard(): + """Serialize entry into Dynamo/Inductor from streaming worker threads. + + PyTorch's compiler tracing state is not reliably isolated across Python + threads. The streaming pipeline deliberately runs encode, decode, and + pseudo-encode in parallel threads, so unguarded first calls can race while + a cached graph is being restored or specialized. The lock only covers + Python-side dispatch/kernel submission; CUDA streams may continue running + asynchronously after the guarded call returns. + """ + if not compile_enabled(): + yield + return + with _call_lock: + yield + + +def _warmup_failed(stage: str, exc: BaseException) -> None: + message = f"{stage} failed: {exc!r}" + _compile_failures.append(message) + print(f"[vae_compile] {message}", flush=True) + if strict_enabled(): + raise RuntimeError(message) from exc + + +def runtime_status() -> dict[str, object]: + """Return a JSON-safe snapshot used by health checks and launch logs.""" + ready = ( + compile_enabled() + and len(_configured_encode) >= 2 + and len(_configured) >= 1 + and not _compile_failures + ) + return { + "enabled": compile_enabled(), + "strict": strict_enabled(), + "ready": ready, + "encode_instances": len(_configured_encode), + "decode_instances": len(_configured), + "dynamic_encode_instances": len(_configured_encode_dynamic), + "thread_call_guard": compile_enabled(), + "failures": list(_compile_failures), + "cache": { + "torchinductor": os.getenv("TORCHINDUCTOR_CACHE_DIR"), + "triton": os.getenv("TRITON_CACHE_DIR"), + "cuda": os.getenv("CUDA_CACHE_PATH"), + }, + } + + +def assert_runtime_ready() -> None: + """Reject the slow eager path when the deployment requires compilation.""" + status = runtime_status() + if status["ready"]: + return + raise RuntimeError(f"compiled VAE did not become ready: {status}") + + +def _skip_compile(stage: str) -> bool: + if compile_enabled(): + return False + if stage not in _skip_notices: + print(f"[vae_compile] disabled by JOYOMNI_VAE_COMPILE=0; skipping {stage}") + _skip_notices.add(stage) + return True def maybe_setup_decode(vae) -> None: - if id(vae) in _configured: + if _skip_compile("decode compilation"): return - n_conv = 0 - for m in vae.modules(): - if isinstance(m, nn.Conv3d): - m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) - n_conv += 1 - if hasattr(vae, "_decode"): - vae._decode = torch.compile(vae._decode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "_decode" - elif hasattr(vae, "decode"): - vae.decode = torch.compile(vae.decode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "decode" - else: - raise RuntimeError("VAE has neither _decode nor decode; cannot compile") - _configured.add(id(vae)) + with call_guard(): + if id(vae) in _configured: + return + n_conv = 0 + for m in vae.modules(): + if isinstance(m, nn.Conv3d): + m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) + n_conv += 1 + if hasattr(vae, "_decode"): + vae._decode = torch.compile(vae._decode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "_decode" + elif hasattr(vae, "decode"): + vae.decode = torch.compile(vae.decode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "decode" + else: + raise RuntimeError("VAE has neither _decode nor decode; cannot compile") + _configured.add(id(vae)) print(f"[vae_compile] converted {n_conv} Conv3d weights to channels_last_3d + compiled vae.{target}") def prep_input(z: torch.Tensor) -> torch.Tensor: + if not compile_enabled(): + return z return z.to(memory_format=torch.channels_last_3d) def maybe_setup_encode(vae) -> None: - if id(vae) in _configured_encode: + if _skip_compile("encode compilation"): return - n_conv = 0 - for m in vae.modules(): - if isinstance(m, nn.Conv3d): - m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) - n_conv += 1 - if hasattr(vae, "_encode"): - vae._encode = torch.compile(vae._encode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "_encode" - elif hasattr(vae, "encode"): - vae.encode = torch.compile(vae.encode, mode="max-autotune-no-cudagraphs", dynamic=False) - target = "encode" - else: - raise RuntimeError("VAE has neither _encode nor encode; cannot compile") - _configured_encode.add(id(vae)) + with call_guard(): + if id(vae) in _configured_encode: + return + n_conv = 0 + for m in vae.modules(): + if isinstance(m, nn.Conv3d): + m.weight.data = m.weight.data.to(memory_format=torch.channels_last_3d) + n_conv += 1 + if hasattr(vae, "_encode"): + vae._encode = torch.compile(vae._encode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "_encode" + elif hasattr(vae, "encode"): + vae.encode = torch.compile(vae.encode, mode="max-autotune-no-cudagraphs", dynamic=False) + target = "encode" + else: + raise RuntimeError("VAE has neither _encode nor encode; cannot compile") + _configured_encode.add(id(vae)) print(f"[vae_compile] converted {n_conv} Conv3d weights to channels_last_3d + compiled vae.{target} (encode)") @@ -63,6 +164,8 @@ def warmup_encode(vae, in_channels: int, h_px: int, w_px: int, device: torch.device, dtype: torch.dtype, temporal_lens: tuple[int, ...] = (1, 9), autocast: bool = False) -> None: + if _skip_compile("encode warmup"): + return maybe_setup_encode(vae) from contextlib import nullcontext dev_type = torch.device(device).type @@ -75,26 +178,29 @@ def warmup_encode(vae, in_channels: int, h_px: int, w_px: int, if use_ac else nullcontext() ) try: - with torch.no_grad(), ctx: + with torch.no_grad(), ctx, call_guard(): _ = vae.encode(x) if torch.cuda.is_available(): torch.cuda.synchronize(device) print(f"[vae_compile] warmup compiled encode shape (1,{in_channels},{t},{h_px},{w_px}) autocast={autocast}") except Exception as exc: # noqa: BLE001 - print(f"[vae_compile] encode warmup failed for t={t}: {exc!r}") + _warmup_failed(f"encode warmup t={t}", exc) def maybe_setup_encode_dynamic(vae) -> None: - if id(vae) in _configured_encode_dynamic: + if _skip_compile("dynamic encode compilation"): return - if hasattr(vae, "_encode"): - core = getattr(vae, "_encode") - elif hasattr(vae, "encode"): - core = getattr(vae, "encode") - else: - raise RuntimeError("VAE has neither _encode nor encode; cannot compile") - vae._encode_dynamic = torch.compile(core, mode="max-autotune-no-cudagraphs", dynamic=True) - _configured_encode_dynamic.add(id(vae)) + with call_guard(): + if id(vae) in _configured_encode_dynamic: + return + if hasattr(vae, "_encode"): + core = getattr(vae, "_encode") + elif hasattr(vae, "encode"): + core = getattr(vae, "encode") + else: + raise RuntimeError("VAE has neither _encode nor encode; cannot compile") + vae._encode_dynamic = torch.compile(core, mode="max-autotune-no-cudagraphs", dynamic=True) + _configured_encode_dynamic.add(id(vae)) print("[vae_compile] compiled vae._encode_dynamic (dynamic=True, reference-image path)") @@ -113,6 +219,8 @@ def encode_via_dynamic(vae, x: torch.Tensor): def warmup_encode_dynamic(vae, in_channels: int, hw_list, device: torch.device, dtype: torch.dtype, temporal_lens: tuple[int, ...] = (1,), autocast: bool = False) -> None: + if _skip_compile("dynamic encode warmup"): + return fn = getattr(vae, "_encode_dynamic", None) if fn is None: print("[vae_compile] warmup_encode_dynamic skipped: _encode_dynamic not set up") @@ -130,13 +238,13 @@ def warmup_encode_dynamic(vae, in_channels: int, hw_list, device: torch.device, if use_ac else nullcontext() ) try: - with torch.no_grad(), ctx: + with torch.no_grad(), ctx, call_guard(): _ = fn(x) if torch.cuda.is_available(): torch.cuda.synchronize(device) n_ok += 1 except Exception as exc: # noqa: BLE001 - print(f"[vae_compile] dynamic encode warmup failed for ({h_px},{w_px},t={t}): {exc!r}") + _warmup_failed(f"dynamic encode warmup ({h_px},{w_px},t={t})", exc) print(f"[vae_compile] dynamic encode warmup done: {n_ok}/{len(hw_list) * len(temporal_lens)} shapes autocast={autocast}") @@ -144,6 +252,8 @@ def warmup_decode(vae, latent_channels: int, h_lat: int, w_lat: int, device: torch.device, dtype: torch.dtype, temporal_lens: tuple[int, ...] = (1, 2), autocast: bool = True) -> None: + if _skip_compile("decode warmup"): + return maybe_setup_decode(vae) from contextlib import nullcontext dev_type = torch.device(device).type @@ -156,10 +266,10 @@ def warmup_decode(vae, latent_channels: int, h_lat: int, w_lat: int, if use_ac else nullcontext() ) try: - with torch.no_grad(), ctx: + with torch.no_grad(), ctx, call_guard(): _ = vae.decode(z, return_dict=False)[0] if torch.cuda.is_available(): torch.cuda.synchronize(device) print(f"[vae_compile] warmup compiled decode shape (1,{latent_channels},{t},{h_lat},{w_lat}) autocast={autocast}") except Exception as exc: # noqa: BLE001 - print(f"[vae_compile] warmup failed for t={t}: {exc!r}") + _warmup_failed(f"decode warmup t={t}", exc) diff --git a/deploy/xvideo/serving/joyomni_streaming.py b/deploy/xvideo/serving/joyomni_streaming.py index d514521..1870a3b 100644 --- a/deploy/xvideo/serving/joyomni_streaming.py +++ b/deploy/xvideo/serving/joyomni_streaming.py @@ -38,9 +38,17 @@ def _autocast_ctx(device_type: str, dtype: torch.dtype, enabled: bool): _FULL_WARMUP_CHUNKS = 4 +_FULL_WARMUP_TIMEOUT_SECONDS = 300.0 _GRAPH_CACHE_CAP = 1 _GRAPH_CAPTURE_MAX_FAILS = 2 + +def _env_on(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + def _vae_compile_module(): from xvideo.models.vae import vae_compile as module return module @@ -179,6 +187,20 @@ def __init__( self.graph_runners: dict[tuple, StreamingGraphRunner] = {} self.graph_capture_failures: dict[tuple, int] = {} + def optimization_status(self) -> dict[str, Any]: + runners = list(self.graph_runners.values()) + return { + "vae_compile": _vae_compile_module().runtime_status(), + "cuda_graph": { + "enabled": graph_env_enabled(), + "ready": any(getattr(runner, "ready", False) for runner in runners), + "ready_runners": sum( + bool(getattr(runner, "ready", False)) for runner in runners + ), + "capture_failures": sum(self.graph_capture_failures.values()), + }, + } + @classmethod def load( cls, @@ -251,7 +273,10 @@ def load( pipeline.transformer.eval() _orientations = [(warmup_height, warmup_width)] - if (warmup_width, warmup_height) != (warmup_height, warmup_width): + if ( + _env_on("JOYOMNI_WARMUP_BOTH_ORIENTATIONS", True) + and (warmup_width, warmup_height) != (warmup_height, warmup_width) + ): _orientations.append((warmup_width, warmup_height)) _stem_mod = pipeline.vae.stem @@ -272,6 +297,8 @@ def load( ) except Exception as _vc_exc: print(f"#####[STREAM] VAE compile warmup skipped: {_vc_exc!r}") + if _env_on("JOYOMNI_VAE_COMPILE_STRICT"): + raise RuntimeError("required VAE decode compilation failed") from _vc_exc try: _vc = _vae_compile_module() @@ -293,20 +320,28 @@ def load( autocast=_vae_ac, ) - _ref_basesize = getattr(cfg, "ref_image_basesize", DEFAULT_REFERENCE_IMG_IV2V_BASESIZE) - _ref_cfgs = generate_video_image_bucket( - img_basesize=_ref_basesize, bs_img=1, bs_vid=0, bs_mimg=0, bs_mvid=0, - ) - _ref_hw = sorted({(c[3], c[4]) for c in _ref_cfgs}) - _vc.maybe_setup_encode_dynamic(_src_vae) - _vc.warmup_encode_dynamic( - _src_vae, 3, _ref_hw, - device=_module_device(_src_vae), dtype=_vae_dt, - temporal_lens=(1,), - autocast=False, - ) + if _env_on("JOYOMNI_WARMUP_REFERENCE_BUCKETS", True): + _ref_basesize = getattr(cfg, "ref_image_basesize", DEFAULT_REFERENCE_IMG_IV2V_BASESIZE) + _ref_cfgs = generate_video_image_bucket( + img_basesize=_ref_basesize, bs_img=1, bs_vid=0, bs_mimg=0, bs_mvid=0, + ) + _ref_hw = sorted({(c[3], c[4]) for c in _ref_cfgs}) + _vc.maybe_setup_encode_dynamic(_src_vae) + _vc.warmup_encode_dynamic( + _src_vae, 3, _ref_hw, + device=_module_device(_src_vae), dtype=_vae_dt, + temporal_lens=(1,), + autocast=False, + ) + else: + print("#####[STREAM] reference-bucket VAE warmup disabled; reference shapes compile lazily") except Exception as _ve_exc: print(f"#####[STREAM] VAE encode compile warmup skipped: {_ve_exc!r}") + if _env_on("JOYOMNI_VAE_COMPILE_STRICT"): + raise RuntimeError("required VAE encode compilation failed") from _ve_exc + + if _env_on("JOYOMNI_VAE_COMPILE_STRICT"): + _vae_compile_module().assert_runtime_ready() runtime = cls( cfg=cfg, @@ -325,7 +360,16 @@ def load( for (_wh, _ww) in _orientations: runtime.warmup_full_pipeline(height=_wh, width=_ww) except Exception as _wexc: - print(f"#####[STREAM] full-pipeline warmup error (non-fatal): {_wexc!r}") + _warmup_strict = _env_on("JOYOMNI_LOAD_WARMUP_STRICT") + _severity = "fatal" if _warmup_strict else "non-fatal" + print(f"#####[STREAM] full-pipeline warmup error ({_severity}): {_wexc!r}") + if _warmup_strict: + raise RuntimeError("required full-pipeline warmup failed") from _wexc + + if _env_on("JOYOMNI_LOAD_WARMUP_STRICT") and graph_env_enabled(): + ready_graphs = [runner for runner in runtime.graph_runners.values() if getattr(runner, "ready", False)] + if not ready_graphs: + raise RuntimeError("required CUDA graph was not captured during full-pipeline warmup") return runtime @@ -390,14 +434,33 @@ def warmup_full_pipeline( if completed >= num_chunks: break - deadline = time.time() + 120.0 - while completed < num_chunks and time.time() < deadline: + timeout_seconds = max( + 30.0, + float( + os.environ.get( + "JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS", + str(_FULL_WARMUP_TIMEOUT_SECONDS), + ) + ), + ) + print( + f"#####[STREAM] waiting up to {timeout_seconds:.0f}s for " + f"{num_chunks - completed} warmup chunk(s)" + ) + deadline = time.monotonic() + timeout_seconds + while completed < num_chunks and time.monotonic() < deadline: r = session.wait_async_result(timeout=0.5) if r is not None: completed += 1 + if completed < num_chunks and _env_on("JOYOMNI_LOAD_WARMUP_STRICT"): + raise RuntimeError( + f"full-pipeline warmup returned only {completed}/{num_chunks} chunks" + ) print(f"#####[STREAM] full-pipeline warmup done: {completed} chunks in {time.time() - t0:.1f}s") except Exception as exc: print(f"#####[STREAM] full-pipeline warmup skipped/failed: {exc!r}") + if _env_on("JOYOMNI_LOAD_WARMUP_STRICT"): + raise finally: if session is not None: try: @@ -672,7 +735,8 @@ def _encode_ref_image_latent(self) -> torch.Tensor | None: ref_img_encoded = ref_img_tensor.to(device=encode_device, dtype=self.vae_dtype) _vc = _vae_compile_module() - encoded = _vc.encode_via_dynamic(self.pipeline.vae, ref_img_encoded) + with _vc.call_guard(): + encoded = _vc.encode_via_dynamic(self.pipeline.vae, ref_img_encoded) if not hasattr(encoded, "latent_dist"): raise TypeError(f"Unsupported VAE encode output type for ref image: {type(encoded)}") ref_img_latent = encoded.latent_dist.sample() @@ -1670,7 +1734,8 @@ def _decode_chunk_pixels( vae_ctx = _autocast_ctx(vae_device_type, self.vae_dtype, self.vae_autocast_enabled) with vae_ctx: started = self._timer_start(vae_device) - chunk_decoded = decode_vae.decode(decode_input, return_dict=False)[0] + with _vc.call_guard(): + chunk_decoded = decode_vae.decode(decode_input, return_dict=False)[0] if profile is not None: self._timer_record(profile, "vae_decode_s", started) @@ -1697,7 +1762,8 @@ def _encode_next_decode_pseudo_latent( prev_pixels = _vc.prep_input(prev_pixels) vae_ctx = _autocast_ctx(pseudo_device_type, self.vae_dtype, self.vae_autocast_enabled) with vae_ctx: - pseudo_enc = pseudo_vae.encode(prev_pixels) + with _vc.call_guard(): + pseudo_enc = pseudo_vae.encode(prev_pixels) if hasattr(pseudo_enc, "latent_dist"): pseudo_latent = pseudo_enc.latent_dist.sample() else: diff --git a/deploy/xvideo/serving/serve_joyomni_streaming.py b/deploy/xvideo/serving/serve_joyomni_streaming.py index 976a547..3de2325 100644 --- a/deploy/xvideo/serving/serve_joyomni_streaming.py +++ b/deploy/xvideo/serving/serve_joyomni_streaming.py @@ -26,6 +26,7 @@ import uvicorn from xvideo.serving.pe import DEFAULT_MODEL as DEFAULT_PE_MODEL +from xvideo.checkpoint_status import checkpoint_status from xvideo.serving.joyomni_streaming import ( JoyOmniRuntime, StreamingSettings, @@ -73,6 +74,10 @@ def release(self, ticket: int) -> None: WS_SEND_TIMEOUT_S = 10.0 +HOLDER_IDLE_TIMEOUT_S = max( + 10.0, + float(os.environ.get("JOYOMNI_SESSION_IDLE_TIMEOUT_SECONDS", "60")), +) REF_IMAGE_DIR = REPO_ROOT / "rv2v_reference" REF_IMAGE_FILES = { @@ -603,6 +608,25 @@ def get_runtime() -> JoyOmniRuntime: warmup_height=args.height, warmup_width=args.width, ) + optimization_status = app.state.runtime.optimization_status() + print( + f"#####[OPTIMIZATIONS] {json.dumps(optimization_status, sort_keys=True)}", + flush=True, + ) + marker_path = os.getenv("JOYOMNI_CACHE_READY_MARKER") + if marker_path: + marker = Path(marker_path) + marker.parent.mkdir(parents=True, exist_ok=True) + marker.write_text( + json.dumps( + { + "ready_at": time.time(), + "optimizations": optimization_status, + }, + indent=2, + sort_keys=True, + ) + ) return app.state.runtime @asynccontextmanager @@ -655,6 +679,7 @@ def health() -> JSONResponse: "ok": True, "runtime_loaded": app.state.runtime is not None, "dit_ckpt": args.dit_ckpt, + "checkpoint": checkpoint_status(args.dit_ckpt), "device": str(app.state.runtime.device) if app.state.runtime is not None else args.device, "vae_device": str(_module_device(app.state.runtime.pipeline.vae)) if app.state.runtime is not None else args.vae_device, "vae_encode_device": ( @@ -683,6 +708,11 @@ def health() -> JSONResponse: "max_temporal_ids": args.max_temporal_ids, "freeze_kv_on_static": args.freeze_kv_on_static, "static_diff_thresh": args.static_diff_thresh, + "optimizations": ( + app.state.runtime.optimization_status() + if app.state.runtime is not None + else None + ), } ) @@ -1295,7 +1325,6 @@ async def _reset_session(reason: str) -> None: app.state.ws_debug = ws_debug - HOLDER_IDLE_TIMEOUT_S = 10.0 last_activity = time.monotonic() last_frames_out = frames_out while True: @@ -1303,6 +1332,11 @@ async def _reset_session(reason: str) -> None: last_frames_out = frames_out last_activity = time.monotonic() if time.monotonic() - last_activity >= HOLDER_IDLE_TIMEOUT_S: + print( + f"#####[WS-GUARD] releasing session after " + f"{HOLDER_IDLE_TIMEOUT_S:.0f}s without browser or output activity", + flush=True, + ) try: await _send_json( { @@ -1320,6 +1354,12 @@ async def _reset_session(reason: str) -> None: if "text" in message and message["text"] is not None: payload = json.loads(message["text"]) msg_type = payload.get("type") + # ACK and ping messages prove that the live browser and its + # camera loop are still reachable. The old 10-second guard + # ignored them, so a short frame/backpressure pause could + # release a healthy session and leave the last frame frozen. + last_activity = time.monotonic() + ws_debug["last_client_activity_at"] = time.time() if msg_type == "start": print(f"#####[RESTART] 'start' received (session {'live' if session is not None else 'none'})", flush=True) last_activity = time.monotonic() diff --git a/runpod/README.md b/runpod/README.md new file mode 100644 index 0000000..2040d50 --- /dev/null +++ b/runpod/README.md @@ -0,0 +1,113 @@ +# RunPod connection contract + +JoyAI uses two HTTP servers inside every RunPod load-balancer worker. The +public application server listens on `PORT` (normally `8080`). A small internal +readiness server listens on `PORT_HEALTH` (normally `8081`). + +| Purpose | Correct route | Caller | +| --- | --- | --- | +| Web interface | `GET /` | Browser or client | +| JoyAI server health | `GET /health` | Browser, client, or diagnostics | +| Load/warm the model | `POST /load` | Startup or authenticated client | +| Real-time video stream | `WS /ws` | Web interface | +| RunPod worker readiness | `GET /ping` on `PORT_HEALTH` | RunPod only | + +Do **not** call `/ping` through the public `*.api.runpod.ai` address. That +address routes to the JoyAI server on `PORT`, where the health route is +`/health`. RunPod calls the separate `/ping` route internally on +`PORT_HEALTH`. + +## Startup sequence + +1. `runpod/start.py` starts the internal readiness server on `PORT_HEALTH`. +2. It starts JoyAI with `--preload`, so the model is loaded immediately from + `JOYOMNI_CKPT_ROOT`. +3. While the model loads, internal `GET /ping` returns `204` (initializing). +4. When public `GET /health` reports `runtime_loaded: true`, internal + `GET /ping` returns `200` and RunPod begins routing public traffic. + +This avoids the deadlock where the worker waited for a public `/load` request +while RunPod waited for the worker to become ready before routing that request. + +## RunPod endpoint settings + +Use these values for a load-balancer endpoint: + +| Setting | Value | +| --- | --- | +| `PORT` | `8080` | +| `PORT_HEALTH` | `8081` | +| `JOYOMNI_PRELOAD` | `1` | +| `JOYOMNI_CKPT_ROOT` | The checkpoint path on the mounted network volume | +| Exposed HTTP ports | `8080,8081` | +| Health-check path | `/ping` | +| Active/min workers | `0` while testing to avoid idle GPU charges | +| Max workers | `1` for one-viewer testing and cost safety | +| Idle timeout | `60` seconds | + +RunPod does not reliably treat a quiet WebSocket as autoscaling activity. The +local proxy therefore sends an authenticated `GET /health` every three seconds +while the stream is open. Ordinary HTTP follow-up requests use soft +`X-Runpod-Worker-Id` affinity, while the WebSocket upgrade uses normal routing: +strict affinity can be held outside the worker when RunPod considers a +long-lived connection at capacity. Limiting the endpoint to one worker is the +hard guarantee that normal routing cannot start a second H200. + +The H200 image default checkpoint path is +`/runpod-volume/joyai/checkpoints`. Mount the model network volume so that the +downloaded checkpoint tree is available at that path. + +## Windows real-time test + +The browser cannot add a RunPod bearer token to normal page navigation or the +WebSocket constructor. Use the included local authenticated proxy: + +```powershell +cd path\to\JoyAI-Video-Edit\runpod +powershell.exe -NoProfile -ExecutionPolicy Bypass -File .\Start-JoyAI-Realtime-Test.ps1 +``` + +The script securely asks for the API key and polls `GET /health` until the +preloaded runtime is ready. RunPod's load balancer may return HTTP `400` with +`timed out waiting for worker` after its own two-minute wait while the worker +is still initializing; the script treats that response as retryable up to its +30-minute safety limit. It then starts the local HTTP/WebSocket proxy and +opens the interface. The key stays in the local process and is removed when +the script exits. + +`POST /load` is not used by this test flow because the container already starts +with `JOYOMNI_PRELOAD=1`, and RunPod cannot route that request until the worker +has passed readiness anyway. + +The 30-minute limit bounds how long the local test waits; it is not a hard +RunPod billing cutoff. RunPod bills from worker start until the worker fully +stops. If the limit expires, stop the current worker (or temporarily set max +workers to zero) before investigating the log. Do not start a second worker or +rebuild the image during that inspection. + +The H200 image defaults to live-only operation. Recording and download +finalization are disabled, the presence gate is off, and the browser starts at +20 FPS with low upload/downlink quality. Its low-latency playback buffer targets +100 ms, and the adaptive downlink can increase quality after the connection +proves stable. Refreshing the browser replaces the previous WebSocket session +instead of waiting behind its stale session ticket. + +The proxy records the `X-Runpod-Worker-Id` returned by the initial local health +check and uses it as a soft preference for page and asset requests. Keep +`Max workers` at `1` during one-viewer testing so the unblocked WebSocket route +still reaches that same worker and cannot launch a different cold worker. + +The H200 image enables the upstream compiled/autotuned VAE path. TorchInductor, +Triton, and CUDA caches are stored under +`/runpod-volume/joyai/cache/h200-torch291-cu128`, so the expensive first compile +is reused by later workers with the same pinned image stack. The health payload +reports `optimizations.vae_compile` and `optimizations.cuda_graph`; the strict +H200 defaults fail startup rather than silently serving the slow eager path. +For the live-only profile, startup precompiles the 840×480 landscape stream and +does not spend GPU time warming unused portrait or reference-image shapes. +After compilation, the server allows up to 300 seconds for its four asynchronous +full-pipeline warm-up chunks. This keeps strict readiness without rejecting a +healthy first compile when later chunks finish shortly after two minutes. + +Stop the test with `Ctrl+C`. With zero active workers and a short idle timeout, +RunPod can scale the worker back to zero after the connection closes. diff --git a/runpod/Start-JoyAI-Realtime-Test.ps1 b/runpod/Start-JoyAI-Realtime-Test.ps1 new file mode 100644 index 0000000..c052b1d --- /dev/null +++ b/runpod/Start-JoyAI-Realtime-Test.ps1 @@ -0,0 +1,195 @@ +param( + [string]$EndpointId = "ex9647vtulowka", + [int]$LocalPort = 9000, + [int]$WarmTimeoutSeconds = 1800, + [int]$RetryDelaySeconds = 10 +) + +$ErrorActionPreference = "Stop" +$proxyProcess = $null +$secret = $null +$apiKey = $null + +function Get-HttpStatusCode { + param([System.Management.Automation.ErrorRecord]$ErrorRecord) + + try { + if ($null -ne $ErrorRecord.Exception.Response.StatusCode) { + return [int]$ErrorRecord.Exception.Response.StatusCode + } + } + catch { + return $null + } + return $null +} + +function Get-ErrorDetail { + param([System.Management.Automation.ErrorRecord]$ErrorRecord) + + if ($null -ne $ErrorRecord.ErrorDetails -and $ErrorRecord.ErrorDetails.Message) { + return $ErrorRecord.ErrorDetails.Message + } + return $ErrorRecord.Exception.Message +} + +try { + Write-Host "Checking the local proxy dependency..." + python.exe -c "import aiohttp" 2>$null + if ($LASTEXITCODE -ne 0) { + python.exe -m pip install --user aiohttp + if ($LASTEXITCODE -ne 0) { + throw "Could not install aiohttp for the local proxy." + } + } + + $secret = Read-Host "Paste your RunPod API key" -AsSecureString + $apiKey = [System.Net.NetworkCredential]::new("", $secret).Password + $headers = @{ Authorization = "Bearer $apiKey" } + $baseUrl = "https://$EndpointId.api.runpod.ai" + + Write-Host "Starting the JoyAI worker through GET /health..." + Write-Host "The H200 must load the 32.5 GB DiT checkpoint before RunPod routes traffic." + Write-Host "The first compiled-VAE start can take up to 30 minutes while it creates the persistent optimization cache." + Write-Host "A RunPod 400 'timed out waiting for worker' response is retried until the $WarmTimeoutSeconds-second safety limit." + + $warmTimer = [System.Diagnostics.Stopwatch]::StartNew() + $health = $null + $workerReady = $false + $attempt = 0 + + while (-not $workerReady -and $warmTimer.Elapsed.TotalSeconds -lt $WarmTimeoutSeconds) { + $attempt += 1 + $remainingSeconds = $WarmTimeoutSeconds - [int]$warmTimer.Elapsed.TotalSeconds + if ($remainingSeconds -le 0) { + break + } + + # RunPod's load balancer can spend about two minutes waiting for a cold + # worker before returning HTTP 400. Keep each request slightly above + # that window, then retry without creating another worker. + $requestTimeout = [Math]::Min(135, $remainingSeconds) + Write-Host "Readiness attempt $attempt (elapsed $([int]$warmTimer.Elapsed.TotalSeconds)s)..." + + try { + $health = Invoke-RestMethod ` + -Uri "$baseUrl/health" ` + -Headers $headers ` + -Method Get ` + -TimeoutSec $requestTimeout + + if ($health.ok -and $health.runtime_loaded) { + $workerReady = $true + break + } + + Write-Host "Worker answered, but the runtime is still initializing." + } + catch { + $statusCode = Get-HttpStatusCode $_ + $detail = Get-ErrorDetail $_ + $retryableStatusCodes = @(400, 408, 425, 429, 500, 502, 503, 504) + $isColdStartResponse = ( + $retryableStatusCodes -contains $statusCode -or + $detail -match "timed out waiting for worker|operation has timed out|cold start|no worker|temporarily unavailable|semaphore timeout|forcibly closed" + ) + + if (-not $isColdStartResponse) { + throw + } + + Write-Host "Worker is still cold-starting: $detail" + } + + $remainingSeconds = $WarmTimeoutSeconds - [int]$warmTimer.Elapsed.TotalSeconds + if ($remainingSeconds -gt 0) { + $sleepSeconds = [Math]::Min($RetryDelaySeconds, $remainingSeconds) + Write-Host "Retrying in $sleepSeconds seconds..." + Start-Sleep -Seconds $sleepSeconds + } + } + + $warmTimer.Stop() + if (-not $workerReady) { + throw "JoyAI did not become ready within $WarmTimeoutSeconds seconds. Stop the current worker in RunPod before inspecting its log; ending this client request alone does not guarantee that GPU billing has stopped. Do not rebuild the image or start another Pod." + } + + Write-Host "JoyAI runtime is ready after $([int]$warmTimer.Elapsed.TotalSeconds) seconds." + if ( + $null -eq $health.checkpoint -or + $health.checkpoint.status -ne "current" + ) { + $checkpointStatus = if ($null -eq $health.checkpoint) { + "missing from the health response" + } else { + [string]$health.checkpoint.status + } + throw "The upgraded RV2V checkpoint was not verified ($checkpointStatus). Stop this worker. Verify the mounted volume with: python3 /opt/joyai/runpod/verify_checkpoint.py --full-hash" + } + Write-Host "Verified: the upgraded RV2V checkpoint is active." + if ( + $null -eq $health.optimizations -or + -not $health.optimizations.vae_compile.ready -or + -not $health.optimizations.cuda_graph.ready + ) { + throw "JoyAI became reachable without the required compiled VAE and CUDA graph. Stop the worker and inspect its optimization log before testing." + } + Write-Host "Verified: compiled VAE and CUDA graph are active." + + $env:RUNPOD_API_KEY = $apiKey + $proxyScript = Join-Path $PSScriptRoot "local_proxy.py" + $proxyProcess = Start-Process ` + -FilePath "python.exe" ` + -ArgumentList @( + ('"{0}"' -f $proxyScript), + "--endpoint-id", $EndpointId, + "--port", $LocalPort + ) ` + -NoNewWindow ` + -PassThru + + $localUrl = "http://127.0.0.1:$LocalPort/" + $localHealthUrl = "${localUrl}health" + $localReady = $false + for ($attempt = 1; $attempt -le 10; $attempt++) { + if ($proxyProcess.HasExited) { + throw "The local JoyAI proxy stopped unexpectedly." + } + try { + $localHealth = Invoke-RestMethod ` + -Uri $localHealthUrl ` + -Method Get ` + -TimeoutSec 30 + if ($localHealth.ok -and $localHealth.runtime_loaded) { + $localReady = $true + break + } + } + catch { + $localDetail = Get-ErrorDetail $_ + Write-Host "Local proxy check $attempt failed: $localDetail" + Start-Sleep -Seconds 1 + } + } + + if (-not $localReady) { + throw "The local proxy could not reach the ready JoyAI worker." + } + + Write-Host "JoyAI is ready: $localUrl" + Write-Host "Keep this window open. Press Ctrl+C when the test is finished." + Start-Process $localUrl + Wait-Process -Id $proxyProcess.Id +} +catch { + Write-Host "TEST FAILED: $($_.Exception.Message)" -ForegroundColor Red + exit 1 +} +finally { + if ($null -ne $proxyProcess -and -not $proxyProcess.HasExited) { + Stop-Process -Id $proxyProcess.Id -Force -ErrorAction SilentlyContinue + } + Remove-Item Env:RUNPOD_API_KEY -ErrorAction SilentlyContinue + $apiKey = $null + $secret = $null +} diff --git a/runpod/download_models.py b/runpod/download_models.py new file mode 100644 index 0000000..ff2fe7a --- /dev/null +++ b/runpod/download_models.py @@ -0,0 +1,62 @@ +import os +import urllib.request +from pathlib import Path + +from huggingface_hub import snapshot_download + + +JOYAI_RV2V_REVISION = "eda14f342ef99c52485bbb8dc271c29b42298089" + + +checkpoint_root = Path( + os.getenv( + "JOYOMNI_CKPT_ROOT", + "/runpod-volume/joyai/checkpoints", + ) +) + +checkpoint_root.mkdir(parents=True, exist_ok=True) + +hf_token = os.getenv("HF_TOKEN") or None + +print("Downloading JoyAI 0811 model and VAE...") + +snapshot_download( + repo_id="jdopensource/JoyAI-Video-Edit", + revision=JOYAI_RV2V_REVISION, + local_dir=checkpoint_root / "JoyAI-Video-Edit", + allow_patterns=[ + "dit/joyai_video_edit_dit_0811.pth", + "vae/*", + ], + token=hf_token, + max_workers=8, +) + +print("Downloading MiMo-VL text and vision encoder...") + +snapshot_download( + repo_id="XiaomiMiMo/MiMo-VL-7B-RL-2508", + local_dir=checkpoint_root / "MiMo-VL-7B-RL-2508", + token=hf_token, + max_workers=8, +) + +face_model = checkpoint_root / "face_detection_yunet_2023mar.onnx" + +if not face_model.exists(): + print("Downloading YuNet face detector...") + + temporary_file = face_model.with_suffix(".download") + + urllib.request.urlretrieve( + "https://media.githubusercontent.com/media/opencv/" + "opencv_zoo/main/models/face_detection_yunet/" + "face_detection_yunet_2023mar.onnx", + temporary_file, + ) + + temporary_file.replace(face_model) + +print("All required model files have been downloaded.") +print(f"Checkpoint location: {checkpoint_root}") diff --git a/runpod/health_server.py b/runpod/health_server.py new file mode 100644 index 0000000..2226194 --- /dev/null +++ b/runpod/health_server.py @@ -0,0 +1,58 @@ +import json +import os +import urllib.error +import urllib.request + +import uvicorn +from fastapi import FastAPI, Response + + +app = FastAPI() + + +@app.get("/ping") +def ping(): + """RunPod-only readiness probe served on PORT_HEALTH. + + The public JoyAI server owns /health. A worker remains in RunPod's + initializing state (204) until that endpoint confirms the runtime loaded. + """ + model_port = os.getenv( + "JOYOMNI_PORT", + os.getenv("PORT", "8080"), + ) + + try: + with urllib.request.urlopen( + f"http://127.0.0.1:{model_port}/health", + timeout=2, + ) as response: + health = json.loads(response.read()) + + if health.get("ok") is True and health.get("runtime_loaded") is True: + return { + "status": "healthy", + "model": "ready", + } + + except ( + urllib.error.URLError, + TimeoutError, + ValueError, + ): + pass + + # RunPod load-balancer health contract: + # 200 = ready, 204 = still initializing, anything else = unhealthy. + return Response(status_code=204) + + +if __name__ == "__main__": + health_port = int(os.getenv("PORT_HEALTH", "8081")) + + uvicorn.run( + app, + host="0.0.0.0", + port=health_port, + log_level="info", + ) diff --git a/runpod/local_proxy.py b/runpod/local_proxy.py new file mode 100644 index 0000000..2e71b45 --- /dev/null +++ b/runpod/local_proxy.py @@ -0,0 +1,354 @@ +"""Local authenticated proxy for the JoyAI RunPod load-balancer endpoint. + +Browsers cannot attach a secret RunPod bearer token to normal navigation or a +WebSocket constructor. This proxy keeps the token in the local process and +adds it to upstream HTTP and WebSocket requests. +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import os + +from aiohttp import ( + ClientError, + ClientSession, + ClientTimeout, + WSCloseCode, + WSMsgType, + web, +) + + +HOP_BY_HOP_HEADERS = { + "connection", + "content-length", + "host", + "keep-alive", + "proxy-authenticate", + "proxy-authorization", + "te", + "trailer", + "transfer-encoding", + "upgrade", +} + +RUNPOD_WORKER_HEADER = "X-Runpod-Worker-Id" +# RunPod's load balancer may not count an otherwise active WebSocket as worker +# activity. Keep this below the platform's 5-second default idle timeout so a +# live stream cannot be scaled down underneath the browser. +RUNPOD_KEEPALIVE_SECONDS = 3.0 + + +def proxy_exception_handler(loop: asyncio.AbstractEventLoop, context: dict) -> None: + """Ignore the harmless Windows Proactor reset logged after a peer closes.""" + error = context.get("exception") + message = context.get("message", "") + if ( + os.name == "nt" + and isinstance(error, ConnectionResetError) + and getattr(error, "winerror", None) == 10054 + and "_call_connection_lost" in message + ): + return + loop.default_exception_handler(context) + + +def filtered_headers(headers) -> dict[str, str]: + return { + name: value + for name, value in headers.items() + if name.lower() not in HOP_BY_HOP_HEADERS + and name.lower() != "authorization" + } + + +def upstream_url(request: web.Request) -> str: + return f"{request.app['upstream']}{request.rel_url}" + + +def upstream_headers( + app: web.Application, + *, + affinity: bool = True, +) -> dict[str, str]: + """Authenticate requests and optionally prefer the selected worker.""" + headers = dict(app["auth_headers"]) + worker_id = app["proxy_state"].get("worker_id") + if affinity and worker_id: + # Soft affinity never waits behind a worker RunPod considers at + # capacity. Max workers=1 remains the hard guarantee that a fallback + # cannot create a second H200 during one-viewer testing. + headers[RUNPOD_WORKER_HEADER] = worker_id + return headers + + +def remember_worker(app: web.Application, headers) -> str | None: + """Remember the worker chosen by RunPod for later HTTP and WS requests.""" + worker_id = str(headers.get(RUNPOD_WORKER_HEADER, "")).strip() + if not worker_id: + return None + state = app["proxy_state"] + if state.get("worker_id") != worker_id: + state["worker_id"] = worker_id + print(f"JoyAI pinned to RunPod worker {worker_id}.", flush=True) + return worker_id + + +async def keep_worker_active(app: web.Application) -> None: + """Reset RunPod's idle timer while the browser has a live WebSocket.""" + session: ClientSession = app["session"] + failures = 0 + while True: + await asyncio.sleep(RUNPOD_KEEPALIVE_SECONDS) + try: + async with session.get( + f"{app['upstream']}/health", + # RunPod can hold an affinity request while a long-lived WS is + # occupying the selected worker. With max workers=1, normal + # routing still reaches the only worker without that wait. + headers=upstream_headers(app, affinity=False), + timeout=ClientTimeout(total=10), + ) as response: + await response.read() + remember_worker(app, response.headers) + if response.status != 200: + raise ClientError( + f"RunPod keepalive returned HTTP {response.status}" + ) + failures = 0 + except asyncio.CancelledError: + raise + except (ClientError, asyncio.TimeoutError, OSError) as error: + failures += 1 + if failures == 1 or failures % 5 == 0: + print( + f"JoyAI RunPod keepalive failed ({failures}): {error}", + flush=True, + ) + + +async def proxy_websocket(request: web.Request) -> web.WebSocketResponse: + downstream = web.WebSocketResponse(heartbeat=30) + await downstream.prepare(request) + + session: ClientSession = request.app["session"] + proxy_state = request.app["proxy_state"] + upstream = None + keepalive_task = None + try: + # Do not put strict/soft affinity on the WebSocket upgrade. RunPod can + # hold an affinity upgrade outside the worker for several minutes even + # when ordinary HTTP requests reach it immediately. Max workers=1 pins + # normal routing to the only worker without blocking the handshake. + upstream = await session.ws_connect( + upstream_url(request), + headers=upstream_headers(request.app, affinity=False), + # The browser already sends an application-level JSON ping every + # second. A protocol heartbeat adds a second liveness mechanism + # that can falsely close with code 1006 when RunPod's load + # balancer does not relay a Pong during graph/session rebuilds. + heartbeat=None, + max_msg_size=0, + ) + response = getattr(upstream, "_response", None) + if response is not None: + remember_worker(request.app, response.headers) + print("JoyAI WebSocket connected through RunPod.", flush=True) + keepalive_task = asyncio.create_task(keep_worker_active(request.app)) + + # The local proxy is intentionally single-viewer. Replacing a stale + # browser socket explicitly tells the server to release its session + # ticket, avoiding a long queue after a page refresh. + async with proxy_state["websocket_lock"]: + previous_upstream = proxy_state["active_upstream"] + previous_downstream = proxy_state["active_downstream"] + proxy_state["active_upstream"] = upstream + proxy_state["active_downstream"] = downstream + + # Close the replaced pair outside the lock. The old handler also uses + # this lock during cleanup, so awaiting its close while holding the + # lock would deadlock a page refresh. + if previous_upstream is not None and not previous_upstream.closed: + try: + await previous_upstream.send_json({"type": "stop"}) + except (ClientError, asyncio.TimeoutError, OSError): + pass + await previous_upstream.close( + code=WSCloseCode.GOING_AWAY, + message=b"replaced by a refreshed browser", + ) + if previous_downstream is not None and not previous_downstream.closed: + await previous_downstream.close( + code=WSCloseCode.GOING_AWAY, + message=b"replaced by a refreshed browser", + ) + + async def browser_to_runpod() -> None: + async for message in downstream: + if message.type == WSMsgType.TEXT: + await upstream.send_str(message.data) + elif message.type == WSMsgType.BINARY: + await upstream.send_bytes(message.data) + elif message.type == WSMsgType.ERROR: + break + + async def runpod_to_browser() -> None: + async for message in upstream: + if message.type == WSMsgType.TEXT: + await downstream.send_str(message.data) + elif message.type == WSMsgType.BINARY: + await downstream.send_bytes(message.data) + elif message.type == WSMsgType.ERROR: + break + + tasks = [ + asyncio.create_task(browser_to_runpod()), + asyncio.create_task(runpod_to_browser()), + ] + done, pending = await asyncio.wait( + tasks, + return_when=asyncio.FIRST_COMPLETED, + ) + for task in pending: + task.cancel() + await asyncio.gather(*done, *pending, return_exceptions=True) + + except (ClientError, asyncio.TimeoutError, OSError) as error: + if not downstream.closed: + await downstream.send_json( + { + "type": "proxy_error", + "error": f"RunPod connection failed: {error}", + } + ) + finally: + if keepalive_task is not None: + keepalive_task.cancel() + await asyncio.gather(keepalive_task, return_exceptions=True) + upstream_code = getattr(upstream, "close_code", None) + downstream_code = getattr(downstream, "close_code", None) + print( + "JoyAI WebSocket closed " + f"(browser={downstream_code}, RunPod={upstream_code}).", + flush=True, + ) + if upstream is not None and not upstream.closed: + try: + await upstream.send_json({"type": "stop"}) + except (ClientError, asyncio.TimeoutError, OSError): + pass + await upstream.close() + async with proxy_state["websocket_lock"]: + if proxy_state["active_upstream"] is upstream: + proxy_state["active_upstream"] = None + if proxy_state["active_downstream"] is downstream: + proxy_state["active_downstream"] = None + if not downstream.closed: + await downstream.close() + + return downstream + + +async def proxy_http(request: web.Request) -> web.Response: + session: ClientSession = request.app["session"] + headers = filtered_headers(request.headers) + # Windows PowerShell 5.1 cannot reliably decode Brotli responses. Ask the + # upstream load balancer for an identity response so both PowerShell's + # readiness probe and the browser receive bytes they can consume directly. + headers["Accept-Encoding"] = "identity" + headers.update(upstream_headers(request.app)) + + try: + async with session.request( + request.method, + upstream_url(request), + headers=headers, + data=await request.read(), + allow_redirects=False, + ) as response: + body = await response.read() + remember_worker(request.app, response.headers) + return web.Response( + body=body, + status=response.status, + reason=response.reason, + headers=filtered_headers(response.headers), + ) + except (ClientError, asyncio.TimeoutError, OSError) as error: + return web.Response( + status=503, + content_type="application/json", + text=json.dumps( + { + "status": 503, + "title": "RunPod connection unavailable", + "detail": str(error), + } + ), + ) + + +async def route_request(request: web.Request): + if request.path == "/ws": + return await proxy_websocket(request) + return await proxy_http(request) + + +async def create_session(app: web.Application) -> None: + if os.name == "nt": + asyncio.get_running_loop().set_exception_handler(proxy_exception_handler) + app["session"] = ClientSession( + timeout=ClientTimeout(total=None, connect=60, sock_connect=60), + # Relay compressed bytes and their Content-Encoding header unchanged. + # The browser can decode Brotli itself; the local Python proxy should + # not require the optional Brotli package just to forward RunPod HTML. + auto_decompress=False, + ) + # aiohttp warns when top-level application state is mutated after startup. + # Keep live socket state inside one mutable object registered at startup. + app["proxy_state"] = { + "websocket_lock": asyncio.Lock(), + "active_upstream": None, + "active_downstream": None, + "worker_id": None, + } + + +async def close_session(app: web.Application) -> None: + await app["session"].close() + + +def create_app(endpoint_id: str, api_key: str) -> web.Application: + app = web.Application(client_max_size=1024**3) + app["upstream"] = f"https://{endpoint_id}.api.runpod.ai" + app["auth_headers"] = {"Authorization": f"Bearer {api_key}"} + app.on_startup.append(create_session) + app.on_cleanup.append(close_session) + app.router.add_route("*", "/{tail:.*}", route_request) + return app + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--endpoint-id", required=True) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", default=9000, type=int) + args = parser.parse_args() + + api_key = os.getenv("RUNPOD_API_KEY") + if not api_key: + raise SystemExit("RUNPOD_API_KEY is required") + + web.run_app( + create_app(args.endpoint_id, api_key), + host=args.host, + port=args.port, + print=lambda message: print(message, flush=True), + ) + + +if __name__ == "__main__": + main() diff --git a/runpod/start.py b/runpod/start.py new file mode 100644 index 0000000..51c1b0b --- /dev/null +++ b/runpod/start.py @@ -0,0 +1,179 @@ +import os +import subprocess +import sys +from pathlib import Path + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +DEPLOY_ROOT = REPOSITORY_ROOT / "deploy" +if str(DEPLOY_ROOT) not in sys.path: + sys.path.insert(0, str(DEPLOY_ROOT)) + +from xvideo.checkpoint_status import checkpoint_status + + +TRUE_VALUES = {"1", "true", "yes", "on"} +FALSE_VALUES = {"0", "false", "no", "off"} + + +def env_enabled(name: str, *, default: bool) -> bool: + """Read a strict boolean environment flag. + + A typo must fail during startup instead of silently disabling model preload + and leaving RunPod's load balancer waiting forever for a ready worker. + """ + value = os.getenv(name) + if value is None: + return default + + normalized = value.strip().lower() + if normalized in TRUE_VALUES: + return True + if normalized in FALSE_VALUES: + return False + + choices = ", ".join(sorted(TRUE_VALUES | FALSE_VALUES)) + raise ValueError(f"{name} must be one of: {choices}") + + +def required_checkpoint_items(checkpoint_root: Path) -> list[Path]: + return [ + checkpoint_root + / "JoyAI-Video-Edit" + / "dit" + / "joyai_video_edit_dit_0811.pth", + checkpoint_root + / "JoyAI-Video-Edit" + / "vae" + / "config.json", + checkpoint_root + / "JoyAI-Video-Edit" + / "vae" + / "diffusion_pytorch_model.safetensors", + checkpoint_root / "MiMo-VL-7B-RL-2508", + ] + + +def build_model_command(repository_root: Path, *, preload: bool) -> list[str]: + return [ + "bash", + str(repository_root / "deploy" / "run_server.sh"), + "--preload" if preload else "--no-preload", + ] + + +def main() -> int: + repository_root = REPOSITORY_ROOT + checkpoint_root = Path( + os.getenv( + "JOYOMNI_CKPT_ROOT", + "/runpod-volume/joyai/checkpoints", + ) + ) + + missing_items = [ + str(item) + for item in required_checkpoint_items(checkpoint_root) + if not item.exists() + ] + + if missing_items: + print("Required model files are missing:", flush=True) + for item in missing_items: + print(f" - {item}", flush=True) + + print( + "\nRun this command once on the attached RunPod volume:", + flush=True, + ) + print( + "python3 /opt/joyai/runpod/download_models.py", + flush=True, + ) + return 1 + + dit_checkpoint = required_checkpoint_items(checkpoint_root)[0] + dit_status = checkpoint_status(dit_checkpoint) + status = dit_status["status"] + print( + "JoyAI RV2V checkpoint status: " + f"{status} ({dit_status['verification']}).", + flush=True, + ) + if status == "stale": + print( + "The mounted volume contains a different 0811 DiT than the " + "upgraded RV2V release. Refusing to load stale weights.", + flush=True, + ) + print( + "Update the volume with: " + "python3 /opt/joyai/runpod/download_models.py", + flush=True, + ) + return 1 + if status == "unknown": + print( + "Hugging Face metadata is unavailable, so this checkpoint cannot " + "be identified without a one-time full hash. Run: " + "python3 /opt/joyai/runpod/verify_checkpoint.py --full-hash", + flush=True, + ) + + environment = os.environ.copy() + + # RunPod supplies PORT for public traffic. PORT_HEALTH is reserved for the + # internal /ping server and must not be used by clients. + model_port = environment.get( + "PORT", + environment.get("JOYOMNI_PORT", "8080"), + ) + environment["JOYOMNI_PORT"] = model_port + + try: + preload = env_enabled("JOYOMNI_PRELOAD", default=True) + except ValueError as error: + print(f"Invalid startup configuration: {error}", flush=True) + return 1 + + print(f"Starting JoyAI on public port {model_port}...", flush=True) + print( + "Routes: GET /, GET /health, POST /load, WS /ws; " + "RunPod health: GET /ping on PORT_HEALTH.", + flush=True, + ) + if preload: + print( + "Preloading the JoyAI runtime before RunPod marks this worker ready...", + flush=True, + ) + + health_process = subprocess.Popen( + [ + sys.executable, + str(repository_root / "runpod" / "health_server.py"), + ], + cwd=repository_root, + env=environment, + ) + + exit_code = 1 + try: + model_process = subprocess.Popen( + build_model_command(repository_root, preload=preload), + cwd=repository_root, + env=environment, + ) + exit_code = model_process.wait() + finally: + health_process.terminate() + try: + health_process.wait(timeout=10) + except subprocess.TimeoutExpired: + health_process.kill() + + return exit_code + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/runpod/verify_checkpoint.py b/runpod/verify_checkpoint.py new file mode 100644 index 0000000..c868541 --- /dev/null +++ b/runpod/verify_checkpoint.py @@ -0,0 +1,42 @@ +import argparse +import json +import os +import sys +from pathlib import Path + + +REPOSITORY_ROOT = Path(__file__).resolve().parents[1] +DEPLOY_ROOT = REPOSITORY_ROOT / "deploy" +if str(DEPLOY_ROOT) not in sys.path: + sys.path.insert(0, str(DEPLOY_ROOT)) + +from xvideo.checkpoint_status import checkpoint_status + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Verify the mounted JoyAI upgraded RV2V DiT checkpoint." + ) + parser.add_argument( + "--full-hash", + action="store_true", + help="Read and SHA-256 hash the complete 32.5 GB file.", + ) + args = parser.parse_args() + + checkpoint_root = Path( + os.getenv("JOYOMNI_CKPT_ROOT", "/runpod-volume/joyai/checkpoints") + ) + checkpoint = ( + checkpoint_root + / "JoyAI-Video-Edit" + / "dit" + / "joyai_video_edit_dit_0811.pth" + ) + report = checkpoint_status(checkpoint, full_hash=args.full_hash) + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 if report["status"] == "current" else 2 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/test_runpod_connection.py b/tests/test_runpod_connection.py new file mode 100644 index 0000000..1c7f93f --- /dev/null +++ b/tests/test_runpod_connection.py @@ -0,0 +1,323 @@ +import importlib.util +import os +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory +from unittest.mock import patch + + +ROOT = Path(__file__).resolve().parents[1] +START_PATH = ROOT / "runpod" / "start.py" + +SPEC = importlib.util.spec_from_file_location("runpod_start", START_PATH) +START = importlib.util.module_from_spec(SPEC) +assert SPEC.loader is not None +SPEC.loader.exec_module(START) + +CHECKPOINT_STATUS_PATH = ROOT / "deploy" / "xvideo" / "checkpoint_status.py" +CHECKPOINT_SPEC = importlib.util.spec_from_file_location( + "checkpoint_status", CHECKPOINT_STATUS_PATH +) +CHECKPOINT_STATUS = importlib.util.module_from_spec(CHECKPOINT_SPEC) +assert CHECKPOINT_SPEC.loader is not None +CHECKPOINT_SPEC.loader.exec_module(CHECKPOINT_STATUS) + + +class RunPodConnectionContractTests(unittest.TestCase): + def test_current_rv2v_checkpoint_is_pinned_and_reported_by_health(self): + downloader = (ROOT / "runpod" / "download_models.py").read_text() + launcher = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + server = ( + ROOT / "deploy" / "xvideo" / "serving" / "serve_joyomni_streaming.py" + ).read_text() + self.assertIn(CHECKPOINT_STATUS.JOYAI_DIT_RELEASE_COMMIT, downloader) + self.assertIn('"checkpoint": checkpoint_status(args.dit_ckpt)', server) + self.assertIn('$health.checkpoint.status -ne "current"', launcher) + self.assertIn("the upgraded RV2V checkpoint is active", launcher) + + def test_checkpoint_metadata_identifies_current_and_stale_weights(self): + with TemporaryDirectory() as temporary_directory: + local_dir = Path(temporary_directory) / "JoyAI-Video-Edit" + checkpoint = local_dir / "dit" / "joyai_video_edit_dit_0811.pth" + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"test checkpoint") + metadata = ( + local_dir + / ".cache" + / "huggingface" + / "download" + / "dit" + / "joyai_video_edit_dit_0811.pth.metadata" + ) + metadata.parent.mkdir(parents=True) + metadata.write_text( + CHECKPOINT_STATUS.JOYAI_DIT_RELEASE_COMMIT + + "\n" + + CHECKPOINT_STATUS.JOYAI_DIT_XET_HASH + + "\n0\n" + ) + + report = CHECKPOINT_STATUS.checkpoint_status(checkpoint) + self.assertEqual(report["status"], "current") + self.assertEqual(report["verification"], "huggingface_metadata") + + metadata.write_text("old-revision\nold-etag\n0\n") + report = CHECKPOINT_STATUS.checkpoint_status(checkpoint) + self.assertEqual(report["status"], "stale") + + def test_checkpoint_without_metadata_is_unknown_until_full_hash(self): + with TemporaryDirectory() as temporary_directory: + checkpoint = ( + Path(temporary_directory) + / "JoyAI-Video-Edit" + / "dit" + / "joyai_video_edit_dit_0811.pth" + ) + checkpoint.parent.mkdir(parents=True) + checkpoint.write_bytes(b"not the official checkpoint") + + report = CHECKPOINT_STATUS.checkpoint_status(checkpoint) + self.assertEqual(report["status"], "unknown") + + report = CHECKPOINT_STATUS.checkpoint_status( + checkpoint, full_hash=True + ) + self.assertEqual(report["status"], "stale") + self.assertEqual(report["verification"], "sha256") + + def test_preload_defaults_to_enabled(self): + with patch.dict(os.environ, {}, clear=True): + self.assertTrue(START.env_enabled("JOYOMNI_PRELOAD", default=True)) + + def test_preload_can_be_disabled_explicitly(self): + with patch.dict(os.environ, {"JOYOMNI_PRELOAD": "0"}, clear=True): + self.assertFalse(START.env_enabled("JOYOMNI_PRELOAD", default=True)) + + def test_invalid_preload_value_fails_fast(self): + with patch.dict(os.environ, {"JOYOMNI_PRELOAD": "maybe"}, clear=True): + with self.assertRaises(ValueError): + START.env_enabled("JOYOMNI_PRELOAD", default=True) + + def test_model_command_makes_preload_explicit(self): + command = START.build_model_command(ROOT, preload=True) + self.assertEqual(command[-1], "--preload") + + def test_runpod_ping_is_internal_and_checks_public_health(self): + text = (ROOT / "runpod" / "health_server.py").read_text() + self.assertIn('@app.get("/ping")', text) + self.assertIn('/health', text) + self.assertIn('status_code=204', text) + + def test_windows_script_uses_public_routes(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + self.assertIn('$baseUrl/health', text) + self.assertNotIn('$baseUrl/load', text) + self.assertNotIn('$baseUrl/ping', text) + self.assertIn('(\'"{0}"\' -f $proxyScript)', text) + + def test_windows_script_prepares_proxy_dependency_before_warming_worker(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + dependency_check = text.index('python.exe -c "import aiohttp"') + warm_request = text.index('$baseUrl/health') + self.assertLess(dependency_check, warm_request) + self.assertIn("$localReady", text) + + def test_windows_script_retries_runpod_cold_start_timeout(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + self.assertIn("[int]$WarmTimeoutSeconds = 1800", text) + self.assertIn("up to 30 minutes", text) + self.assertIn("first compiled-VAE start", text) + self.assertIn("compiled VAE and CUDA graph are active", text) + self.assertIn("$health.optimizations.vae_compile.ready", text) + self.assertIn("$health.optimizations.cuda_graph.ready", text) + self.assertIn("while (-not $workerReady", text) + self.assertIn("timed out waiting for worker", text) + self.assertIn("operation has timed out", text) + self.assertIn("$retryableStatusCodes", text) + self.assertIn("$WarmTimeoutSeconds", text) + self.assertIn("does not guarantee that GPU billing has stopped", text) + + def test_windows_script_reports_local_proxy_failures(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + self.assertIn("Local proxy check $attempt failed: $localDetail", text) + + def test_windows_script_opens_browser_only_after_local_health_is_ready(self): + text = (ROOT / "runpod" / "Start-JoyAI-Realtime-Test.ps1").read_text() + local_health_check = text.index('$localHealthUrl') + browser_open = text.index('Start-Process $localUrl') + self.assertLess(local_health_check, browser_open) + + def test_proxy_handles_websocket_and_does_not_target_public_ping(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn('request.path == "/ws"', text) + self.assertNotIn('api.runpod.ai/ping', text) + self.assertIn('await previous_upstream.send_json({"type": "stop"})', text) + self.assertIn('message=b"replaced by a refreshed browser"', text) + self.assertIn('proxy_state = request.app["proxy_state"]', text) + self.assertNotIn('request.app["active_upstream"] =', text) + + def test_proxy_relays_brotli_without_a_python_decoder(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("auto_decompress=False", text) + self.assertIn('headers["Accept-Encoding"] = "identity"', text) + + def test_proxy_suppresses_only_windows_cleanup_resets(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("proxy_exception_handler", text) + self.assertIn('getattr(error, "winerror", None) == 10054', text) + self.assertIn('"_call_connection_lost" in message', text) + + def test_h200_live_mode_disables_recording_and_uses_low_bandwidth_defaults(self): + dockerfile = (ROOT / "Dockerfile.h200").read_text() + launcher = (ROOT / "deploy" / "run_server.sh").read_text() + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn("JOYOMNI_RECORD_ENABLED=0", dockerfile) + self.assertIn("JOYOMNI_ONLINE_GATE_ENABLED=0", dockerfile) + self.assertIn("JOYOMNI_FPS=20", dockerfile) + self.assertIn("JOYOMNI_VAE_COMPILE=1", dockerfile) + self.assertIn("JOYOMNI_VAE_COMPILE_STRICT=1", dockerfile) + self.assertIn("JOYOMNI_LOAD_WARMUP_STRICT=1", dockerfile) + self.assertIn("JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS=300", dockerfile) + self.assertIn("JOYOMNI_WARMUP_BOTH_ORIENTATIONS=0", dockerfile) + self.assertIn("JOYOMNI_WARMUP_REFERENCE_BUCKETS=0", dockerfile) + self.assertIn( + "JOYOMNI_CACHE_ROOT=/runpod-volume/joyai/cache/h200-torch291-cu128", + dockerfile, + ) + self.assertIn('CACHE_ROOT="${JOYOMNI_CACHE_ROOT:-$HERE/deps/cache}"', launcher) + self.assertIn('$CACHE_ROOT/torchinductor', launcher) + self.assertIn('$CACHE_ROOT/triton', launcher) + self.assertIn('$CACHE_ROOT/nv_compute', launcher) + self.assertIn('EXTRA_ARGS+=(--record-dir "$RECORD_DIR")', launcher) + self.assertNotIn(' --record-dir "$RECORD_DIR" \\\n', launcher) + self.assertNotIn('id="downloadBubble"', html) + self.assertNotIn('id="outputStartOverlay"', html) + self.assertNotIn('id="outputIdleOverlay"', html) + self.assertNotIn('id="outputWaitOverlay"', html) + self.assertNotIn('id="outputToast"', html) + self.assertNotIn('id="sendHint"', html) + self.assertIn('class="kvreset-field keep-min" style="display:none"', html) + self.assertIn('data-upq="0.2" class="on"', html) + self.assertIn('data-fps="20" class="on"', html) + self.assertIn("let autoQTier = 0;", html) + self.assertIn("const DEFAULT_TARGET_OUTPUT_QUEUE_DELAY_MS = 100;", html) + self.assertIn("const DEFAULT_MAX_OUTPUT_QUEUE_DELAY_MS = 200;", html) + self.assertIn("const MAX_BACKEND_PENDING_FRAMES = 16;", html) + self.assertIn("const UPLINK_KEYFRAME_INTERVAL = 20;", html) + self.assertIn('autoQuality = true; autoQTier = 0;', html) + + def test_health_reports_required_runtime_optimizations(self): + server = ( + ROOT / "deploy" / "xvideo" / "serving" / "serve_joyomni_streaming.py" + ).read_text() + runtime = ( + ROOT / "deploy" / "xvideo" / "serving" / "joyomni_streaming.py" + ).read_text() + vae_compile = ( + ROOT / "deploy" / "xvideo" / "models" / "vae" / "vae_compile.py" + ).read_text() + + self.assertIn('"optimizations": (', server) + self.assertIn("app.state.runtime.optimization_status()", server) + self.assertIn("JOYOMNI_CACHE_READY_MARKER", server) + self.assertIn('"vae_compile": _vae_compile_module().runtime_status()', runtime) + self.assertIn('"cuda_graph": {', runtime) + self.assertIn('"JOYOMNI_FULL_WARMUP_TIMEOUT_SECONDS"', runtime) + self.assertIn("time.monotonic() + timeout_seconds", runtime) + self.assertIn("assert_runtime_ready", vae_compile) + self.assertIn("JOYOMNI_VAE_COMPILE_STRICT", vae_compile) + self.assertIn("def call_guard():", vae_compile) + self.assertIn('"thread_call_guard": compile_enabled()', vae_compile) + self.assertIn("with _vc.call_guard():", runtime) + self.assertIn("with _vc.call_guard():", ( + ROOT / "deploy" / "xvideo" / "models" / "pipeline.py" + ).read_text()) + + def test_browser_codec_probes_cannot_lock_the_send_button(self): + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn("const CODEC_PROBE_TIMEOUT_MS = 1200;", html) + self.assertGreaterEqual(html.count("codecProbeWithTimeout("), 3) + self.assertIn('if (!started) throw new Error("session start was not sent")', html) + self.assertIn('setSendBusy(false, "");', html) + + def test_live_session_pings_prevent_false_idle_disconnects(self): + server = ( + ROOT / "deploy" / "xvideo" / "serving" / "serve_joyomni_streaming.py" + ).read_text() + self.assertIn('JOYOMNI_SESSION_IDLE_TIMEOUT_SECONDS", "60"', server) + self.assertIn('ws_debug["last_client_activity_at"] = time.time()', server) + ping_handler = server.index('elif msg_type == "ping":') + client_activity = server.rindex( + 'last_activity = time.monotonic()', 0, ping_handler + ) + payload_decode = server.rindex( + 'payload = json.loads(message["text"])', 0, ping_handler + ) + self.assertGreater(client_activity, payload_decode) + self.assertIn("releasing session after", server) + + def test_browser_recovers_an_unexpected_websocket_close(self): + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn("let streamingWanted = false;", html) + self.assertIn("function scheduleReconnect()", html) + self.assertIn("RECONNECT_MAX_DELAY_MS", html) + self.assertIn("scheduleReconnect();", html) + self.assertIn("streamingWanted = false;\n cancelReconnect();", html) + self.assertIn("streamingWanted = true;\n await start();", html) + + def test_proxy_reports_websocket_close_codes(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("JoyAI WebSocket connected through RunPod.", text) + self.assertIn("JoyAI WebSocket closed", text) + self.assertIn('f"(browser={downstream_code}, RunPod={upstream_code})."', text) + + def test_proxy_does_not_duplicate_browser_application_heartbeat(self): + proxy = (ROOT / "runpod" / "local_proxy.py").read_text() + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn("heartbeat=None", proxy) + self.assertNotIn("heartbeat=10", proxy) + self.assertIn('{ type: "ping", t: Date.now()', html) + + def test_reference_person_replacement_has_an_explicit_prompt(self): + html = (ROOT / "deploy" / "static" / "index.html").read_text() + self.assertIn('label_en: "Reference Image"', html) + self.assertIn('title_en: "My Reference Person"', html) + self.assertIn( + "Replace the main subject in the video with the person shown in the " + "uploaded reference image.", + html, + ) + + def test_proxy_pins_reconnects_to_one_runpod_worker(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn('RUNPOD_WORKER_HEADER = "X-Runpod-Worker-Id"', text) + self.assertIn("headers[RUNPOD_WORKER_HEADER] = worker_id", text) + self.assertIn("remember_worker(request.app, response.headers)", text) + self.assertIn('"worker_id": None', text) + + def test_websocket_and_keepalive_do_not_wait_on_worker_affinity(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertGreaterEqual( + text.count("upstream_headers(request.app, affinity=False)"), + 1, + ) + self.assertIn("upstream_headers(app, affinity=False)", text) + self.assertNotIn('f"strict-resume {worker_id}"', text) + + def test_proxy_keeps_worker_active_during_websocket_stream(self): + text = (ROOT / "runpod" / "local_proxy.py").read_text() + self.assertIn("RUNPOD_KEEPALIVE_SECONDS = 3.0", text) + self.assertIn("async def keep_worker_active", text) + self.assertIn('f"{app[\'upstream\']}/health"', text) + self.assertIn("asyncio.create_task(keep_worker_active(request.app))", text) + self.assertIn("keepalive_task.cancel()", text) + + def test_runpod_settings_limit_one_viewer_tests_to_one_worker(self): + text = (ROOT / "runpod" / "README.md").read_text() + self.assertIn("| Max workers | `1`", text) + self.assertIn("| Idle timeout | `60` seconds |", text) + self.assertIn("soft", text) + + +if __name__ == "__main__": + unittest.main()