deepgeek/gradio-webrtc
2
0
Fork 0
mirror of https://github.com/HumanAIGC-Engineering/gradio-webrtc.git synced 2026-02-04 09:29:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Close Stream from Backend (#222)

Browse Source 
* Close from backend

* Add code
This commit is contained in:
Freddy Boulton
2025-03-28 20:47:34 -04:00
committed by GitHub
parent 71743acb64
commit 8ed27fba78
10 changed files with 2336 additions and 2237 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
backend/fastrtc/__init__.py
View File

@@ -25,6 +25,7 @@ from .tracks import (
)
from .utils import (
AdditionalOutputs,
CloseStream,
Warning,
WebRTCError,
aggregate_bytes_to_16bit,
@@ -75,4 +76,5 @@ __all__ = [
"get_silero_model",
"SileroVadOptions",
"VideoStreamHandler",
"CloseStream",
]

4459
backend/fastrtc/templates/component/index.js
View File

File diff suppressed because one or more lines are too long

40
backend/fastrtc/tracks.py
View File

@@ -37,6 +37,7 @@ from numpy import typing as npt
from fastrtc.utils import (
AdditionalOutputs,
CloseStream,
DataChannel,
WebRTCError,
create_message,
@@ -54,9 +55,14 @@ VideoNDArray: TypeAlias = Union[
]
VideoEmitType = (
VideoNDArray | tuple[VideoNDArray, AdditionalOutputs] | AdditionalOutputs
VideoNDArray
| tuple[VideoNDArray, AdditionalOutputs]
| tuple[VideoNDArray, CloseStream]
| AdditionalOutputs
| CloseStream
)
VideoEventHandler = Callable[[npt.ArrayLike], VideoEmitType]
VideoEventGenerator = Generator[VideoEmitType, None, None]
VideoEventHandler = Callable[[npt.ArrayLike], VideoEmitType | VideoEventGenerator]
@dataclass
@@ -172,6 +178,12 @@ class VideoCallback(VideoStreamTrack):
args = self.add_frame_to_payload(cast(list, self.latest_args), frame_array)
array, outputs = split_output(self.event_handler(*args))
if isinstance(outputs, CloseStream):
cast(DataChannel, self.channel).send(
create_message("end_stream", outputs.msg)
)
self.stop()
return None
if (
isinstance(outputs, AdditionalOutputs)
and self.set_additional_outputs
@@ -444,6 +456,12 @@ class VideoStreamHandler_(VideoCallback):
):
self.set_additional_outputs(outputs)
self.channel.send(create_message("fetch_output", []))
if isinstance(outputs, CloseStream):
cast(DataChannel, self.channel).send(
create_message("end_stream", outputs.msg)
)
self.stop()
return
if array is None and self.mode == "send":
return
@@ -586,6 +604,12 @@ class AudioCallback(AudioStreamTrack):
await self.start()
frame = await self.queue.get()
if isinstance(frame, CloseStream):
cast(DataChannel, self.channel).send(
create_message("end_stream", frame.msg)
)
self.stop()
return
logger.debug("frame %s", frame)
data_time = frame.time
@@ -675,6 +699,12 @@ class ServerToClientVideo(VideoStreamTrack):
)
try:
next_array, outputs = split_output(next(self.generator))
if isinstance(outputs, CloseStream):
cast(DataChannel, self.channel).send(
create_message("end_stream", outputs.msg)
)
self.stop()
return
if (
isinstance(outputs, AdditionalOutputs)
and self.set_additional_outputs
@@ -770,6 +800,12 @@ class ServerToClientAudio(AudioStreamTrack):
await self.start()
data = await self.queue.get()
if isinstance(data, CloseStream):
cast(DataChannel, self.channel).send(
create_message("end_stream", data.msg)
)
self.stop()
return
if data is None:
self.stop()
return

21
backend/fastrtc/utils.py
View File

@@ -32,6 +32,11 @@ class AdditionalOutputs:
self.args = args
class CloseStream:
def __init__(self, msg: str = "Stream closed") -> None:
self.msg = msg
class DataChannel(Protocol):
def send(self, message: str) -> None: ...
@@ -39,6 +44,7 @@ class DataChannel(Protocol):
def create_message(
type: Literal[
"send_input",
"end_stream",
"fetch_output",
"stopword",
"error",
@@ -98,9 +104,13 @@ class WebRTCError(Exception):
_send_log(message, "error")
def split_output(data: tuple | Any) -> tuple[Any, AdditionalOutputs | None]:
def split_output(
data: tuple | Any,
) -> tuple[Any, AdditionalOutputs | CloseStream | None]:
if isinstance(data, AdditionalOutputs):
return None, data
if isinstance(data, CloseStream):
return None, data
if isinstance(data, tuple):
# handle the bare audio case
if 2 <= len(data) <= 3 and isinstance(data[1], np.ndarray):
@@ -109,11 +119,11 @@ def split_output(data: tuple | Any) -> tuple[Any, AdditionalOutputs | None]:
raise ValueError(
"The tuple must have exactly two elements: the data and an instance of AdditionalOutputs."
)
if not isinstance(data[-1], AdditionalOutputs):
if not isinstance(data[-1], (AdditionalOutputs, CloseStream)):
raise ValueError(
"The last element of the tuple must be an instance of AdditionalOutputs."
)
return data[0], cast(AdditionalOutputs, data[1])
return data[0], cast(AdditionalOutputs | CloseStream, data[1])
return data, None
@@ -152,6 +162,8 @@ async def player_worker_decode(
cast(DataChannel, channel()).send(create_message("fetch_output", []))
if frame is None:
if isinstance(outputs, CloseStream):
await queue.put(outputs)
if quit_on_none:
await queue.put(None)
break
@@ -203,7 +215,8 @@ async def player_worker_decode(
processed_frame.time_base = audio_time_base
audio_samples += processed_frame.samples
await queue.put(processed_frame)
if isinstance(outputs, CloseStream):
await queue.put(outputs)
except (TimeoutError, asyncio.TimeoutError):
logger.warning(
"Timeout in frame processing cycle after %s seconds - resetting", 60

2
frontend/Index.svelte
View File

@@ -47,6 +47,8 @@
msg?.type === "error"
) {
gradio.dispatch(msg?.type === "error" ? "error" : "warning", msg.message);
} else if (msg?.type === "end_stream") {
gradio.dispatch("warning", msg.data);
} else if (msg?.type === "fetch_output") {
gradio.dispatch("state_change");
} else if (msg?.type === "send_input") {

6
frontend/shared/InteractiveAudio.svelte
View File

@@ -51,12 +51,16 @@
}
});
let _on_change_cb = (msg: "change" | "tick" | "stopword") => {
let _on_change_cb = (msg: "change" | "tick" | "stopword" | any) => {
if (msg === "stopword") {
stopword_recognized = true;
setTimeout(() => {
stopword_recognized = false;
}, 3000);
} else if (msg.type === "end_stream") {
stream_state = "closed";
stop(pc);
on_change_cb(msg);
} else {
console.debug("calling on_change_cb with msg", msg);
on_change_cb(msg);

13
frontend/shared/StaticAudio.svelte
View File

@@ -29,6 +29,17 @@
let pc: RTCPeerConnection;
let _webrtc_id = Math.random().toString(36).substring(2);
let _on_change_cb = (msg: "change" | "tick" | "stopword" | any) => {
if (msg.type === "end_stream") {
on_change_cb(msg);
stream_state = "closed";
stop(pc);
} else {
console.debug("calling on_change_cb with msg", msg);
on_change_cb(msg);
}
};
const dispatch = createEventDispatcher<{
tick: undefined;
error: string;
@@ -75,7 +86,7 @@
server.offer,
_webrtc_id,
"audio",
on_change_cb,
_on_change_cb,
)
.then((connection) => {
clearTimeout(timeoutId);

13
frontend/shared/StaticVideo.svelte
View File

@@ -25,6 +25,17 @@
tick: undefined;
}>();
let _on_change_cb = (msg: "change" | "tick" | "stopword" | any) => {
if (msg.type === "end_stream") {
on_change_cb(msg);
stream_state = "closed";
stop(pc);
} else {
console.debug("calling on_change_cb with msg", msg);
on_change_cb(msg);
}
};
let stream_state = "closed";
$: if (value === "start_webrtc_stream") {
@@ -62,7 +73,7 @@
server.offer,
_webrtc_id,
"video",
on_change_cb,
_on_change_cb,
)
.then((connection) => {
clearTimeout(timeoutId);

14
frontend/shared/Webcam.svelte
View File

@@ -124,6 +124,18 @@
}
}
let _on_change_cb = (msg: "change" | "tick" | "stopword" | any) => {
if (msg.type === "end_stream") {
on_change_cb(msg);
stream_state = "closed";
stop(pc);
access_webcam();
} else {
console.debug("calling on_change_cb with msg", msg);
on_change_cb(msg);
}
};
let recording = false;
let stream: MediaStream;
@@ -171,7 +183,7 @@
server.offer,
webrtc_id,
"video",
on_change_cb,
_on_change_cb,
rtp_params,
undefined,
reject_cb,

3
frontend/shared/webrtc_utils.ts
View File

@@ -80,7 +80,8 @@ export async function start(
event_json?.type === "error" ||
event_json?.type === "send_input" ||
event_json?.type === "fetch_output" ||
event_json?.type === "stopword"
event_json?.type === "stopword" ||
event_json?.type === "end_stream"
) {
on_change_cb(event_json ?? event.data);
}