File size: 3,247 Bytes
ba798d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | """Predictor-DisCa: current/anchor fusion without AR Context Vision K/V."""
from __future__ import annotations
from typing import Mapping
import torch
from .predictor_v2 import HYWorldPlayPredictorV2, PredictorV2Config
class HYWorldPlayPredictorDisCa(HYWorldPlayPredictorV2):
"""Two-block Predictor using only current/anchor image features and text K/V.
This is the project-local ``predictor-disca`` variant. Unlike Predictor-v4,
it neither fuses previous-chunk hidden states nor injects Context Vision K/V.
The cached text K/V remains available because HY-WorldPlay AR vision blocks
require text conditioning.
"""
uses_context_kv = True
uses_context_vision_kv = False
uses_context_prefeature = False
requires_previous_chunk_hidden = False
def _text_only_cache(
self,
values: Mapping[int, Mapping[str, torch.Tensor]],
) -> dict[int, dict[str, torch.Tensor | None]]:
result: dict[int, dict[str, torch.Tensor | None]] = {}
for block_id in self.predictor_config.source_block_ids:
if block_id not in values:
raise ValueError(f"Missing Predictor-DisCa text block {block_id}")
block = values[block_id]
missing = {"k_txt", "v_txt"}.difference(block)
if missing:
raise ValueError(
f"Predictor-DisCa block {block_id} lacks {sorted(missing)}"
)
result[block_id] = {
"k_txt": block["k_txt"],
"v_txt": block["v_txt"],
"text_valid_mask": block.get("text_valid_mask"),
# Deliberately discard online Context Vision K/V and construct
# zero-length vision caches in HYWorldPlayPredictorV2._make_cache.
"k_vision": None,
"v_vision": None,
}
return result
def forward(
self,
*,
target_model_input: torch.Tensor,
anchor_hidden: torch.Tensor,
target_frame_condition: torch.Tensor,
target_viewmats: torch.Tensor,
target_Ks: torch.Tensor,
rope_temporal_size: torch.Tensor | int,
start_rope_start_idx: torch.Tensor | int,
text_kv: Mapping[int, Mapping[str, torch.Tensor]] | None = None,
context_kv: Mapping[int, Mapping[str, torch.Tensor]] | None = None,
previous_chunk_hidden: torch.Tensor | None = None,
) -> dict[str, torch.Tensor]:
if previous_chunk_hidden is not None:
raise ValueError("Predictor-DisCa does not accept previous_chunk_hidden")
source = text_kv if text_kv is not None else context_kv
if source is None:
raise ValueError("Predictor-DisCa requires text_kv or online context_kv")
return super().forward(
target_model_input=target_model_input,
anchor_hidden=anchor_hidden,
target_frame_condition=target_frame_condition,
target_viewmats=target_viewmats,
target_Ks=target_Ks,
rope_temporal_size=rope_temporal_size,
start_rope_start_idx=start_rope_start_idx,
context_kv=self._text_only_cache(source),
previous_chunk_hidden=None,
)
|