[diffusion] improve: improve realtime webui playback pacing (#27148)
This commit is contained in:
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
|
||||
|
||||
|
||||
+4
@@ -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:
|
||||
|
||||
+161
-22
@@ -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,
|
||||
)
|
||||
|
||||
+70
-4
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user