vllm_omni.diffusion.models.dreamzero.state_dreamzero ¶

DreamZero pipeline persistent state.

FRAMES_PER_CHUNK `module-attribute` ¶

FRAMES_PER_CHUNK = 4

logger = logging.getLogger(__name__)

Pipeline persistent state across forward() calls.

Lifecycle

accumulate_frames(stitched: ndarray) -> ndarray

Accumulate stitched frames and return multi-frame video.

Parameters:

Name	Type	Description	Default
`stitched`	`ndarray`	(H, W, C) single frame or (T, H, W, C) multi-frame, already stitched by transform.	required

Returns:

Type	Description
`ndarray`	(T, H, W, C) ndarray. T=1 for first call, T=FRAMES_PER_CHUNK(4) after.

append_video_latents(video_out: Tensor) -> None

Append one AR chunk of normalized video latents for later decode.

Parameters:

Name	Type	Description	Default
`video_out`	`Tensor`	`(B, T, C, H, W)` tensor returned by the denoise loop (before the `transpose(1, 2)` stored in `DiffusionOutput`).	required

clear_video_latents() -> None

Drop accumulated video latents without resetting KV/frame state.

get_concatenated_video_latents() -> Tensor | None

Return all accumulated chunks concatenated along the time dimension.

reset(*, clear_video_latents: bool = True) -> None

Clear session state.

Parameters:

Name	Type	Description	Default
`clear_video_latents`	`bool`	When `False`, keep `video_latents_across_time` so offline/online export can decode the full AR rollout.	`True`

reset_inference_state() -> None

Reset KV/frame state after local attention rolls without dropping video latents.

reset_reason(
    text_tokens: Tensor | None,
    num_video_frames: int,
    local_attn_size: int,
) -> str | None

Return why state should reset before the next forward(), if any.

reset_vae_encoder_stream() -> None

Clear incremental Wan VAE encoder state used across AR steps.

should_reset(
    text_tokens: Tensor | None,
    num_video_frames: int,
    local_attn_size: int,
) -> bool

Determine if state should be reset before this forward().