Cccccz/HY

1"""Schema for BF16 joint-window Context pre-KV features."""2 3from __future__ import annotations4 5from typing import Mapping6 7import torch8 9from .schema import HIDDEN_SIZE10from .v2_schema import (11    CONTEXT_BLOCK_IDS,12    TOKENS_PER_FRAME,13    validate_case_tensors_v2,14    validate_step_tensors_v2,15)16 17 18SCHEMA_VERSION_PREFEATURE = "predictor_context_prefeature_bf16_v1"19SUPERVISION_PAIRS = ((0, 1), (1, 2), (2, 3))20DEFAULT_ACTIONS = (21    ("w_s", "w-15,s-16"),22    ("a_d", "a-15,d-16"),23    ("up_down", "up-15,down-16"),24    ("left_right", "left-15,right-16"),25)26 27 28def validate_case_tensors(tensors: Mapping[str, torch.Tensor]) -> int:29    return validate_case_tensors_v2(tensors)30 31 32def validate_step_tensors(tensors: Mapping[str, torch.Tensor]) -> None:33    validate_step_tensors_v2(tensors)34 35 36def validate_context_prefeature(37    block_id: int,38    tensors: Mapping[str, torch.Tensor],39) -> int:40    if block_id not in CONTEXT_BLOCK_IDS:41        raise ValueError(f"Unsupported context block: {block_id}")42    required = {43        "img_modulated",44        "context_valid_mask",45        "selected_frame_indices",46        "context_viewmats",47        "context_Ks",48        "rope_temporal_size",49        "start_rope_start_idx",50    }51    missing = required.difference(tensors)52    if missing:53        raise ValueError(f"Block {block_id} missing prefeature keys: {sorted(missing)}")54 55    feature = tensors["img_modulated"]56    if feature.ndim != 3 or feature.shape[0] != 1 or feature.shape[2] != HIDDEN_SIZE:57        raise ValueError(f"Unexpected block {block_id} feature shape: {tuple(feature.shape)}")58    if feature.dtype != torch.bfloat16:59        raise ValueError(f"Block {block_id} img_modulated must be BF16")60    if not torch.isfinite(feature).all():61        raise ValueError(f"Block {block_id} img_modulated contains non-finite values")62    tokens = int(feature.shape[1])63    if tokens % TOKENS_PER_FRAME:64        raise ValueError(f"Block {block_id} token count {tokens} is not frame-aligned")65    frames = tokens // TOKENS_PER_FRAME66 67    mask = tensors["context_valid_mask"]68    if tuple(mask.shape) != (1, tokens) or mask.dtype != torch.bool or not bool(mask.all()):69        raise ValueError("On-disk context_valid_mask must be all-True and unpadded")70    indices = tensors["selected_frame_indices"]71    if tuple(indices.shape) != (frames,) or indices.dtype != torch.int64:72        raise ValueError("selected_frame_indices must be int64 [context_frames]")73    if frames and (int(indices.min()) < 0 or not bool(torch.all(indices[1:] > indices[:-1]))):74        raise ValueError("selected_frame_indices must be non-negative and strictly increasing")75    if tuple(tensors["context_viewmats"].shape) != (1, frames, 4, 4):76        raise ValueError("Unexpected context_viewmats shape")77    if tuple(tensors["context_Ks"].shape) != (1, frames, 3, 3):78        raise ValueError("Unexpected context_Ks shape")79    for name in ("context_viewmats", "context_Ks"):80        value = tensors[name]81        if value.dtype != torch.bfloat16 or not torch.isfinite(value).all():82            raise ValueError(f"{name} must contain finite BF16 values")83    for name in ("rope_temporal_size", "start_rope_start_idx"):84        value = tensors[name]85        if value.dtype != torch.int64 or value.numel() != 1:86            raise ValueError(f"{name} must be one int64 scalar")87    if int(tensors["rope_temporal_size"].item()) != frames:88        raise ValueError("rope_temporal_size must equal the compact context frame count")89    if int(tensors["start_rope_start_idx"].item()) != 0:90        raise ValueError("Context prefill must start RoPE at zero")91    return frames92