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