97 lines
3.1 KiB
Python
97 lines
3.1 KiB
Python
"""Binary stream framing for engine responses.
|
|
|
|
One generation request is answered with a chunked HTTP body made of frames:
|
|
|
|
u32 header_len (little endian) | json header (utf-8) |
|
|
u64 payload_len (little endian) | payload bytes
|
|
|
|
Every frame header has "type" and "request_id". Text-only frames carry an
|
|
empty payload. The javascript reader (ui/src/utils/engineStream.ts) mirrors
|
|
this exactly. Frame types:
|
|
|
|
start resolved model/sample config, output kind, latent preview info
|
|
status a holder stage line (loading / quantizing / encoding ...)
|
|
progress {step, total, elapsed}
|
|
latent raw tensor payload; header carries shape, dtype, layout
|
|
result {path, url, kind, ext, ...} for each produced file
|
|
error {message, cancelled}
|
|
end terminal frame
|
|
"""
|
|
|
|
import json
|
|
import struct
|
|
from typing import Iterator, Optional, Tuple
|
|
|
|
import numpy as np
|
|
|
|
_HDR = struct.Struct("<I")
|
|
_LEN = struct.Struct("<Q")
|
|
|
|
|
|
def encode_frame(header: dict, payload: bytes = b"") -> bytes:
|
|
hdr = json.dumps(header, separators=(",", ":")).encode("utf-8")
|
|
return _HDR.pack(len(hdr)) + hdr + _LEN.pack(len(payload)) + payload
|
|
|
|
|
|
def tensor_payload(tensor, dtype="float16") -> Tuple[dict, bytes]:
|
|
"""Serialize a tensor to raw bytes + the header fields that describe it."""
|
|
import torch
|
|
|
|
t = tensor.detach()
|
|
if t.is_cuda or t.device.type != "cpu":
|
|
t = t.to("cpu")
|
|
if dtype == "float16":
|
|
t = t.to(torch.float16)
|
|
elif dtype == "float32":
|
|
t = t.to(torch.float32)
|
|
arr = t.contiguous().numpy()
|
|
return {"shape": list(arr.shape), "dtype": str(arr.dtype)}, arr.tobytes()
|
|
|
|
|
|
def payload_to_array(header: dict, payload: bytes) -> np.ndarray:
|
|
return np.frombuffer(payload, dtype=np.dtype(header["dtype"])).reshape(header["shape"])
|
|
|
|
|
|
class FrameReader:
|
|
"""Decode frames from an iterator of byte chunks (tests / python clients)."""
|
|
|
|
def __init__(self, chunks: Iterator[bytes]):
|
|
self._chunks = chunks
|
|
self._buf = bytearray()
|
|
self._eof = False
|
|
|
|
def _fill(self, n: int) -> bool:
|
|
while len(self._buf) < n and not self._eof:
|
|
try:
|
|
chunk = next(self._chunks)
|
|
except StopIteration:
|
|
self._eof = True
|
|
break
|
|
if chunk:
|
|
self._buf.extend(chunk)
|
|
return len(self._buf) >= n
|
|
|
|
def _take(self, n: int) -> bytes:
|
|
out = bytes(self._buf[:n])
|
|
del self._buf[:n]
|
|
return out
|
|
|
|
def read(self) -> Optional[Tuple[dict, bytes]]:
|
|
if not self._fill(_HDR.size):
|
|
return None
|
|
(hlen,) = _HDR.unpack(self._take(_HDR.size))
|
|
if not self._fill(hlen + _LEN.size):
|
|
raise EOFError("truncated frame header")
|
|
header = json.loads(self._take(hlen).decode("utf-8"))
|
|
(plen,) = _LEN.unpack(self._take(_LEN.size))
|
|
if not self._fill(plen):
|
|
raise EOFError("truncated frame payload")
|
|
return header, self._take(plen)
|
|
|
|
def __iter__(self):
|
|
while True:
|
|
frame = self.read()
|
|
if frame is None:
|
|
return
|
|
yield frame
|