[diffusion] improve: improve realtime webui playback pacing (#27148)

This commit is contained in:
Mick
2026-06-04 00:33:56 +08:00
committed by GitHub
parent 45a66f4088
commit b0f78bef97
19 changed files with 1980 additions and 310 deletions
File diff suppressed because it is too large Load Diff
@@ -80,24 +80,36 @@ function rawFramesToRgbaBuffers(header, payload) {
return buffers;
}
async function encodedFrameToRgbaBuffers(header, payload) {
if (typeof createImageBitmap === "undefined" || typeof OffscreenCanvas === "undefined") {
function splitEncodedPayload(header, payload) {
const bytes = payload instanceof Uint8Array ? payload : new Uint8Array(payload);
const lengths = Array.isArray(header.payload_lengths) && header.payload_lengths.length
? header.payload_lengths.map(Number)
: [bytes.byteLength];
const payloads = [];
let offset = 0;
for (const length of lengths) {
payloads.push(bytes.buffer.slice(
bytes.byteOffset + offset,
bytes.byteOffset + offset + length,
));
offset += length;
}
return payloads;
}
async function encodedFramesToImageBitmaps(header, payload) {
if (typeof createImageBitmap === "undefined") {
throw new Error("This browser does not support worker image decoding");
}
const blob = new Blob([payload], { type: header.content_type });
const bitmap = await createImageBitmap(blob);
const width = bitmap.width;
const height = bitmap.height;
const canvas = new OffscreenCanvas(width, height);
const ctx = canvas.getContext("2d", { alpha: false });
ctx.drawImage(bitmap, 0, 0);
const image = ctx.getImageData(0, 0, width, height);
bitmap.close?.();
const frames = await Promise.all(splitEncodedPayload(header, payload).map((framePayload) => (
createImageBitmap(new Blob([framePayload], { type: header.content_type }))
)));
return {
width,
height,
frames: [image.data.buffer],
width: frames[0]?.width || 0,
height: frames[0]?.height || 0,
frame_type: "bitmap",
frames,
};
}
@@ -107,12 +119,13 @@ async function decode(header, payload) {
header.content_type === WEBP_FRAME_CONTENT_TYPE ||
header.content_type === JPEG_FRAME_CONTENT_TYPE
) {
const decoded = await encodedFrameToRgbaBuffers(header, payload);
const decoded = await encodedFramesToImageBitmaps(header, payload);
return {
id: header.__decode_id,
width: decoded.width,
height: decoded.height,
chunk: Number(header.chunk_index),
frame_type: decoded.frame_type,
frames: decoded.frames,
};
} else if (header.content_type === RAW_RGB_CONTENT_TYPE) {
@@ -148,10 +161,7 @@ self.onmessage = async (event) => {
return;
}
const result = await decode(message.header, message.payload);
self.postMessage(
{ type: "decoded", ...result },
result.frames,
);
self.postMessage({ type: "decoded", ...result }, result.frames);
} catch (error) {
self.postMessage({
type: "error",
@@ -4,7 +4,7 @@
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1" />
<title>sglang-diffusion Realtime Studio</title>
<link rel="stylesheet" href="./styles.css?v=realtime-sr-v38" />
<link rel="stylesheet" href="./styles.css?v=realtime-record-v49" />
</head>
<body>
<main class="shell">
@@ -60,6 +60,18 @@
</select>
</label>
</div>
<label>SR model
<select id="upscalingModel">
<option value="">Quality x2</option>
<option
value="https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth"
selected
>
Fast general
</option>
<option value="/scratch/realesr-animevideov3.pth">Fast anime</option>
</select>
</label>
<label class="toggle-row"><input id="frameInterpolation" type="checkbox" />Smooth 2x frames</label>
<label class="toggle-row"><input id="continuous" type="checkbox" checked />Continuous session</label>
<div class="actions">
@@ -75,6 +87,11 @@
<span id="statusDot" class="dot"></span>
<span id="statusText">Idle</span>
<span id="chunkText">chunk -</span>
<button id="recordBtn" class="record-button" type="button" aria-pressed="false" title="Record preview">
<span class="record-button-icon" aria-hidden="true"></span>
<span id="recordLabel">Record</span>
<span id="recordDuration" class="record-button-duration">00:00</span>
</button>
<span class="topbar-spacer"></span>
<label class="preview-scale-control">Preview
<input id="previewScale" type="range" min="80" max="170" value="120" />
@@ -82,8 +99,8 @@
</label>
<span class="stage-stat">output <b id="outputSizeText">832x480</b></span>
<span class="stage-stat">render <b id="renderFps">0</b> fps</span>
<span class="stage-stat">theoretical <b id="theoreticalFpsText">-</b></span>
<span class="stage-stat">wait <b id="stageLatencyText">-</b></span>
<span class="stage-stat">source <b id="theoreticalFpsText">-</b></span>
<span class="stage-stat">buffer <b id="stageLatencyText">-</b></span>
</div>
<div class="preview-frame">
<canvas id="viewport" width="1280" height="720"></canvas>
@@ -145,6 +162,7 @@
</section>
</section>
</main>
<script src="./app.js?v=realtime-sr-v38"></script>
</body>
<script src="./playback_controller.js?v=realtime-playback-v13"></script>
<script src="./app.js?v=realtime-record-v72"></script>
</body>
</html>
@@ -0,0 +1,517 @@
(function attachRealtimePlaybackController(global) {
const DEFAULT_CONFIG = {
targetFps: 25,
minSourceFps: 1,
serverFpsAlphaUp: 0.28,
serverFpsAlphaDown: 0.2,
deliveryFpsAlphaUp: 0.08,
deliveryFpsAlphaDown: 0.55,
targetLeadChunkRatio: 1.5,
minTargetLeadMs: 1500,
maxTargetLeadMs: 2600,
maxLeadExtraChunkRatio: 8.0,
startLeadChunkRatio: 1.85,
minStartLeadMs: 1700,
resumeLeadChunkRatio: 2.5,
minResumeLeadMs: 1000,
maxResumeLeadMs: 1800,
rebufferLeadBoostMs: 250,
rebufferLeadBoostDecayMsPerSecond: 120,
deliveryLeadBoostDecayMsPerSecond: 80,
maxDeliveryLeadBoostMs: 2000,
deliveryStallExpectedMultiplier: 1.25,
receiveStallPlaybackRateMin: 0.65,
receiveStallPlaybackRateSlewPerSecond: 0.5,
lowWaterRatio: 0.4,
playbackRateGain: 0.14,
playbackRateMin: 0.92,
playbackRateMax: 1.08,
emergencyPlaybackRateMin: 0.9,
emergencyPlaybackRateMax: 1.12,
playbackRateSlewPerSecond: 0.08,
eventCutoverMaxMs: 420,
eventCutoverMaxFrames: 10,
settleEventCutoverMaxMs: 720,
settleEventCutoverMaxFrames: 18,
startupWarmupMinMs: 1500,
startupWarmupExpectedMultiplier: 3,
};
function clamp(value, min, max) {
return Math.min(max, Math.max(min, value));
}
function finitePositive(value) {
return Number.isFinite(value) && value > 0;
}
class RealtimePlaybackController {
constructor(config = {}) {
this.config = { ...DEFAULT_CONFIG, ...config };
this.reset({ targetFps: this.config.targetFps });
}
reset({ targetFps } = {}) {
this.targetFps = Math.max(1, Number(targetFps || this.config.targetFps));
this.sourceFps = this.targetFps;
this.serverFps = this.targetFps;
this.deliveryFps = this.targetFps;
this.hasServerSample = false;
this.hasDeliverySample = false;
this.latestChunkDurationMs = 1000 / this.targetFps;
this.latestChunkFrames = 1;
this.playbackRate = 1;
this.renderFps = this.targetFps;
this.queue = [];
this.lastDrawAt = 0;
this.lastRateUpdateAt = 0;
this.renderedFrames = 0;
this.droppedFrames = 0;
this.buffering = true;
this.pendingEventId = 0;
this.pendingEventSentAt = 0;
this.pendingEventCutoverMode = "motion";
this.lastDropReason = "";
this.lastDropAt = 0;
this.lastDropCount = 0;
this.rebufferLeadBoostMs = 0;
this.deliveryLeadBoostMs = 0;
this.chunkReceives = new Map();
this.serverStatChunks = new Set();
this.lastFinalReceiveAt = 0;
this.receiveStalled = false;
}
setTargetFps(targetFps) {
const nextTargetFps = Math.max(1, Number(targetFps || this.config.targetFps));
this.targetFps = nextTargetFps;
if (!this.hasServerSample && !this.hasDeliverySample) {
this.serverFps = nextTargetFps;
this.deliveryFps = nextTargetFps;
this.sourceFps = nextTargetFps;
this.renderFps = nextTargetFps;
} else {
this.serverFps = clamp(this.serverFps, this.config.minSourceFps, nextTargetFps);
this.deliveryFps = clamp(this.deliveryFps, this.config.minSourceFps, nextTargetFps);
this.sourceFps = clamp(this.sourceFps, this.config.minSourceFps, nextTargetFps);
this.renderFps = this.sourceFps * this.playbackRate;
}
this.latestChunkDurationMs = Math.max(this.latestChunkDurationMs, 1000 / this.targetFps);
}
clear() {
const frames = this.queue.splice(0);
this.lastDrawAt = 0;
this.buffering = true;
return frames;
}
noteInputEvent(eventId, now, { cutoverMode = "motion" } = {}) {
this.pendingEventId = Number(eventId || 0);
this.pendingEventSentAt = Number(now || 0);
this.pendingEventCutoverMode = cutoverMode;
}
observeServerStats(stats, now) {
const chunkIndex = Number(stats.chunk_index || 0);
const numFrames = Number(stats.num_frames || 0);
const chunkTotalMs = Number(stats.chunk_total_ms || 0);
if (numFrames > 0 && chunkTotalMs > 0) {
this.serverStatChunks.add(chunkIndex);
if (this.serverStatChunks.size > 128) {
this.serverStatChunks.delete(this.serverStatChunks.values().next().value);
}
const expectedMs = numFrames / Math.max(1, this.targetFps) * 1000;
const isStartupWarmup =
chunkIndex === 0 &&
chunkTotalMs > Math.max(
this.config.startupWarmupMinMs,
expectedMs * this.config.startupWarmupExpectedMultiplier,
);
if (isStartupWarmup) return this.snapshot();
this.#observeFpsSample("server", {
fps: numFrames / (chunkTotalMs / 1000),
frameCount: numFrames,
durationMs: chunkTotalMs,
now,
});
}
return this.snapshot();
}
enqueueDecodedFrames(header, frames, now) {
const chunkIndex = Number(header.chunk_index || 0);
const eventId = Number(header.event_id || 0);
const receivedAt = Number(header.__received_at || now);
const preparedFrames = frames.map((frame) => ({
...frame,
chunk: Number(frame.chunk ?? chunkIndex),
chunkIndex,
eventId,
}));
const droppedFrames = [];
let cutover = null;
if (this.pendingEventId && eventId >= this.pendingEventId) {
const oldEventFrameCount = this.#oldEventFrameCount(eventId);
const graceFrames = this.#eventGraceFrames();
const dropCount = Math.max(0, oldEventFrameCount - graceFrames);
if (dropCount > 0) {
droppedFrames.push(...this.queue.splice(graceFrames, dropCount));
this.#recordDrop(dropCount, "event cutover", now);
}
cutover = {
eventId,
latencyMs: this.pendingEventSentAt ? now - this.pendingEventSentAt : 0,
};
this.pendingEventId = 0;
this.pendingEventSentAt = 0;
this.pendingEventCutoverMode = "motion";
}
this.queue.push(...preparedFrames);
this.#observeChunkArrival(header, preparedFrames.length, receivedAt, now);
droppedFrames.push(...this.#trimBacklog(now));
return { droppedFrames, cutover, snapshot: this.snapshot() };
}
render(now, { hasPendingInput = true } = {}) {
this.#decayRebufferBoost(now);
this.#updateReceiveStallGuard(now);
const droppedFrames = this.#trimBacklog(now);
if (!this.queue.length) {
if (this.renderedFrames && hasPendingInput && !this.buffering) {
this.buffering = true;
this.rebufferLeadBoostMs = Math.max(
this.rebufferLeadBoostMs,
this.config.rebufferLeadBoostMs,
);
}
return { action: "hold", droppedFrames, snapshot: this.snapshot() };
}
const bufferMs = this.bufferDurationMs;
if (
hasPendingInput &&
this.receiveStalled &&
this.renderedFrames &&
bufferMs < this.targetLeadMs
) {
this.buffering = true;
this.lastDrawAt = 0;
return { action: "hold", droppedFrames, snapshot: this.snapshot() };
}
if (
hasPendingInput &&
this.buffering &&
bufferMs < (this.renderedFrames ? this.#resumeLeadMs() : this.#startLeadMs())
) {
this.buffering = true;
this.lastDrawAt = 0;
return { action: "hold", droppedFrames, snapshot: this.snapshot() };
}
if (this.buffering) {
this.buffering = false;
this.lastDrawAt = 0;
}
this.#updatePlaybackRate(now);
const targetMs = 1000 / Math.max(1, this.renderFps);
const elapsedMs = this.lastDrawAt ? now - this.lastDrawAt : targetMs;
if (elapsedMs < targetMs) {
return { action: "wait", droppedFrames, snapshot: this.snapshot() };
}
const frame = this.queue.shift();
this.renderedFrames += 1;
this.lastDrawAt = !this.lastDrawAt || elapsedMs > targetMs * 4
? now
: now - (elapsedMs % targetMs);
return { action: "draw", frame, droppedFrames, snapshot: this.snapshot() };
}
get queuedFrames() {
return this.queue.length;
}
get bufferDurationMs() {
return this.queue.length / Math.max(1, this.sourceFps) * 1000;
}
get targetLeadMs() {
const base = clamp(
this.latestChunkDurationMs * this.config.targetLeadChunkRatio,
this.config.minTargetLeadMs,
this.config.maxTargetLeadMs,
);
return clamp(
base + this.rebufferLeadBoostMs + this.deliveryLeadBoostMs,
this.config.minTargetLeadMs,
this.config.maxTargetLeadMs +
this.config.rebufferLeadBoostMs +
this.config.maxDeliveryLeadBoostMs,
);
}
get maxLeadMs() {
return this.targetLeadMs + this.latestChunkDurationMs * this.config.maxLeadExtraChunkRatio;
}
snapshot() {
return {
queueFrames: this.queue.length,
bufferMs: this.bufferDurationMs,
targetLeadMs: this.targetLeadMs,
maxLeadMs: this.maxLeadMs,
sourceFps: this.sourceFps,
serverFps: this.serverFps,
deliveryFps: this.deliveryFps,
targetFps: this.targetFps,
renderFps: this.renderFps,
playbackRate: this.playbackRate,
droppedFrames: this.droppedFrames,
lastDropAt: this.lastDropAt,
lastDropCount: this.lastDropCount,
buffering: this.buffering,
lastDropReason: this.lastDropReason,
};
}
#observeFpsSample(kind, { fps, frameCount, durationMs, now }) {
if (!finitePositive(fps)) return;
const cappedFps = clamp(fps, this.config.minSourceFps, this.targetFps);
const isDelivery = kind === "delivery";
const currentFps = isDelivery ? this.deliveryFps : this.serverFps;
const hasSample = isDelivery ? this.hasDeliverySample : this.hasServerSample;
let nextFps;
if (!hasSample) {
nextFps = cappedFps;
} else {
const alpha = cappedFps > currentFps
? (isDelivery ? this.config.deliveryFpsAlphaUp : this.config.serverFpsAlphaUp)
: (isDelivery ? this.config.deliveryFpsAlphaDown : this.config.serverFpsAlphaDown);
nextFps = currentFps * (1 - alpha) + cappedFps * alpha;
}
if (isDelivery) {
this.deliveryFps = nextFps;
this.hasDeliverySample = true;
this.#observeDeliveryJitter(frameCount, durationMs);
} else {
this.serverFps = nextFps;
this.hasServerSample = true;
}
const effectiveFps = this.hasServerSample
? this.serverFps
: (this.hasDeliverySample ? this.deliveryFps : this.targetFps);
this.sourceFps = clamp(effectiveFps, this.config.minSourceFps, this.targetFps);
if (!isDelivery || !this.hasServerSample) {
this.latestChunkFrames = Math.max(1, Number(frameCount || this.latestChunkFrames));
this.latestChunkDurationMs = clamp(
Number(durationMs || (this.latestChunkFrames / Math.max(1, this.sourceFps) * 1000)),
1000 / Math.max(1, this.targetFps),
2500,
);
}
this.#updatePlaybackRate(now);
}
#observeDeliveryJitter(frameCount, durationMs) {
if (!this.hasServerSample || !finitePositive(durationMs)) return;
const expectedMs = Number(frameCount || 0) / Math.max(1, this.serverFps) * 1000;
if (expectedMs <= 0) return;
if (durationMs <= expectedMs * this.config.deliveryStallExpectedMultiplier) return;
const boostMs = clamp(
durationMs - expectedMs,
0,
this.config.maxDeliveryLeadBoostMs,
);
this.deliveryLeadBoostMs = Math.max(this.deliveryLeadBoostMs, boostMs);
}
#updateReceiveStallGuard(now) {
this.receiveStalled = false;
if (!this.lastFinalReceiveAt || !this.hasServerSample) return;
const elapsedMs = now - this.lastFinalReceiveAt;
const expectedMs = Math.max(
this.latestChunkDurationMs,
this.latestChunkFrames / Math.max(1, this.serverFps) * 1000,
);
if (elapsedMs <= expectedMs * this.config.deliveryStallExpectedMultiplier) return;
this.receiveStalled = true;
this.deliveryLeadBoostMs = Math.max(
this.deliveryLeadBoostMs,
clamp(elapsedMs - expectedMs, 0, this.config.maxDeliveryLeadBoostMs),
);
}
#observeChunkArrival(header, frameCount, receivedAt, now) {
const chunkIndex = Number(header.chunk_index || 0);
const state = this.chunkReceives.get(chunkIndex) || {
firstAt: receivedAt,
frames: 0,
};
state.frames += Number(frameCount || 0);
state.lastAt = receivedAt;
this.chunkReceives.set(chunkIndex, state);
const frameBatchIndex = Number(header.frame_batch_index || 0);
const numFrameBatches = Number(header.num_frame_batches || 1);
const isFinalFrameBatch =
Boolean(header.is_final_frame_batch) ||
frameBatchIndex + 1 >= numFrameBatches;
if (!isFinalFrameBatch) return;
const durationMs = this.lastFinalReceiveAt
? receivedAt - this.lastFinalReceiveAt
: 0;
this.lastFinalReceiveAt = receivedAt;
if (state.frames > 0 && durationMs > 0) {
this.#observeFpsSample("delivery", {
fps: state.frames / (durationMs / 1000),
frameCount: state.frames,
durationMs,
now,
});
}
this.chunkReceives.delete(chunkIndex);
}
#updatePlaybackRate(now) {
const bufferMs = this.bufferDurationMs;
const targetLeadMs = Math.max(1, this.targetLeadMs);
const error = (bufferMs - targetLeadMs) / targetLeadMs;
const emergency =
bufferMs > this.maxLeadMs ||
bufferMs < targetLeadMs * this.config.lowWaterRatio ||
(this.receiveStalled && bufferMs < targetLeadMs);
const minRate = emergency
? (
this.receiveStalled
? this.config.receiveStallPlaybackRateMin
: this.config.emergencyPlaybackRateMin
)
: this.config.playbackRateMin;
const maxRate = this.receiveStalled && bufferMs < targetLeadMs
? 1
: emergency
? this.config.emergencyPlaybackRateMax
: this.config.playbackRateMax;
const desiredRate = clamp(
1 + error * this.config.playbackRateGain,
minRate,
maxRate,
);
if (!this.lastRateUpdateAt) {
this.playbackRate = desiredRate;
} else {
const dtSeconds = Math.max(0.001, (now - this.lastRateUpdateAt) / 1000);
const slewPerSecond = this.receiveStalled
? this.config.receiveStallPlaybackRateSlewPerSecond
: this.config.playbackRateSlewPerSecond;
const maxDelta = slewPerSecond * dtSeconds;
this.playbackRate = clamp(
desiredRate,
this.playbackRate - maxDelta,
this.playbackRate + maxDelta,
);
}
this.lastRateUpdateAt = now;
this.renderFps = clamp(
this.sourceFps * this.playbackRate,
this.config.minSourceFps,
this.targetFps * this.config.emergencyPlaybackRateMax,
);
}
#trimBacklog(now) {
const droppedFrames = [];
while (this.queue.length && this.bufferDurationMs > this.maxLeadMs) {
const firstChunk = this.queue[0].chunkIndex;
let dropCount = 0;
while (
dropCount < this.queue.length &&
this.queue[dropCount].chunkIndex === firstChunk
) {
dropCount += 1;
}
if (!dropCount || dropCount >= this.queue.length) break;
droppedFrames.push(...this.queue.splice(0, dropCount));
this.#recordDrop(dropCount, "backlog", now);
}
return droppedFrames;
}
#oldEventFrameCount(nextEventId) {
let count = 0;
while (count < this.queue.length && this.queue[count].eventId < nextEventId) {
count += 1;
}
return count;
}
#eventGraceFrames() {
const byTime = Math.max(
1,
Math.round(this.sourceFps * this.#eventCutoverMaxMs() / 1000),
);
const byChunkRatio = this.pendingEventCutoverMode === "settle" ? 1.5 : 0.85;
const byChunk = Math.max(1, Math.round(this.latestChunkFrames * byChunkRatio));
return Math.min(this.#eventCutoverMaxFrames(), byTime, byChunk);
}
#eventCutoverMaxMs() {
return this.pendingEventCutoverMode === "settle"
? this.config.settleEventCutoverMaxMs
: this.config.eventCutoverMaxMs;
}
#eventCutoverMaxFrames() {
return this.pendingEventCutoverMode === "settle"
? this.config.settleEventCutoverMaxFrames
: this.config.eventCutoverMaxFrames;
}
#startLeadMs() {
return Math.max(
this.config.minStartLeadMs,
this.latestChunkDurationMs * this.config.startLeadChunkRatio,
this.targetLeadMs,
);
}
#resumeLeadMs() {
return clamp(
this.latestChunkDurationMs * this.config.resumeLeadChunkRatio,
this.config.minResumeLeadMs,
this.config.maxResumeLeadMs,
);
}
#decayRebufferBoost(now) {
if ((!this.rebufferLeadBoostMs && !this.deliveryLeadBoostMs) || !this.lastRateUpdateAt) return;
const dtSeconds = Math.max(0, (now - this.lastRateUpdateAt) / 1000);
this.rebufferLeadBoostMs = Math.max(
0,
this.rebufferLeadBoostMs - dtSeconds * this.config.rebufferLeadBoostDecayMsPerSecond,
);
this.deliveryLeadBoostMs = Math.max(
0,
this.deliveryLeadBoostMs - dtSeconds * this.config.deliveryLeadBoostDecayMsPerSecond,
);
}
#recordDrop(count, reason, now) {
this.droppedFrames += count;
this.lastDropAt = Number(now || 0);
this.lastDropCount = count;
this.lastDropReason = reason;
}
}
global.RealtimePlaybackController = RealtimePlaybackController;
if (typeof module !== "undefined" && module.exports) {
module.exports = { RealtimePlaybackController };
}
})(typeof globalThis !== "undefined" ? globalThis : window);
@@ -0,0 +1,116 @@
const assert = require("node:assert/strict");
const { RealtimePlaybackController } = require("./playback_controller.js");
function frames(count, chunk) {
return Array.from({ length: count }, (_, index) => ({
image: { close() {} },
chunk,
index,
}));
}
function enqueueChunk(controller, {
chunk,
eventId = 0,
frameCount = 12,
durationMs = 480,
now,
}) {
controller.observeServerStats({
chunk_index: chunk,
num_frames: frameCount,
chunk_total_ms: durationMs,
}, now);
return controller.enqueueDecodedFrames({
chunk_index: chunk,
event_id: eventId,
num_frames: frameCount,
__received_at: now,
is_final_frame_batch: true,
}, frames(frameCount, chunk), now);
}
function renderFor(controller, startMs, durationMs) {
let rendered = 0;
for (let now = startMs; now <= startMs + durationMs; now += 16.67) {
const decision = controller.render(now, { hasPendingInput: true });
if (decision.action === "draw") rendered += 1;
}
return rendered;
}
function stableSourceDoesNotDrop() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
let now = 0;
for (let chunk = 0; chunk < 8; chunk += 1) {
now += 480;
enqueueChunk(controller, { chunk, now });
renderFor(controller, now, 480);
}
const snapshot = controller.snapshot();
assert.equal(snapshot.droppedFrames, 0);
assert.ok(snapshot.sourceFps > 24 && snapshot.sourceFps <= 25);
}
function slowServerCapsRenderFps() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
let now = 0;
for (let chunk = 0; chunk < 8; chunk += 1) {
now += 1360;
enqueueChunk(controller, { chunk, durationMs: 1360, now });
renderFor(controller, now, 1360);
}
const snapshot = controller.snapshot();
assert.ok(snapshot.sourceFps > 8 && snapshot.sourceFps < 10);
assert.ok(snapshot.renderFps <= 10);
}
function backlogDropsContiguousOldFrames() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
let now = 100;
for (let chunk = 0; chunk < 13; chunk += 1) {
enqueueChunk(controller, { chunk, now, durationMs: 480 });
now += 20;
}
const snapshot = controller.snapshot();
assert.ok(snapshot.droppedFrames > 0);
assert.equal(snapshot.lastDropReason, "backlog");
}
function eventCutoverKeepsShortGrace() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
enqueueChunk(controller, { chunk: 1, frameCount: 24, durationMs: 960, now: 1000 });
controller.noteInputEvent(5, 1050);
const result = enqueueChunk(controller, {
chunk: 2,
eventId: 5,
frameCount: 12,
durationMs: 480,
now: 1150,
});
assert.ok(result.cutover);
assert.ok(result.droppedFrames.length >= 14);
assert.equal(controller.queue[0].chunk, 1);
assert.equal(controller.queue[0].index, 0);
}
function settleEventCutoverKeepsWiderGrace() {
const controller = new RealtimePlaybackController({ targetFps: 25 });
enqueueChunk(controller, { chunk: 1, frameCount: 24, durationMs: 960, now: 1000 });
controller.noteInputEvent(5, 1050, { cutoverMode: "settle" });
const result = enqueueChunk(controller, {
chunk: 2,
eventId: 5,
frameCount: 12,
durationMs: 480,
now: 1150,
});
assert.ok(result.cutover);
assert.ok(result.droppedFrames.length <= 12);
}
stableSourceDoesNotDrop();
slowServerCapsRenderFps();
backlogDropsContiguousOldFrames();
eventCutoverKeepsShortGrace();
settleEventCutoverKeepsWiderGrace();
@@ -17,6 +17,7 @@
body {
margin: 0;
overflow-x: hidden;
min-height: 100vh;
background:
linear-gradient(90deg, rgba(23, 26, 22, 0.035) 1px, transparent 1px),
@@ -32,8 +33,10 @@ button:disabled { cursor: wait; opacity: 0.64; transform: none; }
.shell {
display: grid;
grid-template-columns: minmax(260px, 320px) minmax(560px, 1fr);
grid-template-columns: minmax(260px, 320px) minmax(0, 1fr);
gap: 18px;
width: 100%;
max-width: 100vw;
min-height: 100vh;
padding: 18px;
}
@@ -219,6 +222,7 @@ button:focus-visible {
display: grid;
gap: 18px;
min-width: 0;
max-width: 100%;
}
.stage {
@@ -227,6 +231,7 @@ button:focus-visible {
grid-template-rows: auto auto auto auto auto;
align-self: start;
justify-self: center;
min-width: 0;
width: 100%;
max-width: min(1500px, 100%);
overflow: hidden;
@@ -273,12 +278,15 @@ button:focus-visible {
.topbar, .timeline {
display: flex;
align-items: center;
gap: 14px;
gap: 10px;
min-width: 0;
height: 44px;
padding: 0 14px;
color: #e8eadf;
background: rgba(17, 20, 15, 0.88);
font-size: 12px;
font-variant-numeric: tabular-nums;
white-space: nowrap;
}
.topbar > * {
@@ -287,11 +295,76 @@ button:focus-visible {
}
.topbar-spacer { flex: 1; }
.record-button {
display: inline-flex;
align-items: center;
justify-content: center;
gap: 6px;
flex: 0 0 118px;
width: 118px;
min-height: 28px;
height: 28px;
padding: 0 9px;
border-color: rgba(232, 234, 223, 0.22);
background: rgba(238, 241, 236, 0.08);
color: #e8eadf;
font-size: 11px;
font-variant-numeric: tabular-nums;
}
.record-button:hover:not(:disabled) {
border-color: rgba(232, 234, 223, 0.44);
background: rgba(238, 241, 236, 0.14);
box-shadow: none;
transform: none;
}
.record-button:active:not(:disabled),
.record-button.is-pressed:not(:disabled),
.record-button:focus-visible {
transform: none;
}
.record-button.is-recording {
border-color: color-mix(in oklch, var(--accent), white 18%);
background: var(--accent);
color: #fffdf7;
}
.record-button.is-saving {
cursor: wait;
opacity: 0.76;
}
#recordLabel {
flex: 0 0 36px;
text-align: left;
}
.record-button-icon {
flex: 0 0 9px;
width: 9px;
height: 9px;
min-width: 9px;
border-radius: 50%;
background: var(--accent);
box-shadow: 0 0 0 3px rgba(185, 84, 60, 0.16);
}
.record-button.is-recording .record-button-icon {
border-radius: 2px;
background: #fffdf7;
box-shadow: none;
}
.record-button-duration {
display: inline-block;
flex: 0 0 34px;
min-width: 34px;
text-align: right;
color: rgba(232, 234, 223, 0.7);
}
.record-button.is-recording .record-button-duration {
color: rgba(255, 253, 247, 0.86);
}
.preview-scale-control {
display: inline-flex;
align-items: center;
gap: 8px;
min-width: 190px;
flex: 0 0 170px;
min-width: 170px;
margin: 0;
color: rgba(232, 234, 223, 0.72);
font-size: 11px;
@@ -324,13 +397,36 @@ button:focus-visible {
display: inline-flex;
align-items: center;
gap: 5px;
min-width: max-content;
flex: 0 1 auto;
min-width: 0;
color: rgba(232, 234, 223, 0.72);
line-height: 1;
}
.stage-stat b {
display: inline-block;
color: #fffdf7;
font-weight: 650;
font-variant-numeric: tabular-nums;
}
#outputSizeText { min-width: 206px; }
#renderFps { min-width: 2ch; text-align: right; }
#theoreticalFpsText { min-width: 116px; }
#stageLatencyText { min-width: 120px; }
@media (max-width: 1180px) {
.topbar {
flex-wrap: wrap;
height: auto;
min-height: 44px;
padding: 8px 14px;
row-gap: 7px;
}
.topbar-spacer { display: none; }
.preview-scale-control { flex-basis: 170px; min-width: 170px; }
#outputSizeText { min-width: 156px; }
#theoreticalFpsText { min-width: 100px; }
#stageLatencyText { min-width: 108px; }
}
.timeline { justify-content: flex-end; border-top: 1px solid rgba(232, 234, 223, 0.12); }
@@ -43,8 +43,10 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
log_batch_completion,
log_generation_timer,
)
from sglang.multimodal_gen.runtime.utils.trace_wrapper import trace_req
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
from sglang.multimodal_gen.runtime.utils.trace_wrapper import (
init_diffusion_tracing,
trace_req,
)
logger = init_logger(__name__)
@@ -123,9 +125,7 @@ class DiffGenerator:
instance = cls(
server_args=server_args,
)
if server_args.enable_trace:
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info("DiffGenerator")
init_diffusion_tracing(server_args, "DiffGenerator")
logger.info(f"Local mode: {local_mode}")
if local_mode:
@@ -151,6 +151,8 @@ class RealtimeVideoGenerationsRequest(VideoGenerationsRequest):
num_profiled_timesteps: Optional[int] = None
profile_all_stages: Optional[bool] = False
realtime_output_format: Optional[Literal["raw", "webp", "jpeg"]] = None
realtime_preview_max_width: Optional[int] = None
realtime_output_pacing: Optional[bool] = False
realtime_causal_sink_size: Optional[int] = None
realtime_causal_kv_cache_num_frames: Optional[int] = None
@@ -430,6 +430,10 @@ class LingBotWorldRealtimeAdapter(RealtimeModelAdapter):
batch.realtime_event_id = self._state(session).latest_sampled_event_id
if session.request is not None:
batch.realtime_output_format = session.request.realtime_output_format
batch.realtime_preview_max_width = (
session.request.realtime_preview_max_width
)
batch.realtime_output_pacing = bool(session.request.realtime_output_pacing)
batch.realtime_causal_sink_size = session.request.realtime_causal_sink_size
batch.realtime_causal_kv_cache_num_frames = (
session.request.realtime_causal_kv_cache_num_frames
@@ -36,6 +36,8 @@ class GenerateSession:
self.realtime_session = RealtimeSession()
self.adapter: RealtimeModelAdapter | None = None
self.adapter_state: Any = None
self.output_pace_next_send_at: float | None = None
self.output_pace_last_event_id: int | None = None
def set_adapter(self, adapter: RealtimeModelAdapter):
self.adapter = adapter
@@ -53,6 +55,8 @@ class GenerateSession:
self.current_chunk = None
self.adapter = None
self.adapter_state = None
self.output_pace_next_send_at = None
self.output_pace_last_event_id = None
self.realtime_session.dispose()
def new_chunk(self) -> RealtimeChunkContext:
@@ -48,6 +48,7 @@ class RealtimeFrameBatchHeader(TypedDict, total=False):
raw_size: int
encoding: str
delta_reference: str
payload_lengths: list[int]
event_id: int
frame_batch_index: int
num_frame_batches: int
@@ -64,6 +65,7 @@ class RealtimeFrameSendStats(TypedDict):
raw_payload_build_ms: float
raw_write_ms: float
ws_write_ms: float
pace_wait_ms: float
raw_bytes: int
ws_payload_bytes: int
num_frames: int
@@ -79,6 +81,7 @@ def empty_frame_send_stats(content_type: str = "") -> RealtimeFrameSendStats:
"raw_payload_build_ms": 0.0,
"raw_write_ms": 0.0,
"ws_write_ms": 0.0,
"pace_wait_ms": 0.0,
"raw_bytes": 0,
"ws_payload_bytes": 0,
"num_frames": 0,
@@ -135,7 +138,7 @@ ENCODED_PREVIEW_FORMATS = {"webp", "jpeg"}
class _TransportPayload:
content_type: str
payload: bytes
metadata: dict[str, int | str | bool]
metadata: dict[str, int | str | bool | list[int]]
last_raw_rgb_frame: bytes | None = None
last_event_id: int | None = None
@@ -155,9 +158,14 @@ def _encode_rgb_frame_to_webp(
width: int,
height: int,
quality: int,
preview_max_width: int | None,
) -> bytes:
buffer = io.BytesIO()
Image.frombytes("RGB", (width, height), frame).save(
image = _resize_preview_image(
Image.frombytes("RGB", (width, height), frame),
preview_max_width=preview_max_width,
)
image.save(
buffer,
format="WEBP",
quality=quality,
@@ -172,9 +180,14 @@ def _encode_rgb_frame_to_jpeg(
width: int,
height: int,
quality: int,
preview_max_width: int | None,
) -> bytes:
buffer = io.BytesIO()
Image.frombytes("RGB", (width, height), frame).save(
image = _resize_preview_image(
Image.frombytes("RGB", (width, height), frame),
preview_max_width=preview_max_width,
)
image.save(
buffer,
format="JPEG",
quality=quality,
@@ -183,6 +196,39 @@ def _encode_rgb_frame_to_jpeg(
return buffer.getvalue()
def _preview_dimensions(
*,
width: int,
height: int,
preview_max_width: int | None,
) -> tuple[int, int]:
if (
preview_max_width is None
or preview_max_width <= 0
or width <= preview_max_width
):
return width, height
preview_width = int(preview_max_width)
preview_height = max(1, round(height * preview_width / width))
return preview_width, preview_height
def _resize_preview_image(
image: Image.Image,
*,
preview_max_width: int | None,
) -> Image.Image:
width, height = image.size
preview_width, preview_height = _preview_dimensions(
width=width,
height=height,
preview_max_width=preview_max_width,
)
if (preview_width, preview_height) == image.size:
return image
return image.resize((preview_width, preview_height), Image.Resampling.BICUBIC)
def _pack_frame_batch_message(
header: RealtimeFrameBatchHeader,
payload: bytes,
@@ -202,11 +248,12 @@ def _build_transport_payload(
metadata: dict[str, int | str],
output_format: str | None,
transport_quality: int | None,
preview_max_width: int | None,
reference_frame: bytes | None,
event_id: int | None,
) -> _TransportPayload:
payload_content_type = content_type
payload_metadata: dict[str, int | str | bool] = {}
payload_metadata: dict[str, int | str | bool | list[int]] = {}
raw_payload = b""
if (
@@ -215,24 +262,45 @@ def _build_transport_payload(
and transport_frames
):
if output_format == "webp":
raw_payload = _encode_rgb_frame_to_webp(
transport_frames[0],
width=int(metadata["width"]),
height=int(metadata["height"]),
quality=int(transport_quality or WEBP_DEFAULT_QUALITY),
)
encoded_frames = [
_encode_rgb_frame_to_webp(
frame,
width=int(metadata["width"]),
height=int(metadata["height"]),
quality=int(transport_quality or WEBP_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
]
payload_content_type = WEBP_FRAME_CONTENT_TYPE
else:
raw_payload = _encode_rgb_frame_to_jpeg(
transport_frames[0],
width=int(metadata["width"]),
height=int(metadata["height"]),
quality=int(transport_quality or JPEG_DEFAULT_QUALITY),
)
encoded_frames = [
_encode_rgb_frame_to_jpeg(
frame,
width=int(metadata["width"]),
height=int(metadata["height"]),
quality=int(transport_quality or JPEG_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
]
payload_content_type = JPEG_FRAME_CONTENT_TYPE
raw_payload = b"".join(encoded_frames)
preview_width, preview_height = _preview_dimensions(
width=int(metadata["width"]),
height=int(metadata["height"]),
preview_max_width=preview_max_width,
)
payload_metadata = {
"format": output_format,
"encoding": output_format,
"source_width": int(metadata["width"]),
"source_height": int(metadata["height"]),
"preview_width": preview_width,
"preview_height": preview_height,
"width": preview_width,
"height": preview_height,
"payload_lengths": [len(frame) for frame in encoded_frames],
}
elif (
output_format == RAW_LOSSLESS_OUTPUT_FORMAT
@@ -302,20 +370,18 @@ async def _build_encoded_preview_payloads(
metadata: dict[str, int | str],
output_format: str,
transport_quality: int | None,
preview_max_width: int | None,
event_id: int | None,
) -> list[_TransportPayload]:
return list(
await asyncio.gather(
*(
asyncio.to_thread(
_build_transport_payload,
_build_encoded_preview_payload(
transport_frames,
content_type=content_type,
metadata=metadata,
output_format=output_format,
transport_quality=transport_quality,
reference_frame=None,
event_id=event_id,
preview_max_width=preview_max_width,
)
for transport_frames in split_batches
)
@@ -323,6 +389,73 @@ async def _build_encoded_preview_payloads(
)
async def _build_encoded_preview_payload(
transport_frames: list[bytes],
*,
metadata: dict[str, int | str],
output_format: str,
transport_quality: int | None,
preview_max_width: int | None,
) -> _TransportPayload:
width = int(metadata["width"])
height = int(metadata["height"])
if output_format == "webp":
encoded_frames = list(
await asyncio.gather(
*(
asyncio.to_thread(
_encode_rgb_frame_to_webp,
frame,
width=width,
height=height,
quality=int(transport_quality or WEBP_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
)
)
)
payload_content_type = WEBP_FRAME_CONTENT_TYPE
else:
encoded_frames = list(
await asyncio.gather(
*(
asyncio.to_thread(
_encode_rgb_frame_to_jpeg,
frame,
width=width,
height=height,
quality=int(transport_quality or JPEG_DEFAULT_QUALITY),
preview_max_width=preview_max_width,
)
for frame in transport_frames
)
)
)
payload_content_type = JPEG_FRAME_CONTENT_TYPE
preview_width, preview_height = _preview_dimensions(
width=width,
height=height,
preview_max_width=preview_max_width,
)
return _TransportPayload(
content_type=payload_content_type,
payload=b"".join(encoded_frames),
metadata={
"format": output_format,
"encoding": output_format,
"source_width": width,
"source_height": height,
"preview_width": preview_width,
"preview_height": preview_height,
"width": preview_width,
"height": preview_height,
"payload_lengths": [len(frame) for frame in encoded_frames],
},
)
class RawRGBRealtimeOutputAdapter:
"""send raw RGB over WebSocket using lossless transport compression"""
@@ -354,6 +487,7 @@ class RawRGBRealtimeOutputAdapter:
else {}
)
output_format = getattr(batch, "realtime_output_format", None)
preview_max_width = getattr(batch, "realtime_preview_max_width", None)
stats = await self._send_frame_batches(
ws,
result.raw_frame_batches,
@@ -364,6 +498,7 @@ class RawRGBRealtimeOutputAdapter:
frame_metadata=frame_metadata,
output_format=output_format,
transport_quality=getattr(batch, "output_compression", None),
preview_max_width=preview_max_width,
)
stats["frame_shape"] = _frame_shape_from_metadata(frame_metadata)
return stats
@@ -380,13 +515,14 @@ class RawRGBRealtimeOutputAdapter:
frame_metadata: dict[str, int | str] | None = None,
output_format: str | None = None,
transport_quality: int | None = None,
preview_max_width: int | None = None,
) -> RealtimeFrameSendStats:
chunk_index = chunk_index_start
metadata = frame_metadata or {}
stats = empty_frame_send_stats(content_type)
for frames in frame_batches:
split_batches = (
[[frame] for frame in frames]
[frames]
if _is_encoded_preview_transport(
content_type=content_type,
output_format=output_format,
@@ -410,6 +546,7 @@ class RawRGBRealtimeOutputAdapter:
metadata=metadata,
output_format=output_format,
transport_quality=transport_quality,
preview_max_width=preview_max_width,
event_id=event_id,
)
stats["raw_payload_build_ms"] += timer.mark_ms()
@@ -434,6 +571,7 @@ class RawRGBRealtimeOutputAdapter:
metadata=metadata,
output_format=output_format,
transport_quality=transport_quality,
preview_max_width=preview_max_width,
reference_frame=reference_frame,
event_id=event_id,
)
@@ -444,6 +582,7 @@ class RawRGBRealtimeOutputAdapter:
metadata=metadata,
output_format=output_format,
transport_quality=transport_quality,
preview_max_width=preview_max_width,
reference_frame=reference_frame,
event_id=event_id,
)
@@ -41,7 +41,7 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
router = APIRouter(prefix="/v1/realtime_video", tags=["realtime"])
_ACTIVE_SESSION_IDS: set[str] = set()
_ACTIVE_SESSION_WAIT_SECONDS = 15.0
_ACTIVE_SESSION_WAIT_SECONDS = 1.0
_ACTIVE_SESSION_WAIT_INTERVAL_SECONDS = 0.1
@@ -73,6 +73,7 @@ def _log_realtime_chunk_timing(
"realtime chunk timing: session_id=%s request_id=%s "
"chunk_idx=%s event_id=%s condition_kinds=%s "
"request_prepare=%.2fms scheduler_forward=%.2fms "
"output_pace=%.2fms "
"header_pack=%.2fms "
"header_write=%.2fms raw_payload_build=%.2fms raw_write=%.2fms "
"ws_write=%.2fms chunk_total=%.2fms batches=%d frames=%d "
@@ -84,6 +85,7 @@ def _log_realtime_chunk_timing(
sorted(batch.condition_inputs) if batch.condition_inputs else [],
request_prepare_ms,
scheduler_forward_ms,
send_stats["pace_wait_ms"],
send_stats["header_pack_ms"],
send_stats["header_write_ms"],
send_stats["raw_payload_build_ms"],
@@ -119,6 +121,7 @@ async def _send_realtime_chunk_stats(
"event_id": getattr(batch, "realtime_event_id", None),
"request_prepare_ms": _transport_ms(request_prepare_ms),
"scheduler_forward_ms": _transport_ms(scheduler_forward_ms),
"pace_wait_ms": _transport_ms(send_stats["pace_wait_ms"]),
"header_write_ms": _transport_ms(send_stats["header_write_ms"]),
"raw_payload_build_ms": _transport_ms(
send_stats["raw_payload_build_ms"]
@@ -178,8 +181,8 @@ async def _generate_loop(ws: WebSocket, session: GenerateSession):
adapter.on_chunk_complete(session, result)
if pending_send_task is not None:
await pending_send_task
pending_send_task = asyncio.create_task(
_send_output_and_log(
if batch.realtime_output_pacing:
await _send_output_and_log(
ws,
session,
chunk,
@@ -189,7 +192,20 @@ async def _generate_loop(ws: WebSocket, session: GenerateSession):
scheduler_forward_ms,
chunk_started,
)
)
pending_send_task = None
else:
pending_send_task = asyncio.create_task(
_send_output_and_log(
ws,
session,
chunk,
batch,
result,
request_prepare_ms,
scheduler_forward_ms,
chunk_started,
)
)
except asyncio.CancelledError:
if pending_send_task is not None:
@@ -241,12 +257,14 @@ async def _send_output_and_log(
) -> RealtimeFrameSendStats:
if session.adapter is None:
raise ValueError("realtime adapter is not initialized")
pace_wait_ms = await _wait_for_realtime_output_slot(session, batch, result)
send_stats = await session.adapter.send_output(
ws,
session,
result,
batch,
)
send_stats["pace_wait_ms"] = pace_wait_ms
chunk_total_ms = (time.perf_counter() - chunk_started) * 1000
_log_realtime_chunk_timing(
session,
@@ -270,6 +288,54 @@ async def _send_output_and_log(
return send_stats
def _result_num_frames(result) -> int:
if result.raw_frame_batches is None:
return 0
return sum(len(frames) for frames in result.raw_frame_batches)
def _output_pacing_fps(batch: "Req") -> float:
fps = float(batch.fps or 0)
if batch.enable_frame_interpolation:
fps *= 2 ** int(batch.frame_interpolation_exp or 1)
return fps
async def _wait_for_realtime_output_slot(
session: GenerateSession,
batch: "Req",
result,
) -> float:
if not batch.realtime_output_pacing:
return 0.0
frame_count = _result_num_frames(result)
output_fps = _output_pacing_fps(batch)
if frame_count <= 0 or output_fps <= 0:
return 0.0
now = time.perf_counter()
next_send_at = session.output_pace_next_send_at
if next_send_at is None:
next_send_at = now
if (
batch.realtime_event_id is not None
and batch.realtime_event_id != session.output_pace_last_event_id
):
next_send_at = min(next_send_at, now)
session.output_pace_last_event_id = batch.realtime_event_id
wait_s = max(0.0, next_send_at - now)
if wait_s > 0:
await asyncio.sleep(wait_s)
send_started_at = time.perf_counter()
session.output_pace_next_send_at = (
max(next_send_at, send_started_at) + frame_count / output_fps
)
return wait_s * 1000
async def _await_realtime_task(task: asyncio.Task | None) -> None:
if task is None:
return
@@ -23,7 +23,7 @@ from sglang.multimodal_gen.runtime.server_args import (
)
from sglang.multimodal_gen.runtime.utils.common import is_port_available
from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
from sglang.multimodal_gen.runtime.utils.trace_wrapper import init_diffusion_tracing
def _find_available_port(
@@ -443,9 +443,7 @@ def _run_disagg_role_process(
def launch_http_server_only(server_args):
if server_args.enable_trace:
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info("DiffHTTPServer")
init_diffusion_tracing(server_args, "DiffHTTPServer")
# set for endpoints to access global_server_args
set_global_server_args(server_args)
@@ -72,8 +72,11 @@ from sglang.multimodal_gen.runtime.utils.realtime_video import (
RAW_RGB_CONTENT_TYPE,
build_raw_rgb_frame_batches,
)
from sglang.multimodal_gen.runtime.utils.trace_wrapper import DiffStage, trace_slice
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
from sglang.multimodal_gen.runtime.utils.trace_wrapper import (
DiffStage,
init_diffusion_tracing,
trace_slice,
)
from sglang.srt.utils.network import NetworkAddress
logger = init_logger(__name__)
@@ -989,9 +992,7 @@ def run_scheduler_process(
elif current_platform.is_musa():
set_musa_arch()
if server_args.enable_trace:
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info(f"DiffWorker_rank{rank}")
init_diffusion_tracing(server_args, f"DiffWorker_rank{rank}")
port_args = PortArgs.from_server_args(server_args)
@@ -211,6 +211,8 @@ class Req:
realtime_chunk_size: int | None = None
realtime_event_id: int | None = None
realtime_output_format: str | None = None
realtime_preview_max_width: int | None = None
realtime_output_pacing: bool = False
realtime_causal_sink_size: int | None = None
realtime_causal_kv_cache_num_frames: int | None = None
# return websocket-friendly raw RGB frame bytes instead of rwa tensors
@@ -12,7 +12,9 @@ The ImageUpscaler wrapper and integration code are original work.
import math
import os
import time
from hashlib import sha256
from typing import Optional
from urllib.parse import unquote, urlparse
import numpy as np
import torch
@@ -689,6 +691,7 @@ def _resolve_model_path(model_path: str) -> str:
Accepts:
- An existing local file path (pass-through).
- An http(s) URL to a .pth file, downloaded into the local cache.
- A HuggingFace ``repo_id`` → downloads the default weight file
(``RealESRGAN_x4.pth``).
- A HuggingFace ``repo_id:filename`` → downloads *filename* from *repo_id*,
@@ -702,6 +705,25 @@ def _resolve_model_path(model_path: str) -> str:
_RESOLVED_MODEL_PATH_CACHE[model_path] = model_path
return model_path
parsed_url = urlparse(model_path)
if parsed_url.scheme in ("http", "https"):
filename = (
os.path.basename(unquote(parsed_url.path)) or _DEFAULT_REALESRGAN_FILENAME
)
cache_dir = os.path.join(
os.path.expanduser("~"), ".cache", "sglang", "realesrgan"
)
os.makedirs(cache_dir, exist_ok=True)
cache_key = sha256(model_path.encode("utf-8")).hexdigest()[:12]
local_path = os.path.join(cache_dir, f"{cache_key}-{filename}")
if not os.path.isfile(local_path):
tmp_path = f"{local_path}.tmp"
logger.info("Downloading Real-ESRGAN weights from URL %s", model_path)
torch.hub.download_url_to_file(model_path, tmp_path, progress=False)
os.replace(tmp_path, local_path)
_RESOLVED_MODEL_PATH_CACHE[model_path] = local_path
return local_path
# Parse optional "repo_id:filename" syntax; fall back to default filename.
if ":" in model_path and not model_path.startswith("/"):
repo_id, filename = model_path.split(":", 1)
@@ -11,6 +11,7 @@ The FrameInterpolator wrapper and integration code are original work.
"""
import os
import time
from typing import Optional
import numpy as np
@@ -28,6 +29,7 @@ _DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite"
# Module-level cache: model_path -> Model instance
_MODEL_CACHE: dict[str, "Model"] = {}
_MAX_RIFE_BATCH_PAIRS = 16
# ---------------------------------------------------------------------------
@@ -392,6 +394,20 @@ class FrameInterpolator:
arr = t.squeeze(0).permute(1, 2, 0).clamp(0.0, 1.0).cpu().numpy()
return (arr * 255.0).astype(np.uint8)
@staticmethod
def _frames_to_tensor(
frames: list[np.ndarray], device: torch.device
) -> torch.Tensor:
t = torch.from_numpy(np.stack(frames, axis=0))
t = t.permute(0, 3, 1, 2).contiguous().float() / 255.0
return t.to(device, non_blocking=True)
@staticmethod
def _tensor_to_frames(t: torch.Tensor) -> list[np.ndarray]:
arr = t.permute(0, 2, 3, 1).clamp(0.0, 1.0).cpu().numpy()
arr = (arr * 255.0).astype(np.uint8)
return [arr[i] for i in range(arr.shape[0])]
def _make_inference(
self, model: Model, I0: torch.Tensor, I1: torch.Tensor, n: int, scale: float
) -> list[torch.Tensor]:
@@ -409,6 +425,30 @@ class FrameInterpolator:
+ self._make_inference(model, mid, I1, n // 2, scale)
)
def _interpolate_2x_batched(
self, model: Model, frames: list[np.ndarray], scale: float
) -> list[np.ndarray]:
device = model.device()
source = self._frames_to_tensor(frames, device)
intermediate_frames: list[np.ndarray] = []
with torch.inference_mode():
for start in range(0, len(frames) - 1, _MAX_RIFE_BATCH_PAIRS):
end = min(start + _MAX_RIFE_BATCH_PAIRS, len(frames) - 1)
mids = model.inference(
source[start:end],
source[start + 1 : end + 1],
scale=scale,
)
intermediate_frames.extend(self._tensor_to_frames(mids))
result: list[np.ndarray] = []
for i, mid in enumerate(intermediate_frames):
result.append(frames[i])
result.append(mid)
result.append(frames[-1])
return result
def interpolate(
self,
frames: list[np.ndarray],
@@ -436,6 +476,16 @@ class FrameInterpolator:
device = model.device()
n_intermediate = 2**exp // 2 # intermediates per adjacent pair
start_time = time.perf_counter()
if n_intermediate == 1:
result = self._interpolate_2x_batched(model, frames, scale)
logger.info(
"RIFE batched interpolation completed in %.3f seconds for %d frames",
time.perf_counter() - start_time,
len(frames),
)
return result, 2**exp
result: list[np.ndarray] = []
for i in range(len(frames) - 1):
@@ -452,6 +502,11 @@ class FrameInterpolator:
result.append(frames[-1])
multiplier = 2**exp
logger.info(
"RIFE interpolation completed in %.3f seconds for %d frames",
time.perf_counter() - start_time,
len(frames),
)
return result, multiplier
@@ -10,6 +10,8 @@ from __future__ import annotations
from contextlib import contextmanager
from dataclasses import dataclass
DIFFUSION_TRACE_MODULE = "diffusion"
@dataclass(frozen=True)
class DiffStageConfig:
@@ -26,6 +28,39 @@ class DiffStage:
GPU_FORWARD = DiffStageConfig("gpu_forward", level=2)
def init_diffusion_tracing(server_args, thread_label: str):
if not server_args.enable_trace:
return
from types import SimpleNamespace
from sglang.srt import server_args as srt_server_args_module
from sglang.srt.observability.trace import (
process_tracing_init,
trace_set_thread_info,
)
from sglang.srt.server_args import set_global_server_args_for_scheduler
# srt owns TraceReqContext and filters spans through its global trace_modules
try:
srt_server_args = srt_server_args_module.get_global_server_args()
except ValueError:
srt_server_args = SimpleNamespace(trace_modules=DIFFUSION_TRACE_MODULE)
set_global_server_args_for_scheduler(srt_server_args)
trace_modules = [
module.strip()
for module in getattr(srt_server_args, "trace_modules", "").split(",")
if module.strip()
]
if DIFFUSION_TRACE_MODULE not in trace_modules:
trace_modules.append(DIFFUSION_TRACE_MODULE)
srt_server_args.trace_modules = ",".join(trace_modules)
process_tracing_init(server_args.otlp_traces_endpoint, "sglang-diffusion")
trace_set_thread_info(thread_label)
@contextmanager
def trace_req(trace_ctx):
"""Ensure ``trace_req_finish()`` is called when a request scope exits.
@@ -19,6 +19,8 @@ from __future__ import annotations
import json
import unittest
from contextlib import contextmanager
from types import SimpleNamespace
import torch
@@ -28,8 +30,10 @@ from sglang.multimodal_gen.runtime.disaggregation.scheduler_mixin import (
)
from sglang.multimodal_gen.runtime.disaggregation.transport.codec import pack_tensors
from sglang.multimodal_gen.runtime.pipelines_core import Req
from sglang.srt import server_args as srt_server_args_module
from sglang.srt.observability import trace as srt_trace
from sglang.srt.observability.trace import TraceNullContext, TraceReqContext
from sglang.srt.server_args import set_global_server_args_for_scheduler
try:
from opentelemetry import propagate as otel_propagate
@@ -56,6 +60,16 @@ def _enable_minimal_otel() -> None:
srt_trace.trace_set_thread_info("TestThread")
@contextmanager
def _srt_trace_server_args():
prev_server_args = srt_server_args_module._global_server_args
set_global_server_args_for_scheduler(SimpleNamespace(trace_modules="request"))
try:
yield
finally:
srt_server_args_module._global_server_args = prev_server_args
def _traceparent_from(ctx) -> str | None:
"""Re-inject a W3C carrier from an OTel Context and return the traceparent.
@@ -133,35 +147,37 @@ class TestDisaggTracePropagation(unittest.TestCase):
TraceReqContext with an OTel Context (not the raw dict)."""
_enable_minimal_otel()
ctx = TraceReqContext(rid="test-on", role="server", module_name="request")
ctx.trace_req_start()
self.assertTrue(ctx.tracing_enable)
self.assertFalse(ctx.is_copy)
with _srt_trace_server_args():
ctx = TraceReqContext(rid="test-on", role="server", module_name="request")
ctx.trace_req_start()
self.assertTrue(ctx.tracing_enable)
self.assertFalse(ctx.is_copy)
req = Req(request_id="test-on", prompt="x")
req.trace_ctx = ctx
req = Req(request_id="test-on", prompt="x")
req.trace_ctx = ctx
_, scalar_fields = extract_transfer_fields(req)
self.assertNotIn("trace_ctx", scalar_fields)
self.assertIn("_trace_state", scalar_fields)
state = scalar_fields["_trace_state"]
self.assertTrue(state.get("tracing_enable"))
# W3C carrier must be present so downstream roles can nest spans.
self.assertIn("traceparent", state.get("root_span_context", {}))
_, scalar_fields = extract_transfer_fields(req)
self.assertNotIn("trace_ctx", scalar_fields)
self.assertIn("_trace_state", scalar_fields)
state = scalar_fields["_trace_state"]
self.assertTrue(state.get("tracing_enable"))
# W3C carrier must be present so downstream roles can nest spans.
self.assertIn("traceparent", state.get("root_span_context", {}))
decoded = _roundtrip_scalar_fields(scalar_fields)
self.assertEqual(decoded["_trace_state"], state)
decoded = _roundtrip_scalar_fields(scalar_fields)
self.assertEqual(decoded["_trace_state"], state)
rebuilt = object.__new__(TraceReqContext)
rebuilt.__setstate__(decoded["_trace_state"])
self.assertTrue(rebuilt.tracing_enable)
self.assertTrue(rebuilt.is_copy)
# The sender's traceparent must survive into the rebuilt Context so
# downstream role spans nest under the original trace_id.
self.assertEqual(
_traceparent_from(rebuilt.root_span_context),
state["root_span_context"]["traceparent"],
)
rebuilt = object.__new__(TraceReqContext)
rebuilt.__setstate__(decoded["_trace_state"])
self.assertTrue(rebuilt.tracing_enable)
self.assertTrue(rebuilt.is_copy)
# The sender's traceparent must survive into the rebuilt Context so
# downstream role spans nest under the original trace_id.
self.assertEqual(
_traceparent_from(rebuilt.root_span_context),
state["root_span_context"]["traceparent"],
)
ctx.trace_req_finish()
@unittest.skipUnless(_OTEL_AVAILABLE, "opentelemetry SDK not installed")
def test_build_disagg_req_installs_rebuilt_ctx(self):
@@ -170,22 +186,26 @@ class TestDisaggTracePropagation(unittest.TestCase):
the Req as a stray attribute."""
_enable_minimal_otel()
ctx = TraceReqContext(rid="test-brq", role="server", module_name="request")
ctx.trace_req_start()
with _srt_trace_server_args():
ctx = TraceReqContext(rid="test-brq", role="server", module_name="request")
ctx.trace_req_start()
req = Req(request_id="test-brq", prompt="x")
req.trace_ctx = ctx
_, scalar_fields = extract_transfer_fields(req)
self.assertIn("_trace_state", scalar_fields)
req = Req(request_id="test-brq", prompt="x")
req.trace_ctx = ctx
_, scalar_fields = extract_transfer_fields(req)
self.assertIn("_trace_state", scalar_fields)
# _build_disagg_req is an instance method but its body does not touch
# ``self``; call via __func__ to avoid needing a real Scheduler.
rebuilt = SchedulerDisaggMixin._build_disagg_req(None, dict(scalar_fields), {})
# _build_disagg_req is an instance method but its body does not
# touch ``self``; call via __func__ to avoid needing a real Scheduler.
rebuilt = SchedulerDisaggMixin._build_disagg_req(
None, dict(scalar_fields), {}
)
self.assertIsInstance(rebuilt.trace_ctx, TraceReqContext)
self.assertTrue(rebuilt.trace_ctx.tracing_enable)
self.assertTrue(rebuilt.trace_ctx.is_copy)
self.assertFalse(hasattr(rebuilt, "_trace_state"))
self.assertIsInstance(rebuilt.trace_ctx, TraceReqContext)
self.assertTrue(rebuilt.trace_ctx.tracing_enable)
self.assertTrue(rebuilt.trace_ctx.is_copy)
self.assertFalse(hasattr(rebuilt, "_trace_state"))
ctx.trace_req_finish()
@unittest.skipUnless(_OTEL_AVAILABLE, "opentelemetry SDK not installed")
def test_build_disagg_req_falls_back_when_tracing_off(self):