from __future__ import annotations from types import SimpleNamespace import torch from hyvideo.commons.infer_state import initialize_infer_state from models import HYWorldPlayPredictorV2, PredictorV2Config from predictor_data.v2_schema import validate_case_tensors_v2, validate_vision_context_v2 def _init_infer_state() -> None: initialize_infer_state( SimpleNamespace( sage_blocks_range="0-1", use_sageattn=False, enable_torch_compile=False, use_fp8_gemm=False, quant_type="fp8-per-block", include_patterns="double_blocks", use_vae_parallel=False, ) ) def test_v2_zero_init_ar_context_and_backward() -> None: _init_infer_state() cfg = PredictorV2Config( hidden_size=128, heads_num=1, mlp_width_ratio=1.0, attn_mode="torch", source_block_ids=(0, 1), teacher_block_count=2, latent_frames=1, latent_height=2, latent_width=2, ) model = HYWorldPlayPredictorV2(cfg) torch.nn.init.normal_(model.final_layer.linear.weight, std=0.01) model.train() anchor = torch.randn(1, 4, 128) context = { block_id: { "k_txt": torch.randn(1, 1, 3, 128), "v_txt": torch.randn(1, 1, 3, 128), "k_vision": torch.empty(2, 1, 0, 128), "v_vision": torch.empty(2, 1, 0, 128), } for block_id in (0, 1) } output = model( target_model_input=torch.randn(1, 65, 1, 2, 2), anchor_hidden=anchor, target_frame_condition=torch.randn(1, 1, 128), target_viewmats=torch.eye(4).reshape(1, 1, 4, 4), target_Ks=torch.eye(3).reshape(1, 1, 3, 3), rope_temporal_size=1, start_rope_start_idx=0, context_kv=context, ) assert torch.equal(output["pred_hidden"], anchor) assert output["pred_velocity"].shape == (1, 32, 1, 2, 2) loss = output["pred_velocity"].float().square().mean() loss.backward() assert model.residual_out.weight.grad is not None assert model.img_in.proj.weight.grad is None assert model.final_layer.linear.weight.grad is None assert all( not parameter.requires_grad for block in model.double_blocks for name, parameter in block.named_parameters() if name.startswith("txt_") ) def test_v2_bf16_context_schema() -> None: case = {"image_condition_latent": torch.zeros(1, 32, 1, 30, 52, dtype=torch.bfloat16)} for block_id in (0, 1, 52, 53): case[f"block_{block_id:02d}_k_txt"] = torch.zeros( 1, 16, 7, 128, dtype=torch.bfloat16 ) case[f"block_{block_id:02d}_v_txt"] = torch.zeros( 1, 16, 7, 128, dtype=torch.bfloat16 ) case["text_valid_mask"] = torch.ones(1, 7, dtype=torch.bool) validate_case_tensors_v2(case) empty_context = { "k_vision": torch.empty(2, 16, 0, 128, dtype=torch.bfloat16), "v_vision": torch.empty(2, 16, 0, 128, dtype=torch.bfloat16), "context_valid_mask": torch.empty(1, 0, dtype=torch.bool), } assert validate_vision_context_v2(0, empty_context) == 0 def test_v2_padded_masked_batch_matches_compact_single_samples() -> None: _init_infer_state() torch.manual_seed(7) cfg = PredictorV2Config( hidden_size=128, heads_num=1, mlp_width_ratio=1.0, attn_mode="torch", source_block_ids=(0, 1), teacher_block_count=2, latent_frames=1, latent_height=2, latent_width=2, ) model = HYWorldPlayPredictorV2(cfg).eval() torch.nn.init.normal_(model.residual_out.weight, std=0.01) torch.nn.init.normal_(model.residual_out.bias, std=0.01) samples = [] for sample_id, (text_tokens, context_tokens, rope_size, rope_start) in enumerate( ((2, 0, 1, 0), (3, 2, 2, 1)) ): context = {} for block_id in (0, 1): context[block_id] = { "k_txt": torch.randn(1, 1, text_tokens, 128), "v_txt": torch.randn(1, 1, text_tokens, 128), "k_vision": torch.randn(2, 1, context_tokens, 128), "v_vision": torch.randn(2, 1, context_tokens, 128), } samples.append( { "target_model_input": torch.randn(1, 65, 1, 2, 2), "anchor_hidden": torch.randn(1, 4, 128), "target_frame_condition": torch.randn(1, 1, 128), "target_viewmats": torch.eye(4).reshape(1, 1, 4, 4), "target_Ks": torch.eye(3).reshape(1, 1, 3, 3), "rope_temporal_size": rope_size, "start_rope_start_idx": rope_start, "context_kv": context, } ) with torch.no_grad(): singles = [model(**sample) for sample in samples] batched_context = {} for block_id in (0, 1): first, second = (sample["context_kv"][block_id] for sample in samples) def pad_text(tensor: torch.Tensor) -> torch.Tensor: result = tensor.new_zeros(1, 1, 3, 128) result[:, :, : tensor.shape[2]] = tensor return result def pad_vision(tensor: torch.Tensor) -> torch.Tensor: result = tensor.new_zeros(2, 1, 2, 128) result[:, :, : tensor.shape[2]] = tensor return result first_vision_k, second_vision_k = pad_vision(first["k_vision"]), pad_vision( second["k_vision"] ) first_vision_v, second_vision_v = pad_vision(first["v_vision"]), pad_vision( second["v_vision"] ) batched_context[block_id] = { "k_txt": torch.cat([pad_text(first["k_txt"]), pad_text(second["k_txt"])], 0), "v_txt": torch.cat([pad_text(first["v_txt"]), pad_text(second["v_txt"])], 0), "text_valid_mask": torch.tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.bool), "k_vision": torch.cat( [first_vision_k[:1], second_vision_k[:1], first_vision_k[1:], second_vision_k[1:]], 0, ), "v_vision": torch.cat( [first_vision_v[:1], second_vision_v[:1], first_vision_v[1:], second_vision_v[1:]], 0, ), "context_valid_mask": torch.tensor([[0, 0], [1, 1]], dtype=torch.bool), } with torch.no_grad(): batched = model( target_model_input=torch.cat([sample["target_model_input"] for sample in samples]), anchor_hidden=torch.cat([sample["anchor_hidden"] for sample in samples]), target_frame_condition=torch.cat( [sample["target_frame_condition"] for sample in samples] ), target_viewmats=torch.cat([sample["target_viewmats"] for sample in samples]), target_Ks=torch.cat([sample["target_Ks"] for sample in samples]), rope_temporal_size=torch.tensor([1, 2]), start_rope_start_idx=torch.tensor([0, 1]), context_kv=batched_context, ) for index, single in enumerate(singles): torch.testing.assert_close( batched["pred_hidden"][index : index + 1], single["pred_hidden"], rtol=2e-5, atol=2e-5, )