Diffusers
MuseTalk1.5 / server /web /src /hooks /useWebSocket.js
Marcos
Add H.264 WebSocket streaming and React.js web interface
32bba92
Raw
History Blame Contribute Delete
2.11 kB
import { useState, useEffect, useRef, useCallback } from 'react'
export function useWebSocket(onMessage) {
const [status, setStatus] = useState('disconnected')
const wsRef = useRef(null)
const reconnectTimeoutRef = useRef(null)
const connect = useCallback(() => {
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'
// Always use root WebSocket path (works with proxy)
const wsUrl = `${protocol}//${window.location.host}/ws/stream`
console.log('[WS] Connecting to:', wsUrl)
const ws = new WebSocket(wsUrl)
ws.binaryType = 'arraybuffer'
ws.onopen = () => {
console.log('[WS] Connected')
setStatus('connected')
}
ws.onclose = () => {
console.log('[WS] Disconnected')
setStatus('disconnected')
reconnectTimeoutRef.current = setTimeout(connect, 3000)
}
ws.onerror = (e) => {
console.error('[WS] Error:', e)
setStatus('error')
}
ws.onmessage = async (event) => {
if (event.data instanceof ArrayBuffer) {
onMessage?.({ type: 'binary', data: event.data })
} else {
try {
const data = JSON.parse(event.data)
onMessage?.(data)
} catch (e) {
console.error('[WS] Parse error:', e)
}
}
}
wsRef.current = ws
// Keep alive
const pingInterval = setInterval(() => {
if (ws.readyState === WebSocket.OPEN) {
ws.send(JSON.stringify({ type: 'ping' }))
}
}, 10000)
return () => {
clearInterval(pingInterval)
ws.close()
}
}, [onMessage])
useEffect(() => {
connect()
return () => {
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current)
}
if (wsRef.current) {
wsRef.current.close()
}
}
}, [connect])
const send = useCallback((data) => {
if (wsRef.current?.readyState === WebSocket.OPEN) {
if (typeof data === 'string') {
wsRef.current.send(data)
} else {
wsRef.current.send(JSON.stringify(data))
}
}
}, [])
return { status, send }
}