258 lines
10 KiB
Python
258 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Bit-equivalence tests: DreamZeroStateAdapter vs bespoke DreamZeroState.
|
|
|
|
These drive both state objects directly with tiny CPU tensors -- no model, no
|
|
GPU -- and assert the adapter (backed by the SessionStateManager) stores and
|
|
returns exactly what the bespoke DreamZeroState does. This gates routing
|
|
DreamZero through the shared session state manager (RFC #4480).
|
|
|
|
DreamZero's attention KV is engine-owned (AR-Diffusion paged pool, PR #4534)
|
|
and out of the adapter's scope; the state surface tested here is the model's
|
|
non-KV session state.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm_omni.diffusion.models.dreamzero.state_dreamzero import (
|
|
FRAMES_PER_CHUNK,
|
|
DreamZeroState,
|
|
)
|
|
from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import DreamZeroStateAdapter
|
|
from vllm_omni.experimental.world_models.session_state import SessionStateManager
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _create_both() -> tuple[DreamZeroState, DreamZeroStateAdapter]:
|
|
return DreamZeroState(), DreamZeroStateAdapter("session-0", SessionStateManager())
|
|
|
|
|
|
def test_accumulate_frames_and_call_count_match() -> None:
|
|
bespoke, adapter = _create_both()
|
|
rng = np.random.default_rng(0)
|
|
for _ in range(2 * FRAMES_PER_CHUNK + 1):
|
|
frame = rng.integers(0, 255, size=(5, 6, 3), dtype=np.uint8)
|
|
b_out = bespoke.accumulate_frames(frame)
|
|
a_out = adapter.accumulate_frames(frame)
|
|
np.testing.assert_array_equal(b_out, a_out)
|
|
assert bespoke.call_count == adapter.call_count
|
|
assert len(bespoke.stitched_buffer) == len(adapter.stitched_buffer)
|
|
# Deque rollover capped at FRAMES_PER_CHUNK.
|
|
assert len(adapter.stitched_buffer) == FRAMES_PER_CHUNK
|
|
|
|
|
|
def test_accumulate_frames_4d_extend_matches() -> None:
|
|
bespoke, adapter = _create_both()
|
|
clip = np.random.default_rng(1).integers(0, 255, size=(3, 5, 6, 3), dtype=np.uint8)
|
|
np.testing.assert_array_equal(bespoke.accumulate_frames(clip), adapter.accumulate_frames(clip))
|
|
assert len(bespoke.stitched_buffer) == len(adapter.stitched_buffer)
|
|
|
|
|
|
def test_video_latents_roundtrip_matches() -> None:
|
|
bespoke, adapter = _create_both()
|
|
assert bespoke.get_concatenated_video_latents() is None
|
|
assert adapter.get_concatenated_video_latents() is None
|
|
|
|
rng = torch.Generator().manual_seed(0)
|
|
for _ in range(3):
|
|
# (B, T, C, H, W) as returned by the denoise loop.
|
|
chunk = torch.randn(1, 2, 3, 4, 4, generator=rng)
|
|
bespoke.append_video_latents(chunk)
|
|
adapter.append_video_latents(chunk)
|
|
torch.testing.assert_close(
|
|
bespoke.get_concatenated_video_latents(),
|
|
adapter.get_concatenated_video_latents(),
|
|
)
|
|
|
|
bespoke.clear_video_latents()
|
|
adapter.clear_video_latents()
|
|
assert bespoke.get_concatenated_video_latents() is None
|
|
assert adapter.get_concatenated_video_latents() is None
|
|
|
|
|
|
def test_vae_stream_fields_match_bespoke() -> None:
|
|
bespoke, adapter = _create_both()
|
|
for state in (bespoke, adapter):
|
|
assert state.vae_stream_initialized is False
|
|
assert state.vae_enc_feat_map is None
|
|
assert state.vae_encoder_out is None
|
|
assert state.vae_pending_body_frames is None
|
|
|
|
feat_map: list[torch.Tensor | None] = [torch.ones(1, 2), None]
|
|
encoder_out = torch.randn(1, 4, 2, 3, 3)
|
|
pending = torch.randn(1, 3, 2, 8, 8)
|
|
for state in (bespoke, adapter):
|
|
state.vae_enc_feat_map = feat_map
|
|
state.vae_encoder_out = encoder_out
|
|
state.vae_pending_body_frames = pending
|
|
state.vae_stream_initialized = True
|
|
|
|
assert adapter.vae_stream_initialized is bespoke.vae_stream_initialized is True
|
|
torch.testing.assert_close(adapter.vae_encoder_out, bespoke.vae_encoder_out)
|
|
torch.testing.assert_close(adapter.vae_pending_body_frames, bespoke.vae_pending_body_frames)
|
|
assert adapter.vae_enc_feat_map is feat_map
|
|
|
|
for state in (bespoke, adapter):
|
|
state.reset_vae_encoder_stream()
|
|
assert state.vae_stream_initialized is False
|
|
assert state.vae_enc_feat_map is None
|
|
assert state.vae_encoder_out is None
|
|
assert state.vae_pending_body_frames is None
|
|
|
|
|
|
def test_reset_clears_state_and_should_reset_matches() -> None:
|
|
bespoke, adapter = _create_both()
|
|
tokens = torch.tensor([1, 2, 3])
|
|
# language unset -> both want reset.
|
|
assert bespoke.should_reset(tokens, 4, -1) is adapter.should_reset(tokens, 4, -1) is True
|
|
|
|
for state in (bespoke, adapter):
|
|
state.language = tokens
|
|
state.current_start_frame = 2
|
|
# same language -> no reset; both agree.
|
|
assert bespoke.should_reset(tokens, 4, -1) == adapter.should_reset(tokens, 4, -1) is False
|
|
# language change -> reset.
|
|
other = torch.tensor([9, 9])
|
|
assert bespoke.should_reset(other, 4, -1) == adapter.should_reset(other, 4, -1) is True
|
|
# local_attn_size exceeded -> reset ("inference", not "session").
|
|
assert bespoke.reset_reason(tokens, 4, 1) == adapter.reset_reason(tokens, 4, 1) == "inference"
|
|
|
|
# A session reset clears everything, including the prompt-embed cache and
|
|
# the VAE encoder stream (mirrors DreamZeroState.reset()).
|
|
for state in (bespoke, adapter):
|
|
state.prompt_embeds = torch.ones(1, 2)
|
|
state.vae_encoder_out = torch.ones(1, 4, 1, 2, 2)
|
|
state.vae_stream_initialized = True
|
|
state.reset()
|
|
assert state.language is None
|
|
assert state.prompt_embeds is None
|
|
assert state.current_start_frame == 0
|
|
assert state.call_count == 0
|
|
assert state.vae_stream_initialized is False
|
|
assert state.vae_encoder_out is None
|
|
|
|
|
|
def test_window_reset_preserves_latents_prompt_and_vae_stream() -> None:
|
|
# reset(clear_video_latents=False) is the window ("inference") reset: the
|
|
# prompt is unchanged and the VAE feat_cache history is independent of the
|
|
# DiT attention window, so both survive alongside the video latents.
|
|
bespoke, adapter = _create_both()
|
|
tokens = torch.tensor([1, 2, 3])
|
|
embeds = torch.randn(1, 4)
|
|
encoder_out = torch.randn(1, 4, 2, 3, 3)
|
|
chunk = torch.randn(1, 2, 3, 4, 4)
|
|
for state in (bespoke, adapter):
|
|
state.language = tokens
|
|
state.prompt_embeds = embeds
|
|
state.vae_encoder_out = encoder_out
|
|
state.vae_stream_initialized = True
|
|
state.current_start_frame = 7
|
|
state.append_video_latents(chunk)
|
|
|
|
state.reset_inference_state()
|
|
|
|
assert state.current_start_frame == 0
|
|
assert state.call_count == 0
|
|
assert state.clip_feas is None
|
|
torch.testing.assert_close(state.language, tokens)
|
|
torch.testing.assert_close(state.prompt_embeds, embeds)
|
|
assert state.vae_stream_initialized is True
|
|
torch.testing.assert_close(state.vae_encoder_out, encoder_out)
|
|
assert state.get_concatenated_video_latents() is not None
|
|
|
|
torch.testing.assert_close(
|
|
bespoke.get_concatenated_video_latents(),
|
|
adapter.get_concatenated_video_latents(),
|
|
)
|
|
|
|
|
|
def test_metadata_fields_persist_across_adapter_instances() -> None:
|
|
# A fresh adapter for the same session must see prior metadata (the manager
|
|
# is the single source of truth).
|
|
manager = SessionStateManager()
|
|
first = DreamZeroStateAdapter("s", manager)
|
|
first.current_start_frame = 5
|
|
first.clip_feas = torch.ones(2, 2)
|
|
first.append_video_latents(torch.ones(1, 2, 3, 4, 4))
|
|
|
|
second = DreamZeroStateAdapter("s", manager)
|
|
assert second.current_start_frame == 5
|
|
torch.testing.assert_close(second.clip_feas, torch.ones(2, 2))
|
|
# Video latents appended via `first` are visible via `second`.
|
|
assert second.get_concatenated_video_latents() is not None
|
|
|
|
|
|
def test_session_lru_caps_retained_sessions() -> None:
|
|
manager = SessionStateManager(max_sessions=3)
|
|
for i in range(5):
|
|
DreamZeroStateAdapter(f"s{i}", manager)
|
|
assert len(manager) == 3
|
|
# Oldest two evicted; newest three retained.
|
|
assert "s0" not in manager
|
|
assert "s1" not in manager
|
|
for i in (2, 3, 4):
|
|
assert f"s{i}" in manager
|
|
|
|
|
|
def test_adapter_keeps_state_when_its_session_is_evicted() -> None:
|
|
# Regression: an adapter mid-generation must not lose its data when the
|
|
# manager evicts its session to make room for newer ones. The adapter pins
|
|
# the session, so its stored latents survive eviction from the lookup table.
|
|
manager = SessionStateManager(max_sessions=2)
|
|
adapter = DreamZeroStateAdapter("active", manager)
|
|
chunk = torch.randn(1, 2, 3, 4, 4)
|
|
adapter.append_video_latents(chunk)
|
|
adapter.language = torch.tensor([1, 2, 3])
|
|
|
|
# Push "active" out of the manager's table with newer sessions.
|
|
for i in range(3):
|
|
DreamZeroStateAdapter(f"other{i}", manager)
|
|
assert "active" not in manager
|
|
|
|
# The in-use adapter still returns its own data.
|
|
assert adapter.get_concatenated_video_latents() is not None
|
|
torch.testing.assert_close(adapter.language, torch.tensor([1, 2, 3]))
|
|
|
|
|
|
def test_session_reset_clears_metadata() -> None:
|
|
# Regression: resetting a session must not leave stale metadata behind, so a
|
|
# fresh adapter for the same session sees defaults, not the old values.
|
|
manager = SessionStateManager()
|
|
adapter = DreamZeroStateAdapter("s", manager)
|
|
adapter.language = torch.tensor([1, 2, 3])
|
|
adapter.current_start_frame = 5
|
|
adapter.clip_feas = torch.ones(2, 2)
|
|
adapter.prompt_embeds = torch.ones(1, 2)
|
|
|
|
adapter.reset()
|
|
|
|
fresh = DreamZeroStateAdapter("s", manager)
|
|
assert fresh.language is None
|
|
assert fresh.current_start_frame == 0
|
|
assert fresh.clip_feas is None
|
|
assert fresh.prompt_embeds is None
|
|
|
|
|
|
def test_bespoke_path_without_memory_manager() -> None:
|
|
# Regression: a pipeline built via __new__ (lightweight fixtures) seeds only
|
|
# the bespoke fields and never sets _memory_manager; _get_or_create_state
|
|
# must fall through to the bespoke DreamZeroState path without AttributeError.
|
|
from collections import OrderedDict
|
|
|
|
from vllm_omni.diffusion.models.dreamzero.pipeline_dreamzero import (
|
|
MAX_DREAMZERO_SESSIONS,
|
|
DreamZeroPipeline,
|
|
)
|
|
|
|
pipe = DreamZeroPipeline.__new__(DreamZeroPipeline)
|
|
pipe._states = OrderedDict()
|
|
pipe._max_session_states = MAX_DREAMZERO_SESSIONS
|
|
|
|
state = DreamZeroPipeline._get_or_create_state(pipe, "x")
|
|
assert isinstance(state, DreamZeroState)
|