387 lines
15 KiB
Python
387 lines
15 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Typed continuation state for the LTX-2 streaming pipeline.
|
|
|
|
Segment N+1 conditions on segment N's trailing decoded frames and
|
|
denoised audio latents. The streaming runtime used to hold this state as
|
|
per-worker globals; lifting it into a typed, JSON-serializable object
|
|
lets clients snapshot, migrate, or round-trip it through an HTTP/RPC
|
|
boundary. The envelope ``ContinuationState(kind, payload)`` is the
|
|
shared public API; the typed class here owns the LTX-2 payload shape.
|
|
|
|
Serialization contract:
|
|
|
|
* Video frames → PNG bytes + base64, or a :class:`BlobStore` id.
|
|
* Audio latents → a self-describing safetensors blob + base64, or a
|
|
:class:`BlobStore` id. safetensors preserves ``bfloat16``, which a
|
|
raw-numpy round-trip cannot.
|
|
* The returned payload is always a plain JSON-serializable dict.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
from collections.abc import Mapping
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from fastvideo.api.compat import register_continuation_kind
|
|
from fastvideo.api.schema import ContinuationState
|
|
|
|
if TYPE_CHECKING:
|
|
import numpy as np
|
|
import torch
|
|
|
|
from fastvideo.entrypoints.streaming.session_store import BlobStore
|
|
|
|
LTX2_CONTINUATION_KIND = "ltx2.v1"
|
|
"""Public ``ContinuationState.kind`` for LTX-2 payloads."""
|
|
|
|
LTX2_CONTINUATION_SCHEMA_VERSION = 1
|
|
"""Payload schema version carried inside ``payload.schema_version``."""
|
|
|
|
DEFAULT_INLINE_THRESHOLD_BYTES = 2 * 1024 * 1024
|
|
"""Tensors larger than this go to the blob store (if available). 2 MiB
|
|
is below typical single-JSON-message limits (Dynamo: 4 MiB, Postgres
|
|
TOAST: 1 GiB) and well above per-frame PNG payloads (~200 KiB at
|
|
512x512)."""
|
|
|
|
|
|
@dataclass
|
|
class LTX2ContinuationState:
|
|
"""Typed LTX-2 continuation state carried between streaming segments.
|
|
|
|
``video_frames`` hold trailing decoded RGB frames (uint8 HxWx3) from
|
|
segment N for conditioning segment N+1 via the VAE encode path.
|
|
``audio_latents`` is the cached denoised audio latent tensor of shape
|
|
``[B, C, T, mel]`` that segment N+1 will copy into the overlap
|
|
region of its clean-latent conditioning.
|
|
|
|
Most fields map 1:1 onto the internal gpu_pool's per-worker state;
|
|
the only new concept is the ``*_blob_id`` fields, which allow large
|
|
tensors to live outside the JSON payload. See module docstring.
|
|
"""
|
|
|
|
segment_index: int = 0
|
|
"""Index of the *just-completed* segment. Segment 0 has no history;
|
|
state returned after segment 0 carries ``segment_index=0`` and the
|
|
caller uses ``segment_index + 1`` as the next segment number."""
|
|
|
|
video_frames: list[np.ndarray] | None = None
|
|
"""Trailing decoded frames, each an RGB uint8 ``np.ndarray`` shaped
|
|
``(H, W, 3)``. ``None`` when the state is blob-backed or unset."""
|
|
|
|
video_frames_blob_id: str | None = None
|
|
"""Blob store id when the frames live outside the payload."""
|
|
|
|
video_conditioning_frame_idx: int = 0
|
|
"""Target frame index inside the next segment that the trailing
|
|
frames align with (matches the LTX-2 ``ltx2_video_conditions``
|
|
tuple's ``frame_idx`` slot)."""
|
|
|
|
video_conditioning_strength: float = 1.0
|
|
"""Conditioning strength in [0, 1]. Matches the ``ltx2_video_
|
|
conditions`` tuple's strength slot."""
|
|
|
|
audio_latents: torch.Tensor | None = None
|
|
"""Denoised audio latent tensor of shape ``[B, C, T, mel]``.
|
|
``None`` when the state is blob-backed or unset."""
|
|
|
|
audio_latents_blob_id: str | None = None
|
|
"""Blob store id when audio latents live outside the payload."""
|
|
|
|
audio_sample_rate: int | None = None
|
|
"""Sample rate for the audio side (e.g. 24000)."""
|
|
|
|
audio_conditioning_num_frames: int = 0
|
|
"""Number of trailing audio frames that carry over as clean context
|
|
into segment N+1."""
|
|
|
|
audio_conditioning_strength: float = 1.0
|
|
"""Clean-latent mask value applied to the overlap region; 0.0 keeps
|
|
the cached audio entirely, 1.0 renoises from scratch."""
|
|
|
|
video_position_offset_sec: float = 0.0
|
|
"""Seconds by which video RoPE is shifted forward so the audio
|
|
prefix can sit at ``t >= 0`` when audio conditioning is longer than
|
|
video conditioning."""
|
|
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
"""Opaque metadata bag for forward-compat fields that don't need
|
|
their own typed slot yet (e.g. custom knob experiments)."""
|
|
|
|
def to_continuation_state(
|
|
self,
|
|
*,
|
|
blob_store: BlobStore | None = None,
|
|
inline_threshold_bytes: int = DEFAULT_INLINE_THRESHOLD_BYTES,
|
|
) -> ContinuationState:
|
|
"""Serialize into a public :class:`ContinuationState`.
|
|
|
|
When ``blob_store`` is given, tensors larger than
|
|
``inline_threshold_bytes`` are stored via
|
|
:meth:`BlobStore.put` and referenced by id; otherwise all data
|
|
is base64-encoded inline. The payload is always a plain
|
|
JSON-serializable dict.
|
|
"""
|
|
payload: dict[str, Any] = {
|
|
"schema_version": LTX2_CONTINUATION_SCHEMA_VERSION,
|
|
"segment_index": int(self.segment_index),
|
|
"video_conditioning_frame_idx": int(self.video_conditioning_frame_idx),
|
|
"video_conditioning_strength": float(self.video_conditioning_strength),
|
|
"audio_conditioning_num_frames": int(self.audio_conditioning_num_frames),
|
|
"audio_conditioning_strength": float(self.audio_conditioning_strength),
|
|
"video_position_offset_sec": float(self.video_position_offset_sec),
|
|
"metadata": dict(self.metadata),
|
|
}
|
|
if self.audio_sample_rate is not None:
|
|
payload["audio_sample_rate"] = int(self.audio_sample_rate)
|
|
|
|
video_payload = self._encode_video_frames(
|
|
blob_store=blob_store,
|
|
inline_threshold_bytes=inline_threshold_bytes,
|
|
)
|
|
if video_payload is not None:
|
|
payload["video"] = video_payload
|
|
|
|
audio_payload = self._encode_audio_latents(
|
|
blob_store=blob_store,
|
|
inline_threshold_bytes=inline_threshold_bytes,
|
|
)
|
|
if audio_payload is not None:
|
|
payload["audio"] = audio_payload
|
|
|
|
return ContinuationState(
|
|
kind=LTX2_CONTINUATION_KIND,
|
|
payload=payload,
|
|
)
|
|
|
|
@classmethod
|
|
def from_continuation_state(
|
|
cls,
|
|
state: ContinuationState,
|
|
*,
|
|
blob_store: BlobStore | None = None,
|
|
) -> LTX2ContinuationState:
|
|
"""Rebuild a typed state from a public :class:`ContinuationState`.
|
|
|
|
Raises :class:`ValueError` when the kind doesn't match or the
|
|
schema version is unsupported.
|
|
"""
|
|
if state.kind != LTX2_CONTINUATION_KIND:
|
|
raise ValueError(f"Expected ContinuationState.kind={LTX2_CONTINUATION_KIND!r}, "
|
|
f"got {state.kind!r}")
|
|
payload = state.payload or {}
|
|
version = int(payload.get("schema_version", LTX2_CONTINUATION_SCHEMA_VERSION))
|
|
if version != LTX2_CONTINUATION_SCHEMA_VERSION:
|
|
raise ValueError(f"Unsupported LTX-2 continuation schema_version={version}; "
|
|
f"this build expects {LTX2_CONTINUATION_SCHEMA_VERSION}")
|
|
|
|
out = cls(
|
|
segment_index=int(payload.get("segment_index", 0)),
|
|
video_conditioning_frame_idx=int(payload.get("video_conditioning_frame_idx", 0)),
|
|
video_conditioning_strength=float(payload.get("video_conditioning_strength", 1.0)),
|
|
audio_sample_rate=(int(payload["audio_sample_rate"]) if "audio_sample_rate" in payload else None),
|
|
audio_conditioning_num_frames=int(payload.get("audio_conditioning_num_frames", 0)),
|
|
audio_conditioning_strength=float(payload.get("audio_conditioning_strength", 1.0)),
|
|
video_position_offset_sec=float(payload.get("video_position_offset_sec", 0.0)),
|
|
metadata=dict(payload.get("metadata") or {}),
|
|
)
|
|
|
|
video = payload.get("video")
|
|
if isinstance(video, Mapping):
|
|
cls._decode_video_frames(out, video, blob_store=blob_store)
|
|
|
|
audio = payload.get("audio")
|
|
if isinstance(audio, Mapping):
|
|
cls._decode_audio_latents(out, audio, blob_store=blob_store)
|
|
|
|
return out
|
|
|
|
# ------------------------------------------------------------------
|
|
# Video frame helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _encode_video_frames(
|
|
self,
|
|
*,
|
|
blob_store: BlobStore | None,
|
|
inline_threshold_bytes: int,
|
|
) -> dict[str, Any] | None:
|
|
if self.video_frames_blob_id is not None:
|
|
return {"blob_id": self.video_frames_blob_id}
|
|
if not self.video_frames:
|
|
return None
|
|
|
|
encoded = [_encode_png(frame) for frame in self.video_frames]
|
|
total = sum(len(b) for b in encoded)
|
|
if blob_store is not None and total > inline_threshold_bytes:
|
|
concatenated = _pack_frame_blobs(encoded)
|
|
blob_id = blob_store.put(
|
|
concatenated,
|
|
mime="application/x-fastvideo-frames+png",
|
|
)
|
|
return {"blob_id": blob_id, "frame_count": len(encoded)}
|
|
return {
|
|
"frames_b64": [base64.b64encode(b).decode("ascii") for b in encoded],
|
|
}
|
|
|
|
@staticmethod
|
|
def _decode_video_frames(
|
|
out: LTX2ContinuationState,
|
|
video: Mapping[str, Any],
|
|
*,
|
|
blob_store: BlobStore | None,
|
|
) -> None:
|
|
blob_id = video.get("blob_id")
|
|
if isinstance(blob_id, str):
|
|
if blob_store is None:
|
|
out.video_frames_blob_id = blob_id
|
|
return
|
|
raw = blob_store.get(blob_id)
|
|
encoded = _unpack_frame_blobs(raw)
|
|
out.video_frames = [_decode_png(b) for b in encoded]
|
|
return
|
|
frames_b64 = video.get("frames_b64")
|
|
if isinstance(frames_b64, list):
|
|
decoded = [_decode_png(base64.b64decode(b)) for b in frames_b64 if isinstance(b, str)]
|
|
out.video_frames = decoded or None
|
|
|
|
# ------------------------------------------------------------------
|
|
# Audio latent helpers
|
|
# ------------------------------------------------------------------
|
|
|
|
def _encode_audio_latents(
|
|
self,
|
|
*,
|
|
blob_store: BlobStore | None,
|
|
inline_threshold_bytes: int,
|
|
) -> dict[str, Any] | None:
|
|
if self.audio_latents_blob_id is not None:
|
|
return {"blob_id": self.audio_latents_blob_id}
|
|
if self.audio_latents is None:
|
|
return None
|
|
raw = _tensor_to_safetensors_bytes(self.audio_latents)
|
|
if blob_store is not None and len(raw) > inline_threshold_bytes:
|
|
blob_id = blob_store.put(
|
|
raw,
|
|
mime="application/x-fastvideo-tensor+safetensors",
|
|
)
|
|
return {"blob_id": blob_id}
|
|
return {"safetensors_b64": base64.b64encode(raw).decode("ascii")}
|
|
|
|
@staticmethod
|
|
def _decode_audio_latents(
|
|
out: LTX2ContinuationState,
|
|
audio: Mapping[str, Any],
|
|
*,
|
|
blob_store: BlobStore | None,
|
|
) -> None:
|
|
blob_id = audio.get("blob_id")
|
|
if isinstance(blob_id, str):
|
|
if blob_store is None:
|
|
out.audio_latents_blob_id = blob_id
|
|
return
|
|
raw = blob_store.get(blob_id)
|
|
out.audio_latents = _safetensors_bytes_to_tensor(raw)
|
|
return
|
|
data_b64 = audio.get("safetensors_b64")
|
|
if isinstance(data_b64, str):
|
|
out.audio_latents = _safetensors_bytes_to_tensor(base64.b64decode(data_b64))
|
|
|
|
|
|
def _encode_png(frame: np.ndarray) -> bytes:
|
|
"""Encode an ``(H, W, 3)`` uint8 RGB frame as PNG bytes."""
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
if not isinstance(frame, np.ndarray):
|
|
raise TypeError(f"LTX2 continuation frame must be a numpy ndarray, got {type(frame).__name__}")
|
|
if frame.dtype != np.uint8 or frame.ndim != 3 or frame.shape[-1] != 3:
|
|
raise ValueError("LTX2 continuation frame must be uint8 HxWx3 RGB; got "
|
|
f"dtype={frame.dtype}, shape={frame.shape}")
|
|
import io
|
|
|
|
buffer = io.BytesIO()
|
|
Image.fromarray(frame).save(buffer, format="PNG")
|
|
return buffer.getvalue()
|
|
|
|
|
|
def _decode_png(data: bytes) -> np.ndarray:
|
|
import io
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
img = Image.open(io.BytesIO(data)).convert("RGB")
|
|
return np.array(img, dtype=np.uint8)
|
|
|
|
|
|
def _pack_frame_blobs(encoded: list[bytes]) -> bytes:
|
|
"""Pack multiple PNG blobs into a single blob for blob-store storage.
|
|
|
|
Format: ``[4-byte big-endian count][4-byte len][png][4-byte len][png]...``.
|
|
"""
|
|
parts: list[bytes] = [len(encoded).to_bytes(4, "big")]
|
|
for blob in encoded:
|
|
parts.append(len(blob).to_bytes(4, "big"))
|
|
parts.append(blob)
|
|
return b"".join(parts)
|
|
|
|
|
|
def _unpack_frame_blobs(raw: bytes) -> list[bytes]:
|
|
if len(raw) < 4:
|
|
raise ValueError("frame blob truncated: missing count header")
|
|
count = int.from_bytes(raw[:4], "big")
|
|
# Each frame contributes at least a 4-byte length prefix, so a
|
|
# declared count larger than (len(raw) - 4) // 4 cannot fit and
|
|
# would otherwise cause an O(count) allocation loop on malformed
|
|
# input.
|
|
if count > (len(raw) - 4) // 4:
|
|
raise ValueError(f"frame blob declares {count} frames but buffer holds at most "
|
|
f"{(len(raw) - 4) // 4}")
|
|
out: list[bytes] = []
|
|
cursor = 4
|
|
for index in range(count):
|
|
if cursor + 4 > len(raw):
|
|
raise ValueError(f"frame blob truncated at frame {index} length header")
|
|
length = int.from_bytes(raw[cursor:cursor + 4], "big")
|
|
cursor += 4
|
|
if cursor + length > len(raw):
|
|
raise ValueError(f"frame blob truncated at frame {index} payload")
|
|
out.append(raw[cursor:cursor + length])
|
|
cursor += length
|
|
return out
|
|
|
|
|
|
def _tensor_to_safetensors_bytes(tensor: Any) -> bytes:
|
|
"""Serialize a torch tensor to a self-describing safetensors blob.
|
|
|
|
Uses the in-memory safetensors API so the wire format preserves
|
|
dtype (including ``bfloat16``, which a raw-numpy path cannot) and
|
|
shape without needing sidecar metadata.
|
|
"""
|
|
import torch
|
|
from safetensors.torch import save as st_save
|
|
|
|
if isinstance(tensor, torch.Tensor):
|
|
return st_save({"t": tensor.detach().cpu()})
|
|
import numpy as np
|
|
if isinstance(tensor, np.ndarray):
|
|
return st_save({"t": torch.from_numpy(np.ascontiguousarray(tensor))})
|
|
raise TypeError("LTX2 audio_latents must be a torch.Tensor or numpy.ndarray, got "
|
|
f"{type(tensor).__name__}")
|
|
|
|
|
|
def _safetensors_bytes_to_tensor(raw: bytes) -> Any:
|
|
from safetensors.torch import load as st_load
|
|
return st_load(raw)["t"]
|
|
|
|
|
|
register_continuation_kind(LTX2_CONTINUATION_KIND)
|
|
|
|
__all__ = [
|
|
"DEFAULT_INLINE_THRESHOLD_BYTES",
|
|
"LTX2ContinuationState",
|
|
"LTX2_CONTINUATION_KIND",
|
|
"LTX2_CONTINUATION_SCHEMA_VERSION",
|
|
]
|