| """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"), |
| |
| |
| "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, |
| ) |
|
|
