Turn committed VAE chunks into MP4 bytes without materializing the video.
The producer side is the model's business: a Wan VAE walks a causal frame-by-frame loop, a MiniMax-H3 VAE walks overlapping clips, and Wan S2V decodes one clip per autoregressive iteration. What happens to a chunk once it is final is not: every one of them quantizes to uint8, moves to the host once, and lands in a bounded encoder. That consumer lives here so a model only has to publish finished chunks -- :class:SupportsChunkedVAEDecode for producers that can be driven, :class:ChunkedVideoMP4Session for producers that drive themselves.
The move to the host runs as its own scheduled stage. A chunk is quantized on the accelerator, copied without blocking into one of a fixed number of reusable pinned host slots on a dedicated stream, and only read once its copy-completion event has fired. The producer therefore returns to decoding while the copy is still in flight, and the slot count bounds how many chunks can be in flight at once.
One chunk in flight: a host slot, its copy event, and its readers.
The slot returns to the ring only after the last reader has finished with it. Being handed to an encoder is not enough -- the encoder's queue holds a view of this buffer until its muxing thread reads it -- so release is driven by the encoder rather than by dequeuing.
Drop an unread chunk during teardown, keeping the copy's memory alive.
The copy may still be writing into the slot and reading from the source tensor, so this waits for it rather than letting either be reused underneath an in-flight transfer.
dataclass ¶
What the host transfer cost, for checking that it actually overlapped.
copy_wait_seconds staying near zero while transfers climbs is the signal that copies completed behind the producer; it rising with them means the producer is outrunning the copy stream. slot_wait_seconds rising instead means the encoder is the bottleneck and backpressure has reached the producer.
Encode committed video chunks into one progressive MP4 per batch entry.
Push finished BCTHW chunks as the producer commits them; each batch entry gets its own bounded encoder, so host transfer and H.264 encoding overlap whatever the producer is still decoding.
audio_waveforms holds one waveform per batch entry (None for a silent entry), so a caller with several outputs per prompt repeats a request's waveform across its entries. batch_frames coalesces transfers for producers that publish finer than a transfer is worth; crop trims the decoder's padding to the requested output size. transfer_slots sets how many chunks may be in flight to the host at once.
property ¶
ring: PinnedChunkRing | None
The transfer ring this session settled on, once a chunk has flushed.
transfer_stats() -> ChunkTransferStats | None
Per-request transfer accounting, or None if no ring was used.
Fixed-depth pool of reusable host slots for chunk transfers.
A slot walks free -> copying -> ready -> encoding -> free and holds exactly one chunk for that whole trip, so pending transfer bytes are bounded by the depth times the largest chunk of the request. Waiting for a free slot is where encoder backpressure reaches the producer.
property ¶
slots_in_use: int
Slots currently holding a chunk, whether copying, ready or encoding.
abort(error: BaseException | None = None) -> None
Fail every present and future wait for a slot.
transfer(frames: Tensor, *, readers: int) -> ChunkLease
Start one non-blocking copy of frames into a free slot.
Quantize a BCTHW chunk and block until it reaches the host.
The path a session takes when no pinned ring is available, and the one a caller outside a session gets. A session on an accelerator uses :class:PinnedChunkRing instead so the copy does not stall the producer.
Decode z straight into one progressive MP4 per batch entry.
Drives a VAE that declares :class:SupportsChunkedVAEDecode, so host transfer and encoding overlap the remaining decode and the full video is never materialized. Ranks that own no decode output receive no chunks and get an empty list, matching the empty tensor the full-decode path returns there.
Quantize a BCTHW chunk to contiguous BTHWC uint8 on its own device.
value_range is the interval the producer publishes, which differs per checkpoint (see SupportsChunkedVAEDecode.chunk_value_range). Quantizing on the accelerator means the transfer that follows moves the final bytes rather than float frames, and making the result contiguous there keeps the copy a single flat move instead of a strided host-side gather.