diff --git a/python/pyproject.toml b/python/pyproject.toml index 36073b40a..9d1ca121b 100755 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -111,6 +111,7 @@ diffusion = [ "av==16.1.0", "scikit-image==0.25.2", "trimesh>=4.0.0", + "websockets", "xatlas", ] diff --git a/python/sglang/multimodal_gen/apps/realtime_webui/app.js b/python/sglang/multimodal_gen/apps/realtime_webui/app.js index 57eace86e..e42cb30a1 100644 --- a/python/sglang/multimodal_gen/apps/realtime_webui/app.js +++ b/python/sglang/multimodal_gen/apps/realtime_webui/app.js @@ -10,6 +10,8 @@ const DEFAULT_PREVIEW_OUTPUT_QUALITY = 95; const DEFAULT_TARGET_FPS = 25; const DEFAULT_FRAME_INTERPOLATION_EXP = 1; const DEFAULT_FRAME_INTERPOLATION_SCALE = 1.0; +const DEFAULT_UPSCALING_SCALE = 2; +const DEFAULT_PREVIEW_SCALE = 120; const RECONNECT_CLOSE_TIMEOUT_MS = 15000; const LIVE_QUEUE_SECONDS = 0.45; const LOW_LATENCY_FPS_FLOOR = 10; @@ -222,9 +224,13 @@ let encodedDecodeErrors = 0; let socketHadError = false; let socketCloseExpected = false; let socketServerError = ""; +let renderedPreviewFrames = 0; +let previewScaleFrame = 0; const decodeRequests = new Map(); let controlStateController = null; +const stage = document.querySelector(".stage"); +const previewFrame = document.querySelector(".preview-frame"); const canvas = $("viewport"); const ctx = canvas.getContext("2d", { alpha: false }); const scratchCanvas = document.createElement("canvas"); @@ -235,6 +241,12 @@ function setStatus(text, kind = "") { $("statusDot").className = "dot" + (kind ? ` ${kind}` : ""); } +function setPreviewState(state) { + if (!stage) return; + stage.dataset.previewState = state; + canvas.setAttribute("aria-busy", state === "waiting" ? "true" : "false"); +} + function addHistory(text) { const item = document.createElement("span"); item.textContent = text; @@ -243,19 +255,40 @@ function addHistory(text) { } function drawIdle() { - const w = canvas.width, h = canvas.height; - const g = ctx.createLinearGradient(0, 0, w, h); - g.addColorStop(0, "#171a16"); - g.addColorStop(0.58, "#253628"); - g.addColorStop(1, "#8f4a37"); - ctx.fillStyle = g; + const w = 1280, h = 720; + if (canvas.width !== w || canvas.height !== h) { + canvas.width = w; + canvas.height = h; + } + setPreviewState("idle"); + renderedPreviewFrames = 0; + ctx.fillStyle = "#11140f"; ctx.fillRect(0, 0, w, h); - ctx.fillStyle = "rgba(238,241,236,.82)"; - ctx.font = "600 46px Avenir Next, sans-serif"; - ctx.fillText("sglang-diffusion", 56, 86); - for (let i = 0; i < 18; i++) { - ctx.fillStyle = `rgba(238,241,236,${0.05 + i * 0.015})`; - ctx.fillRect(56 + i * 64, h - 86 - i * 7, 42, 42 + i * 4); + + const surface = ctx.createLinearGradient(0, 0, 0, h); + surface.addColorStop(0, "rgba(238,241,236,0.045)"); + surface.addColorStop(0.5, "rgba(238,241,236,0.012)"); + surface.addColorStop(1, "rgba(0,0,0,0.16)"); + ctx.fillStyle = surface; + ctx.fillRect(0, 0, w, h); + + ctx.strokeStyle = "rgba(238,241,236,0.11)"; + ctx.lineWidth = 1; + ctx.strokeRect(0.5, 0.5, w - 1, h - 1); + ctx.strokeStyle = "rgba(238,241,236,0.08)"; + ctx.beginPath(); + if (ctx.roundRect) { + ctx.roundRect(w * 0.38, h * 0.42, w * 0.24, h * 0.16, 18); + } else { + ctx.rect(w * 0.38, h * 0.42, w * 0.24, h * 0.16); + } + ctx.stroke(); + + ctx.fillStyle = "rgba(238,241,236,0.22)"; + for (let i = -1; i <= 1; i++) { + ctx.beginPath(); + ctx.arc(w * 0.5 + i * 22, h * 0.5, 4.5, 0, Math.PI * 2); + ctx.fill(); } } @@ -277,6 +310,7 @@ function resetStreamStats() { currentReceiveChunk = null; currentReceiveChunkFrames = 0; encodedDecodeErrors = 0; + renderedPreviewFrames = 0; controlStateController?.reset({ sendRelease: false }); resetDecoderState(); updateStats(); @@ -290,6 +324,7 @@ function resetStreamStats() { $("theoreticalFpsText").textContent = "-"; $("chunkText").textContent = "chunk -"; $("payloadMode").textContent = selectedTransportLabel(); + updateOutputSizeText(); } function rejectPendingDecodes(message) { @@ -669,27 +704,22 @@ function drawFrame(image) { drawSource = scratchCanvas; } - const rect = canvas.getBoundingClientRect(); - const dpr = Math.max(1, Math.min(window.devicePixelRatio || 1, 2)); - const targetWidth = Math.max(1, Math.round(rect.width * dpr)); - const targetHeight = Math.max( - 1, - Math.round((targetWidth * sourceHeight) / sourceWidth), - ); - if (canvas.width !== targetWidth || canvas.height !== targetHeight) { - canvas.width = targetWidth; - canvas.height = targetHeight; + if (canvas.width !== sourceWidth || canvas.height !== sourceHeight) { + canvas.width = sourceWidth; + canvas.height = sourceHeight; } ctx.imageSmoothingEnabled = true; ctx.imageSmoothingQuality = "high"; - ctx.drawImage(drawSource, 0, 0, canvas.width, canvas.height); + ctx.drawImage(drawSource, 0, 0, sourceWidth, sourceHeight); + renderedPreviewFrames += 1; + setPreviewState("live"); if (!(image instanceof ImageData)) image.close?.(); } function renderLoop(now) { const targetFps = playbackFps || Number($("fps").value || DEFAULT_TARGET_FPS); const queueSeconds = queue.length / Math.max(1, targetFps); - const catchupFps = queueSeconds > LOW_LATENCY_QUEUE_SECONDS + const catchupFps = !$("superResolution").checked && queueSeconds > LOW_LATENCY_QUEUE_SECONDS ? Math.min(MAX_CATCHUP_FPS, Math.ceil(queue.length / LOW_LATENCY_QUEUE_SECONDS)) : targetFps; const targetMs = 1000 / Math.max(1, catchupFps); @@ -770,6 +800,7 @@ async function setPresetReference(preset) { function showError(error) { setStatus("Reference load failed", "error"); + if (!renderedPreviewFrames) setPreviewState("idle"); addHistory(error.message || "reference load failed"); } @@ -795,10 +826,12 @@ function abortCurrentSession(reason = "session closed by client", { clearQueueOnClose = false; if (!keepConnectDisabled) $("connectBtn").disabled = false; setStatus("Closed"); + if (!renderedPreviewFrames) setPreviewState("idle"); return null; } if (!keepConnectDisabled) $("connectBtn").disabled = false; setStatus(expectedClose ? "Closing" : "Aborting"); + if (!renderedPreviewFrames) setPreviewState("idle"); addHistory(reason); socket.close(expectedClose ? 1000 : 1011, reason.slice(0, 120)); return socket; @@ -828,6 +861,7 @@ function waitForSocketClose(socket, timeoutMs = RECONNECT_CLOSE_TIMEOUT_MS) { async function connect() { $("connectBtn").disabled = true; setStatus("Preparing"); + setPreviewState("waiting"); addHistory("preparing session"); try { if (ws && ws.readyState !== WebSocket.CLOSED) { @@ -845,12 +879,14 @@ async function connect() { const firstFrame = await readFirstFrame(); if (!firstFrame) { setStatus("Pick a reference", "error"); + setPreviewState("idle"); addHistory("reference image required"); $("connectBtn").disabled = false; return; } const previewTransportParams = readPreviewTransportParams(); const frameInterpolationParams = readFrameInterpolationParams(); + const superResolutionParams = readSuperResolutionParams(); const init = compact({ type: "init", model: $("model").value, @@ -867,6 +903,7 @@ async function connect() { first_frame: firstFrame, ...previewTransportParams, ...frameInterpolationParams, + ...superResolutionParams, }); document.activeElement?.blur?.(); canvas.tabIndex = 0; @@ -908,6 +945,7 @@ async function connect() { setStatus("Closed"); addHistory(closeText); } + if (!renderedPreviewFrames) setPreviewState("idle"); socketCloseExpected = false; }; socket.onerror = () => { @@ -928,6 +966,7 @@ async function connect() { } catch (error) { $("connectBtn").disabled = false; setStatus("Init failed", "error"); + if (!renderedPreviewFrames) setPreviewState("idle"); addHistory(error.message || "init failed"); } } @@ -1019,6 +1058,7 @@ async function decodeAndEnqueueFrameBatch(header, data, epoch) { frames += chunkFrameCount; bytes += payloadBytes; $("payloadMode").textContent = header.encoding || "raw RGB"; + updateOutputSizeFromHeader(header); updatePlaybackPace(header, performance.now(), chunkFrameCount); setStatus("Live", "live"); updateStats(header); @@ -1059,9 +1099,10 @@ function updatePlaybackPace(header, now, frameCount) { if (waitSeconds > 0) { const generatedFps = currentReceiveChunkFrames / Math.max(0.001, waitSeconds); const requestedFps = Number($("fps").value || DEFAULT_TARGET_FPS); + const playbackFloor = $("superResolution").checked ? 1 : LOW_LATENCY_FPS_FLOOR; playbackFps = Math.min( requestedFps, - Math.max(LOW_LATENCY_FPS_FLOOR, generatedFps * 1.05), + Math.max(playbackFloor, generatedFps * 1.05), ); const latencyText = `${waitSeconds.toFixed(1)}s · ${playbackFps.toFixed(1)}fps`; $("latencyText").textContent = latencyText; @@ -1114,6 +1155,7 @@ async function applyPreset(preset, options = {}) { $("prompt").value = preset.prompt; $("size").value = preset.size; $("fps").value = preset.fps; + updateOutputSizeText(); await setPresetReference(preset); if (sendRuntimeEvents) { sendEvent("prompt", preset.prompt, `prompt update · ${preset.name}`); @@ -1249,6 +1291,67 @@ function readFrameInterpolationParams() { }; } +function readUpscalingScale() { + return Number($("upscalingScale").value || DEFAULT_UPSCALING_SCALE); +} + +function readSuperResolutionParams() { + if (!$("superResolution").checked) return {}; + return { + enable_upscaling: true, + upscaling_scale: readUpscalingScale(), + }; +} + +function parseSizeValue(sizeText) { + const match = /^(\d+)\s*x\s*(\d+)$/i.exec(String(sizeText || "").trim()); + if (!match) return null; + return { + width: Number(match[1]), + height: Number(match[2]), + }; +} + +function updateOutputSizeText(width = null, height = null) { + let outputWidth = Number(width || 0); + let outputHeight = Number(height || 0); + const srEnabled = $("superResolution").checked; + const scale = srEnabled ? readUpscalingScale() : 1; + if (!outputWidth || !outputHeight) { + const base = parseSizeValue($("size").value); + if (base) { + outputWidth = base.width * scale; + outputHeight = base.height * scale; + } + } + $("outputSizeText").textContent = outputWidth && outputHeight + ? `${outputWidth}x${outputHeight}${srEnabled ? ` · SR ${scale}x` : ""}` + : "-"; +} + +function updateOutputSizeFromHeader(header) { + const width = Number(header.width || 0); + const height = Number(header.height || 0); + if (width && height) updateOutputSizeText(width, height); +} + +function updateSuperResolutionControls() { + $("upscalingScale").disabled = !$("superResolution").checked; + updateOutputSizeText(); +} + +function setPreviewScale(value) { + if (!previewFrame) return; + const scale = Math.max(80, Math.min(170, Number(value || DEFAULT_PREVIEW_SCALE))); + $("previewScale").value = String(scale); + $("previewScaleText").textContent = `${scale}%`; + if (previewScaleFrame) cancelAnimationFrame(previewScaleFrame); + previewScaleFrame = requestAnimationFrame(() => { + previewScaleFrame = 0; + previewFrame.style.setProperty("--preview-scale", String(scale / 100)); + }); +} + function selectedTransportLabel() { const select = $("transportFormat"); return select.options[select.selectedIndex]?.textContent || "raw RGB"; @@ -1299,6 +1402,11 @@ async function applyQueryParams() { if (model) $("model").value = model; $("transportFormat").value = params.get("transport") || DEFAULT_PREVIEW_OUTPUT_FORMAT; $("transportQuality").value = params.get("quality") || String(DEFAULT_PREVIEW_OUTPUT_QUALITY); + const srParam = params.get("sr"); + $("superResolution").checked = srParam === "1" || srParam === "true"; + $("upscalingScale").value = params.get("sr_scale") || String(DEFAULT_UPSCALING_SCALE); + setPreviewScale(params.get("preview_scale") || params.get("zoom")); + updateSuperResolutionControls(); return { model: Boolean(model), preset: Boolean(presetKey && appliedPreset), @@ -1394,6 +1502,8 @@ function unpack(buf) { renderPresets(); drawIdle(); +setPreviewScale(DEFAULT_PREVIEW_SCALE); +updateSuperResolutionControls(); applyPreset(presets[0], { sendRuntimeEvents: false }) .then(applyQueryParams) .then((query) => queryServerModelInfo({ @@ -1406,6 +1516,10 @@ $("stopBtn").onclick = () => closeSession(); $("sendPromptBtn").onclick = () => sendEvent("prompt", $("prompt").value); $("enhanceBtn").onclick = enhancePrompt; $("firstFrame").onchange = () => drawReferencePreview($("firstFrame").files[0]); +$("size").addEventListener("input", () => updateOutputSizeText()); +$("superResolution").addEventListener("change", updateSuperResolutionControls); +$("upscalingScale").addEventListener("change", () => updateOutputSizeText()); +$("previewScale").addEventListener("input", () => setPreviewScale($("previewScale").value)); $("serverUrl").addEventListener("change", () => { queryServerModelInfo({ applyPresetForModel: true }).catch(showError); }); diff --git a/python/sglang/multimodal_gen/apps/realtime_webui/index.html b/python/sglang/multimodal_gen/apps/realtime_webui/index.html index 4647f6bcc..c0cfadac8 100644 --- a/python/sglang/multimodal_gen/apps/realtime_webui/index.html +++ b/python/sglang/multimodal_gen/apps/realtime_webui/index.html @@ -4,7 +4,7 @@ sglang-diffusion Realtime Studio - +
@@ -51,6 +51,15 @@ +
+ + +
@@ -60,70 +69,82 @@ -
-
- - Idle - chunk - - - render 0 fps - theoretical - - wait - -
- -
-
- Move -
- - - - - - +
+
+
+ + Idle + chunk - + + + output 832x480 + render 0 fps + theoretical - + wait - +
+
+ +
-
- Look -
- - - - - - +
+
+ Move +
+ + + + + + +
+
+
+ Look +
+ + + + + + +
-
-
- queue 0 - frames 0 - 0 MB -
-
- Payloadwebp - Server send- - Chunk bytes- - Chunk wait- - Decode- - Display lag- -
-
+
+ queue 0 + frames 0 + 0 MB +
+
+ Payloadwebp + Server send- + Chunk bytes- + Chunk wait- + Decode- + Display lag- +
+
-
-
LingBot
-
- 25 fps target - chunked stream - 480p/720p - Cam + Act -
-
Presets
-
-
History
-
+
+
LingBot
+
+ 25 fps target + chunked stream + 480p/720p + Cam + Act +
+
Presets
+
+
History
+
+
- + diff --git a/python/sglang/multimodal_gen/apps/realtime_webui/styles.css b/python/sglang/multimodal_gen/apps/realtime_webui/styles.css index 5481adc19..77bf5454b 100644 --- a/python/sglang/multimodal_gen/apps/realtime_webui/styles.css +++ b/python/sglang/multimodal_gen/apps/realtime_webui/styles.css @@ -32,7 +32,7 @@ button:disabled { cursor: wait; opacity: 0.64; transform: none; } .shell { display: grid; - grid-template-columns: minmax(260px, 320px) minmax(420px, 1fr) minmax(260px, 320px); + grid-template-columns: minmax(260px, 320px) minmax(560px, 1fr); gap: 18px; min-height: 100vh; padding: 18px; @@ -136,6 +136,13 @@ textarea { resize: vertical; line-height: 1.45; } input:focus, textarea:focus, select:focus { border-color: var(--accent); box-shadow: 0 0 0 3px rgba(185, 84, 60, 0.12); } .split { display: grid; grid-template-columns: 1fr 1fr; gap: 10px; } +.output-options { + align-items: end; +} +.output-options .toggle-row { + min-height: 40px; + margin: 12px 0; +} .actions { display: grid; grid-template-columns: 1fr 0.7fr; gap: 10px; margin-top: 16px; } .toggle-row { display: flex; @@ -208,6 +215,12 @@ button:focus-visible { } .wide { width: 100%; margin-top: 10px; } +.workspace { + display: grid; + gap: 18px; + min-width: 0; +} + .stage { position: relative; display: grid; @@ -215,7 +228,7 @@ button:focus-visible { align-self: start; justify-self: center; width: 100%; - max-width: 1040px; + max-width: min(1500px, 100%); overflow: hidden; border: 1px solid #11140f; border-radius: 8px; @@ -223,6 +236,40 @@ button:focus-visible { box-shadow: var(--shadow); } +.preview-frame { + position: relative; + display: grid; + place-items: center; + justify-self: center; + width: min(calc(1040px * var(--preview-scale, 1.2)), 100%); + overflow: hidden; + background: #11140f; + contain: paint; + isolation: isolate; +} + +.preview-frame::before { + content: ""; + position: absolute; + inset: 0; + z-index: 0; + pointer-events: none; + background: linear-gradient( + 180deg, + rgba(238, 241, 236, 0.045), + transparent 34%, + rgba(0, 0, 0, 0.18) + ); +} + +.preview-frame::after { + content: none; +} + +.stage[data-preview-state="waiting"] .preview-frame::after { + animation: none; +} + .topbar, .timeline { display: flex; align-items: center; @@ -235,6 +282,28 @@ button:focus-visible { } .topbar-spacer { flex: 1; } +.preview-scale-control { + display: inline-flex; + align-items: center; + gap: 8px; + min-width: 190px; + margin: 0; + color: rgba(232, 234, 223, 0.72); + font-size: 11px; +} +.preview-scale-control input { + width: 92px; + min-width: 72px; + padding: 0; + border: 0; + background: transparent; + accent-color: #eef1ec; +} +.preview-scale-control b { + min-width: 36px; + color: #fffdf7; + font-weight: 650; +} #statusText { display: inline-block; min-width: 92px; @@ -261,13 +330,61 @@ button:focus-visible { .dot.error { background: var(--accent); } #viewport { + position: relative; + z-index: 1; display: block; width: 100%; height: auto; - max-height: min(56vh, 560px); + max-height: min(calc(56vh * var(--preview-scale, 1.2)), 82vh); min-height: 0; object-fit: contain; image-rendering: auto; + transform: translateZ(0); +} + +.preview-overlay { + position: absolute; + inset: 0; + z-index: 3; + display: none; + place-items: center; + pointer-events: none; + background: transparent; +} + +.stage[data-preview-state="waiting"] .preview-overlay { + display: grid; +} + +.preview-loader { + position: relative; + width: 7px; + height: 7px; + border-radius: 50%; + background: rgba(238, 241, 236, 0.88); + animation: previewDotPulse 1.05s ease-in-out infinite; +} + +.preview-loader::before, +.preview-loader::after { + content: ""; + position: absolute; + top: 0; + width: 7px; + height: 7px; + border-radius: 50%; + background: rgba(238, 241, 236, 0.88); + animation: previewDotPulse 1.05s ease-in-out infinite; +} + +.preview-loader::before { + left: -16px; + animation-delay: -0.18s; +} + +.preview-loader::after { + left: 16px; + animation-delay: 0.18s; } .stage-controls { @@ -366,7 +483,7 @@ button:focus-visible { .spec-grid { display: grid; - grid-template-columns: repeat(2, 1fr); + grid-template-columns: repeat(4, minmax(120px, 1fr)); gap: 8px; margin-bottom: 12px; } @@ -390,18 +507,18 @@ button:focus-visible { } .presets { - position: sticky; - top: 18px; - max-height: calc(100vh - 36px); - overflow: auto; + position: static; + max-height: none; + overflow: visible; scrollbar-gutter: stable; } .preset-list { display: grid; + grid-template-columns: repeat(auto-fill, minmax(220px, 1fr)); gap: 7px; - min-height: 320px; - max-height: min(54vh, 560px); + min-height: 0; + max-height: 230px; margin-bottom: 12px; overflow: auto; padding-right: 3px; @@ -483,7 +600,7 @@ button:focus-visible { .history-list { display: grid; gap: 7px; - max-height: 76px; + max-height: 92px; overflow: auto; } @@ -499,10 +616,17 @@ button:focus-visible { @media (max-width: 980px) { .shell { grid-template-columns: 1fr; } .presets { position: static; max-height: none; overflow: visible; } + .spec-grid { grid-template-columns: repeat(2, 1fr); } .preset-list { min-height: 260px; max-height: 420px; } #viewport { max-height: 420px; } .stage-controls { grid-template-columns: 1fr; } .stage-telemetry { grid-template-columns: repeat(2, minmax(0, 1fr)); } .topbar { flex-wrap: wrap; height: auto; min-height: 44px; padding: 10px 14px; } .topbar-spacer { display: none; } + .preview-scale-control { min-width: 160px; } +} + +@keyframes previewDotPulse { + 0%, 80%, 100% { opacity: 0.32; } + 40% { opacity: 1; } } diff --git a/python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py b/python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py index 873e9d9c1..83d5295f7 100644 --- a/python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py +++ b/python/sglang/multimodal_gen/runtime/postprocess/realesrgan_upscaler.py @@ -27,12 +27,26 @@ logger = init_logger(__name__) # Default HuggingFace repo and filename for Real-ESRGAN weights _DEFAULT_REALESRGAN_HF_REPO = "ai-forever/Real-ESRGAN" _DEFAULT_REALESRGAN_FILENAME = "RealESRGAN_x4.pth" +_DEFAULT_REALESRGAN_FILENAMES_BY_SCALE = { + 2: "RealESRGAN_x2.pth", + 4: "RealESRGAN_x4.pth", + 8: "RealESRGAN_x8.pth", +} _LOW_MEMORY_TILED_UPSCALE_FREE_BYTES = 2 * 1024**3 _REALESRGAN_TILE_SIZE = 256 _REALESRGAN_TILE_PAD = 32 # Module-level cache: model_path -> UpscalerModel instance _MODEL_CACHE: dict[str, "UpscalerModel"] = {} +_RESOLVED_MODEL_PATH_CACHE: dict[str, str] = {} + + +def _default_model_path_for_scale(scale: int) -> str: + filename = _DEFAULT_REALESRGAN_FILENAMES_BY_SCALE.get( + int(scale), + _DEFAULT_REALESRGAN_FILENAME, + ) + return f"{_DEFAULT_REALESRGAN_HF_REPO}:{filename}" # --------------------------------------------------------------------------- @@ -196,6 +210,17 @@ def _build_net_from_state_dict(state_dict: dict) -> nn.Module: if "conv_first.weight" in state_dict: # RRDBNet (e.g., RealESRGAN_x4plus) num_feat = state_dict["conv_first.weight"].shape[0] + in_channels = state_dict["conv_first.weight"].shape[1] + if in_channels == 3: + scale = 4 + elif in_channels == 12: + scale = 2 + elif in_channels == 48: + scale = 1 + else: + raise ValueError( + f"Unsupported RRDBNet conv_first input channels: {in_channels}" + ) num_block = sum( 1 for k in state_dict @@ -203,15 +228,16 @@ def _build_net_from_state_dict(state_dict: dict) -> nn.Module: ) num_grow_ch = state_dict["body.0.rdb1.conv1.weight"].shape[0] logger.info( - "Detected RRDBNet: num_feat=%d, num_block=%d, num_grow_ch=%d", + "Detected RRDBNet: num_feat=%d, num_block=%d, num_grow_ch=%d, scale=%d", num_feat, num_block, num_grow_ch, + scale, ) return RRDBNet( num_in_ch=3, num_out_ch=3, - scale=4, + scale=scale, num_feat=num_feat, num_block=num_block, num_grow_ch=num_grow_ch, @@ -544,7 +570,7 @@ class ImageUpscaler: def _ensure_model_loaded(self) -> UpscalerModel: """Download/load Real-ESRGAN weights, detect arch, and cache globally.""" - model_path = self._model_path or _DEFAULT_REALESRGAN_HF_REPO + model_path = self._model_path or _default_model_path_for_scale(self._scale) # Resolve: local .pth pass-through, or HF repo → download single file resolved_path = _resolve_model_path(model_path) @@ -668,7 +694,12 @@ def _resolve_model_path(model_path: str) -> str: - A HuggingFace ``repo_id:filename`` → downloads *filename* from *repo_id*, allowing users to specify custom weight files hosted on HF. """ + cached_path = _RESOLVED_MODEL_PATH_CACHE.get(model_path) + if cached_path is not None: + return cached_path + if os.path.isfile(model_path): + _RESOLVED_MODEL_PATH_CACHE[model_path] = model_path return model_path # Parse optional "repo_id:filename" syntax; fall back to default filename. @@ -704,6 +735,7 @@ def _resolve_model_path(model_path: str) -> str: f"'repo_id:filename' format (e.g. 'my-org/my-esrgan:weights.pth'). " f"Original error: {e}" ) from e + _RESOLVED_MODEL_PATH_CACHE[model_path] = local_path return local_path diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_output_materialization.py b/python/sglang/multimodal_gen/test/unit/realtime/test_output_materialization.py index 37f9fc7e4..038fc551c 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_output_materialization.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_output_materialization.py @@ -114,3 +114,60 @@ def test_raw_rgb_frame_batches_convert_batched_video_tensor_to_thwc_bytes(): assert np.all(first[..., 0] == 255) assert np.all(first[..., 1] == 127) assert np.all(first[..., 2] == 0) + + +def test_raw_rgb_frame_batches_apply_realtime_upscaling(monkeypatch): + calls = [] + + def fake_batch_upscale_frames(frames, *, model_path, scale): + calls.append((model_path, scale, [frame.shape for frame in frames])) + return [ + np.repeat(np.repeat(frame, scale, axis=0), scale, axis=1) + for frame in frames + ] + + from sglang.multimodal_gen.runtime import postprocess + + monkeypatch.setattr(postprocess, "batch_upscale_frames", fake_batch_upscale_frames) + + req = type( + "Req", + (), + { + "data_type": DataType.VIDEO, + "fps": 24, + "output_compression": None, + "enable_frame_interpolation": False, + "frame_interpolation_exp": 1, + "frame_interpolation_scale": 1.0, + "frame_interpolation_model_path": None, + "enable_upscaling": True, + "upscaling_model_path": "mock-sr", + "upscaling_scale": 2, + "request_id": "req", + "block_idx": 0, + }, + )() + output_batch = OutputBatch(audio_sample_rate=None) + + def post_process_sample(_sample, *_args, **kwargs): + assert kwargs["enable_upscaling"] is False + return [np.array([[[1, 2, 3]]], dtype=np.uint8)] + + frame_batches, metadata = build_raw_rgb_frame_batches( + torch.zeros(1, 3, 1, 1, 1), + req, + output_batch, + post_process_sample, + ) + + assert calls == [("mock-sr", 2, [(1, 1, 3)])] + assert metadata == { + "format": "rgb24", + "width": 2, + "height": 2, + "channels": 3, + "bytes_per_frame": 12, + } + assert len(frame_batches) == 1 + assert frame_batches[0][0] == bytes([1, 2, 3] * 4) diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_webui.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_webui.py index 884013ed7..73cef4b84 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_webui.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_webui.py @@ -22,6 +22,14 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts(): assert 'const DEFAULT_PREVIEW_OUTPUT_FORMAT = "webp";' in app_js assert 'id="transportFormat"' in index_html assert 'id="fps" type="number" value="25"' in index_html + assert 'id="superResolution" type="checkbox"' in index_html + assert 'id="upscalingScale"' in index_html + assert 'class="workspace"' in index_html + assert 'class="preview-frame"' in index_html + assert 'id="previewOverlay" class="preview-overlay"' in index_html + assert 'id="previewScale" type="range" min="80" max="170" value="120"' in index_html + assert 'id="previewScaleText"' in index_html + assert 'id="outputSizeText"' in index_html assert 'id="frameInterpolation" type="checkbox" />' in index_html assert ( 'id="serverUrl" value="ws://127.0.0.1:30000/v1/realtime_video/generate"' @@ -40,15 +48,29 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts(): assert "Info" not in index_html assert 'id="steps" type="number" value="4"' in index_html assert 'id="guidance" type="number" value="1"' in index_html - assert "styles.css?v=realtime-fixes-v25" in index_html - assert "app.js?v=realtime-fixes-v30" in index_html + assert "styles.css?v=realtime-sr-v35" in index_html + assert "app.js?v=realtime-sr-v35" in index_html assert 'const DECODER_WORKER_URL = "./decoder_worker.js?v=rgb-worker-v6";' in app_js assert "const DEFAULT_TARGET_FPS = 25;" in app_js assert "const DEFAULT_FRAME_INTERPOLATION_EXP = 1;" in app_js assert "const DEFAULT_FRAME_INTERPOLATION_SCALE = 1.0;" in app_js + assert "const DEFAULT_UPSCALING_SCALE = 2;" in app_js + assert "const DEFAULT_PREVIEW_SCALE = 120;" in app_js + assert 'setPreviewState("waiting")' in app_js + assert "stage.dataset.previewState = state" in app_js + assert 'document.querySelector(".preview-frame")' in app_js + assert 'previewFrame.style.setProperty("--preview-scale"' in app_js + assert "cancelAnimationFrame(previewScaleFrame)" in app_js assert "enable_frame_interpolation: true" in app_js assert "frame_interpolation_exp: DEFAULT_FRAME_INTERPOLATION_EXP" in app_js assert "frame_interpolation_scale: DEFAULT_FRAME_INTERPOLATION_SCALE" in app_js + assert "readSuperResolutionParams()" in app_js + assert "enable_upscaling: true" in app_js + assert "upscaling_scale: readUpscalingScale()" in app_js + assert "updateOutputSizeFromHeader(header)" in app_js + assert "setPreviewScale(DEFAULT_PREVIEW_SCALE)" in app_js + assert "preview_scale" in app_js + assert "sr_scale" in app_js assert "elapsedMs % targetMs" in app_js assert "liveQueueFrameFloor(header, chunkFrameCount)" in app_js assert ( @@ -77,3 +99,9 @@ def test_realtime_webui_presets_do_not_emit_camera_scripts(): assert 'message.type === "chunk_stats"' in app_js assert "chunkTotal > 0 ? numFrames / chunkTotal" in app_js assert ".stage-stat" in styles_css + assert ".workspace" in styles_css + assert ".preview-frame" in styles_css + assert ".preview-overlay" in styles_css + assert "@keyframes previewSweep" in styles_css + assert ".preview-scale-control" in styles_css + assert "--preview-scale" in styles_css