Cccccz/HY

1"""Dataset schema constants and tensor validation."""2 3from __future__ import annotations4 5from typing import Mapping6 7import torch8 9 10SCHEMA_VERSION = "predictor_v1_txt_features_no_context_kv"11NUM_STEPS = 412LATENT_CHANNELS = 3213MODEL_INPUT_CHANNELS = 6514CHUNK_LATENT_FRAMES = 415LATENT_HEIGHT = 3016LATENT_WIDTH = 5217HIDDEN_SIZE = 204818TOKENS_PER_CHUNK = CHUNK_LATENT_FRAMES * LATENT_HEIGHT * LATENT_WIDTH19 20STEP_FIELDS = (21    "timestep",22    "noisy_sample",23    "frame_condition",24    "final_hidden",25    "velocity",26)27 28 29def expected_chunk_keys() -> set[str]:30    keys = {31        "action_labels",32        "target_viewmats",33        "target_Ks",34        "rope_temporal_size",35        "start_rope_start_idx",36    }37    for step in range(NUM_STEPS):38        keys.update(f"step_{step}_{field}" for field in STEP_FIELDS)39    return keys40 41 42def validate_case_tensors(tensors: Mapping[str, torch.Tensor]) -> None:43    required = {44        "image_condition_latent",45        "current_txt",46        "cached_txt",47        "vec_txt",48    }49    missing = required.difference(tensors)50    if missing:51        raise ValueError(f"Missing case tensor keys: {sorted(missing)}")52 53    image_condition = tensors["image_condition_latent"]54    if tuple(image_condition.shape) != (1, 32, 1, 30, 52):55        raise ValueError(56            "image_condition_latent must be [1,32,1,30,52], got "57            f"{tuple(image_condition.shape)}"58        )59 60    current_txt = tensors["current_txt"]61    cached_txt = tensors["cached_txt"]62    if current_txt.ndim != 3 or current_txt.shape[0] != 1 or current_txt.shape[-1] != HIDDEN_SIZE:63        raise ValueError(f"current_txt must be [1,S,2048], got {tuple(current_txt.shape)}")64    if tuple(cached_txt.shape) != tuple(current_txt.shape):65        raise ValueError(66            f"cached_txt {tuple(cached_txt.shape)} != current_txt {tuple(current_txt.shape)}"67        )68    if tuple(tensors["vec_txt"].shape) != (1, HIDDEN_SIZE):69        raise ValueError(f"vec_txt must be [1,2048], got {tuple(tensors['vec_txt'].shape)}")70 71    for name, tensor in tensors.items():72        if not torch.isfinite(tensor).all():73            raise ValueError(f"Non-finite values in case tensor {name}")74 75 76def validate_chunk_tensors(tensors: Mapping[str, torch.Tensor]) -> None:77    missing = expected_chunk_keys().difference(tensors)78    if missing:79        raise ValueError(f"Missing chunk tensor keys: {sorted(missing)}")80 81    if tuple(tensors["action_labels"].shape) != (1, CHUNK_LATENT_FRAMES):82        raise ValueError(f"Unexpected action_labels shape: {tuple(tensors['action_labels'].shape)}")83    if tuple(tensors["target_viewmats"].shape) != (1, CHUNK_LATENT_FRAMES, 4, 4):84        raise ValueError(f"Unexpected target_viewmats shape: {tuple(tensors['target_viewmats'].shape)}")85    if tuple(tensors["target_Ks"].shape) != (1, CHUNK_LATENT_FRAMES, 3, 3):86        raise ValueError(f"Unexpected target_Ks shape: {tuple(tensors['target_Ks'].shape)}")87 88    for step in range(NUM_STEPS):89        noisy = tensors[f"step_{step}_noisy_sample"]90        hidden = tensors[f"step_{step}_final_hidden"]91        condition = tensors[f"step_{step}_frame_condition"]92        velocity = tensors[f"step_{step}_velocity"]93        timestep = tensors[f"step_{step}_timestep"]94        if tuple(noisy.shape) != (1, 32, 4, 30, 52):95            raise ValueError(f"step {step} noisy shape: {tuple(noisy.shape)}")96        if tuple(hidden.shape) != (1, TOKENS_PER_CHUNK, HIDDEN_SIZE):97            raise ValueError(f"step {step} hidden shape: {tuple(hidden.shape)}")98        if tuple(condition.shape) != (1, CHUNK_LATENT_FRAMES, HIDDEN_SIZE):99            raise ValueError(f"step {step} condition shape: {tuple(condition.shape)}")100        if tuple(velocity.shape) != (1, 32, 4, 30, 52):101            raise ValueError(f"step {step} velocity shape: {tuple(velocity.shape)}")102        if timestep.numel() != 1:103            raise ValueError(f"step {step} timestep must be scalar, got {tuple(timestep.shape)}")104 105    for name, tensor in tensors.items():106        if tensor.is_floating_point() and not torch.isfinite(tensor).all():107            raise ValueError(f"Non-finite values in chunk tensor {name}")108