| 1 | import json |
| 2 | import sys |
| 3 | from pathlib import Path |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | ROOT = Path(__file__).resolve().parents[2] |
| 9 | WM_ROOT = Path(__file__).resolve().parents[1] |
| 10 | sys.path[:0] = [str(WM_ROOT), str(WM_ROOT / "ltx-core" / "src")] |
| 11 | |
| 12 | from helpers.action_camera import ( |
| 13 | DEFAULT_ROTATION_SPEED_DEG, |
| 14 | DEFAULT_TRANSLATION_SPEED, |
| 15 | default_k_pix, |
| 16 | parse_action_string, |
| 17 | ) |
| 18 | from helpers.action_condition import action_config, build_action_condition |
| 19 | from helpers.moge_fov import effective_fov_x |
| 20 | from helpers.action_overlay import _normalised_rotation, _translation_keys |
| 21 | from ltx_core.model.transformer.transformer import BasicAVTransformerBlock, TransformerConfig |
| 22 | |
| 23 | |
| 24 | def test_action_dsl_parses_combined_and_idle_segments(): |
| 25 | frames = parse_action_string("w-2, wj-1, none-2") |
| 26 | assert frames == [["w"], ["w"], ["j", "w"], [], []] |
| 27 | |
| 28 | |
| 29 | @pytest.mark.parametrize("value", ["", "w", "x-3", "w-0", "w-nope"]) |
| 30 | def test_action_dsl_rejects_invalid_input(value): |
| 31 | with pytest.raises(ValueError): |
| 32 | parse_action_string(value) |
| 33 | |
| 34 | |
| 35 | def test_action_condition_shape_dtype_and_keys(): |
| 36 | condition = build_action_condition( |
| 37 | "w-4", num_frames=17, width=64, height=64, |
| 38 | translation_speed=0.025, rotation_speed_deg=0.6, |
| 39 | pitch_limit_deg=60.0, fov_deg=70.0, device=torch.device("cpu"), fps=24.0, |
| 40 | ) |
| 41 | assert set(condition) == {"ucpe_viewmats", "ucpe_Ks"} |
| 42 | assert condition["ucpe_viewmats"].shape == (1, 3, 4, 4) |
| 43 | assert condition["ucpe_Ks"].shape == (1, 3, 3, 3) |
| 44 | assert all(value.dtype == torch.bfloat16 for value in condition.values()) |
| 45 | |
| 46 | |
| 47 | def test_ucpe_attention_shape_and_device(): |
| 48 | cfg = action_config(width=64, height=64, num_blocks=1) |
| 49 | cfg.ucpe_attn_dim = 32 |
| 50 | cfg.ucpe_num_heads = 2 |
| 51 | block = BasicAVTransformerBlock( |
| 52 | idx=0, num_layers=1, |
| 53 | video=TransformerConfig(dim=64, heads=2, d_head=32, context_dim=64), |
| 54 | ) |
| 55 | block._init_action_params(TransformerConfig(dim=64, heads=2, d_head=32, context_dim=0), cfg) |
| 56 | condition = build_action_condition( |
| 57 | "w-8", num_frames=9, width=64, height=64, |
| 58 | translation_speed=0.025, rotation_speed_deg=0.6, |
| 59 | pitch_limit_deg=60.0, fov_deg=70.0, device=torch.device("cpu"), fps=24.0, |
| 60 | ) |
| 61 | x = torch.randn(1, 8, 64) |
| 62 | out = block._apply_ucpe_attention( |
| 63 | x, condition["ucpe_viewmats"].float(), condition["ucpe_Ks"].float() |
| 64 | ) |
| 65 | assert out.shape == x.shape |
| 66 | assert out.dtype == x.dtype and out.device == x.device |
| 67 | |
| 68 | |
| 69 | def test_ucpe_attention_accepts_fp32_cameras_with_bf16_hidden_states(): |
| 70 | cfg = action_config(width=64, height=64, num_blocks=1) |
| 71 | cfg.ucpe_attn_dim = 32 |
| 72 | cfg.ucpe_num_heads = 2 |
| 73 | block = BasicAVTransformerBlock( |
| 74 | idx=0, num_layers=1, |
| 75 | video=TransformerConfig(dim=64, heads=2, d_head=32, context_dim=64), |
| 76 | ).to(torch.bfloat16) |
| 77 | block._init_action_params(TransformerConfig(dim=64, heads=2, d_head=32, context_dim=0), cfg) |
| 78 | block = block.to(torch.bfloat16) |
| 79 | block.ucpe_prope.coeffs_x_0 = block.ucpe_prope.coeffs_x_1 = None |
| 80 | block.ucpe_prope.coeffs_y_0 = block.ucpe_prope.coeffs_y_1 = None |
| 81 | condition = build_action_condition( |
| 82 | "w-8", num_frames=9, width=64, height=64, |
| 83 | translation_speed=0.025, rotation_speed_deg=0.6, |
| 84 | pitch_limit_deg=60.0, fov_deg=70.0, device=torch.device("cpu"), fps=24.0, |
| 85 | ) |
| 86 | x = torch.randn(1, 8, 64, dtype=torch.bfloat16) |
| 87 | out = block._apply_ucpe_attention( |
| 88 | x, condition["ucpe_viewmats"].float(), condition["ucpe_Ks"].float() |
| 89 | ) |
| 90 | assert out.shape == x.shape and out.dtype == torch.bfloat16 |
| 91 | |
| 92 | |
| 93 | def test_fov_crop_and_default_intrinsics(): |
| 94 | effective, factor = effective_fov_x(90.0, 200, 100, 100, 100) |
| 95 | assert factor == pytest.approx(0.5) |
| 96 | assert effective == pytest.approx(53.130102, rel=1e-6) |
| 97 | K = default_k_pix(1280, 704, 70.0) |
| 98 | assert K[0, 2].item() == 640 and K[1, 2].item() == 352 |
| 99 | |
| 100 | |
| 101 | |
| 102 | def test_prompt_skill_has_required_six_fields(): |
| 103 | skill = (WM_ROOT / "PROMPT_SKILL.md").read_text() |
| 104 | for field in ("Environment:", "Character:", "Style:", "Perspective:", "Sounds:", "Speech:"): |
| 105 | assert field in skill |
| 106 | assert "Do not request subtitles" in skill |
| 107 | |
| 108 | |
| 109 | def test_public_config_has_only_semantic_action_controls(): |
| 110 | text = (WM_ROOT / "configs" / "inference_wm.yaml").read_text() |
| 111 | lowered = text.lower() |
| 112 | assert "translation_speed:" in text |
| 113 | assert "rotation_speed_deg:" in text |
| 114 | assert DEFAULT_TRANSLATION_SPEED > 0 |
| 115 | assert DEFAULT_ROTATION_SPEED_DEG > 0 |
| 116 | assert "global_trans" not in lowered and "normalize_mode" not in lowered |
| 117 | assert "video_cfg: 4.0" in text |
| 118 | assert "audio_cfg: 2.0" in text |
| 119 | assert "negative_prompt:" in text |
| 120 | assert "game UI" in text |
| 121 | assert "crosshair" in text |
| 122 | |
| 123 | |
| 124 | def test_action_overlay_derives_stable_hud_controls(): |
| 125 | trans = __import__("numpy").array([[0.0, 0.0, 0.2], [0.1, 0.0, 0.0]], dtype=float) |
| 126 | keys = _translation_keys(trans) |
| 127 | yaw, pitch = _normalised_rotation(__import__("numpy").zeros((2, 3), dtype=float)) |
| 128 | assert len(keys) == 3 and keys[0] == ["W"] |
| 129 | assert yaw.shape == pitch.shape == (3,) |
| 130 |