Spaces:
Running on Zero
Running on Zero
| """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 | |