223 lines
8.7 KiB
Python
223 lines
8.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""DreamZero pipeline persistent state."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from collections import deque
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Number of frames per chunk for subsequent calls (first call uses 1)
|
|
FRAMES_PER_CHUNK = 4
|
|
|
|
|
|
class DreamZeroState:
|
|
"""Pipeline persistent state across forward() calls.
|
|
|
|
Lifecycle:
|
|
- Created once in DreamZeroPipeline.__init__()
|
|
- Mutated every forward() call (frame append, KV cache grow)
|
|
- reset() on new session / language change
|
|
- reset_inference_state() on local_attn_size exceeded (KV only)
|
|
"""
|
|
|
|
def __init__(self) -> None:
|
|
self.reset()
|
|
|
|
# ------------------------------------------------------------------
|
|
# Frame accumulation (single stitched buffer)
|
|
# Transform outputs stitched single frame per call.
|
|
# We accumulate here to build multi-frame video for AR inference.
|
|
# ------------------------------------------------------------------
|
|
|
|
def accumulate_frames(self, stitched: np.ndarray) -> np.ndarray:
|
|
"""Accumulate stitched frames and return multi-frame video.
|
|
|
|
Args:
|
|
stitched: (H, W, C) single frame or (T, H, W, C) multi-frame,
|
|
already stitched by transform.
|
|
|
|
Returns:
|
|
(T, H, W, C) ndarray. T=1 for first call, T=FRAMES_PER_CHUNK(4) after.
|
|
"""
|
|
if stitched.ndim == 3:
|
|
self.stitched_buffer.append(stitched)
|
|
elif stitched.ndim == 4:
|
|
self.stitched_buffer.extend(list(stitched))
|
|
else:
|
|
raise ValueError(f"Expected 3D or 4D stitched, got {stitched.ndim}D")
|
|
|
|
num_frames = 1 if self.call_count == 0 else FRAMES_PER_CHUNK
|
|
|
|
buffer_frames = list(self.stitched_buffer)
|
|
if len(buffer_frames) >= num_frames:
|
|
frames = buffer_frames[-num_frames:]
|
|
else:
|
|
frames = buffer_frames
|
|
while len(frames) < num_frames:
|
|
frames.insert(0, buffer_frames[0])
|
|
|
|
self.call_count += 1
|
|
return np.stack(frames, axis=0) # (T, H, W, C)
|
|
|
|
def append_video_latents(self, video_out: torch.Tensor) -> None:
|
|
"""Append one AR chunk of normalized video latents for later decode.
|
|
|
|
Args:
|
|
video_out: ``(B, T, C, H, W)`` tensor returned by the denoise loop
|
|
(before the ``transpose(1, 2)`` stored in ``DiffusionOutput``).
|
|
"""
|
|
if video_out.dim() != 5:
|
|
raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}")
|
|
# Upstream ``torch.cat(self.video_across_time, dim=2)`` uses (B, C, T, H, W).
|
|
chunk = video_out.transpose(1, 2).detach().cpu()
|
|
self.video_latents_across_time.append(chunk)
|
|
logger.info(
|
|
"append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d",
|
|
tuple(chunk.shape),
|
|
len(self.video_latents_across_time),
|
|
int(sum(c.shape[2] for c in self.video_latents_across_time)),
|
|
)
|
|
|
|
def get_concatenated_video_latents(self) -> torch.Tensor | None:
|
|
"""Return all accumulated chunks concatenated along the time dimension."""
|
|
if not self.video_latents_across_time:
|
|
return None
|
|
if len(self.video_latents_across_time) == 1:
|
|
return self.video_latents_across_time[0]
|
|
return torch.cat(self.video_latents_across_time, dim=2)
|
|
|
|
def clear_video_latents(self) -> None:
|
|
"""Drop accumulated video latents without resetting KV/frame state."""
|
|
self.video_latents_across_time = []
|
|
|
|
# ------------------------------------------------------------------
|
|
# Reset / should_reset
|
|
# ------------------------------------------------------------------
|
|
|
|
def reset(self, *, clear_video_latents: bool = True) -> None:
|
|
"""Clear session state.
|
|
|
|
Args:
|
|
clear_video_latents: When ``False``, keep ``video_latents_across_time``
|
|
so offline/online export can decode the full AR rollout.
|
|
"""
|
|
saved_video_latents = [] if clear_video_latents else list(self.video_latents_across_time)
|
|
self.stitched_buffer = deque(maxlen=FRAMES_PER_CHUNK)
|
|
self.call_count = 0
|
|
# Normalized VAE latents per AR chunk, each (B, C, T, H, W) on CPU.
|
|
# Concatenated along T before decode (matches upstream video_across_time).
|
|
self.video_latents_across_time = saved_video_latents
|
|
|
|
# KV cache once robot-policy diffusion supports that integration.
|
|
self.kv_cache = None
|
|
self.kv_cache_neg = None
|
|
self.crossattn_cache = None
|
|
self.crossattn_cache_neg = None
|
|
self.current_start_frame = 0
|
|
|
|
self.clip_feas = None
|
|
self.ys = None
|
|
if clear_video_latents:
|
|
self.language = None
|
|
|
|
def reset_inference_state(self) -> None:
|
|
"""Reset KV/frame state after local attention rolls without dropping video latents."""
|
|
self.reset(clear_video_latents=False)
|
|
|
|
def reset_reason(
|
|
self,
|
|
text_tokens: torch.Tensor | None,
|
|
num_video_frames: int,
|
|
local_attn_size: int,
|
|
) -> str | None:
|
|
"""Return why state should reset before the next forward(), if any."""
|
|
if self.language is None:
|
|
logger.info("language is None, resetting")
|
|
return "session"
|
|
|
|
if text_tokens is not None and not torch.equal(self.language, text_tokens):
|
|
logger.info("language changed, resetting")
|
|
return "session"
|
|
|
|
# NOTE: after accumulate_frames, num_video_frames is the accumulated T
|
|
# (1 for first call, 4 for subsequent). Only reset on true single-frame
|
|
# which happens when the stitched_buffer was cleared externally.
|
|
if num_video_frames == 1 and self.call_count > 1:
|
|
logger.info("single frame input after first call, resetting")
|
|
return "session"
|
|
|
|
if local_attn_size != -1 and self.current_start_frame >= local_attn_size:
|
|
logger.info(
|
|
"current_start_frame %d >= local_attn_size %d, resetting inference state",
|
|
self.current_start_frame,
|
|
local_attn_size,
|
|
)
|
|
return "inference"
|
|
|
|
return None
|
|
|
|
def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool:
|
|
"""Determine if state should be reset before this forward()."""
|
|
return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None
|
|
|
|
# ------------------------------------------------------------------
|
|
# KV cache management
|
|
# ------------------------------------------------------------------
|
|
|
|
def create_kv_caches(
|
|
self,
|
|
batch_size: int,
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
num_layers: int,
|
|
num_heads: int,
|
|
head_dim: int,
|
|
) -> None:
|
|
"""Initialize empty KV caches and cross-attention caches."""
|
|
self.kv_cache = [
|
|
torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
|
|
]
|
|
self.kv_cache_neg = [
|
|
torch.zeros(2, batch_size, 0, num_heads, head_dim, dtype=dtype, device=device) for _ in range(num_layers)
|
|
]
|
|
|
|
self.crossattn_cache = [
|
|
{"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
|
|
]
|
|
self.crossattn_cache_neg = [
|
|
{"is_init": False, "k": None, "v": None, "k_img": None, "v_img": None} for _ in range(num_layers)
|
|
]
|
|
|
|
def update_kv_cache(
|
|
self,
|
|
layer_index: int,
|
|
updated_kv: torch.Tensor,
|
|
is_negative: bool = False,
|
|
) -> None:
|
|
"""Update a single layer's KV cache after prefill."""
|
|
cache = self.kv_cache_neg if is_negative else self.kv_cache
|
|
if cache is None:
|
|
raise RuntimeError("KV caches not initialized, call create_kv_caches first.")
|
|
cache[layer_index] = updated_kv.clone()
|
|
|
|
def get_kv_caches(self, is_negative: bool = False) -> list[torch.Tensor]:
|
|
"""Get KV caches for the specified branch."""
|
|
cache = self.kv_cache_neg if is_negative else self.kv_cache
|
|
if cache is None:
|
|
raise RuntimeError("KV caches not initialized.")
|
|
return cache
|
|
|
|
def get_crossattn_caches(self, is_negative: bool = False) -> list[dict[str, bool | torch.Tensor | None]]:
|
|
"""Get cross-attention caches for the specified branch."""
|
|
cache = self.crossattn_cache_neg if is_negative else self.crossattn_cache
|
|
if cache is None:
|
|
raise RuntimeError("Cross-attn caches not initialized.")
|
|
return cache
|