| 1 | import json |
| 2 | import subprocess |
| 3 | import sys |
| 4 | from pathlib import Path |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | import yaml |
| 9 | |
| 10 | WM_ROOT = Path(__file__).resolve().parents[1] |
| 11 | sys.path[:0] = [ |
| 12 | str(WM_ROOT), |
| 13 | str(WM_ROOT / "ltx-core" / "src"), |
| 14 | str(WM_ROOT / "ltx-causal" / "src"), |
| 15 | ] |
| 16 | |
| 17 | from helpers.action_condition import build_causal_action_condition # noqa: E402 |
| 18 | from ltx_core.model.transformer.attention import update_kv_cache # noqa: E402 |
| 19 | from ltx_core.model.transformer.transformer import rebase_viewmat_translation # noqa: E402 |
| 20 | import ltx_causal as causal # noqa: E402 |
| 21 | |
| 22 | |
| 23 | def test_four_step_schedule_has_four_student_steps_and_no_zero(): |
| 24 | sigmas = causal.resolve_causal_sigmas() |
| 25 | assert len(sigmas) == 4 |
| 26 | assert all(a > b for a, b in zip(sigmas, sigmas[1:])) |
| 27 | assert sigmas[-1] > 0 |
| 28 | |
| 29 | |
| 30 | def test_default_241_frame_block_layout_and_audio_mapping(): |
| 31 | blocks = causal.causal_video_blocks(31) |
| 32 | assert blocks == [(0, 1), *[(start, start + 3) for start in range(1, 31, 3)]] |
| 33 | assert causal.causal_audio_frames(31) == 252 |
| 34 | assert causal.causal_audio_blocks(31)[-1] == (227, 252) |
| 35 | cache = causal.CausalCacheConfig() |
| 36 | assert ( |
| 37 | cache.video_local_attn_size, |
| 38 | cache.video_sink_size, |
| 39 | cache.video_chunk_size, |
| 40 | ) == (19, 7, 3) |
| 41 | assert (cache.audio_local_attn_size, cache.audio_sink_size) == (152, 52) |
| 42 | |
| 43 | |
| 44 | def test_audio_cache_sizes_follow_video_cache_alignment(): |
| 45 | cache = causal.CausalCacheConfig( |
| 46 | video_local_attn_size=25, |
| 47 | video_sink_size=7, |
| 48 | ) |
| 49 | cache.validate() |
| 50 | assert (cache.audio_local_attn_size, cache.audio_sink_size) == (202, 52) |
| 51 | |
| 52 | invalid = causal.CausalCacheConfig(video_local_attn_size=20) |
| 53 | with pytest.raises(ValueError, match="audio alignment"): |
| 54 | invalid.validate() |
| 55 | |
| 56 | |
| 57 | def test_flash_rejects_unsupported_video_chunk_size(): |
| 58 | cache = causal.CausalCacheConfig(video_chunk_size=4) |
| 59 | with pytest.raises(ValueError, match="requires video_chunk_size=3"): |
| 60 | cache.validate() |
| 61 | with pytest.raises(ValueError, match="requires video_chunk_size=3"): |
| 62 | causal.causal_video_blocks(31, chunk_size=0) |
| 63 | with pytest.raises(ValueError, match="latent video length"): |
| 64 | causal.causal_audio_frames(30) |
| 65 | |
| 66 | |
| 67 | def test_sink_plus_fifo_cache_rollover_and_block_replacement(): |
| 68 | cache = { |
| 69 | "k": torch.zeros(1, 7, 1), "v": torch.zeros(1, 7, 1), |
| 70 | "positions": torch.full((7,), -1, dtype=torch.long), "length": 0, |
| 71 | "local_attn_size": 7, "sink_tokens": 2, |
| 72 | } |
| 73 | with torch.no_grad(): |
| 74 | for start in (0, 2, 5, 8): |
| 75 | values = torch.arange(start, start + 3).view(1, 3, 1).float() |
| 76 | update_kv_cache(cache, start, values, values) |
| 77 | assert cache["positions"][: cache["length"]].tolist() == [0, 1, 6, 7, 8, 9, 10] |
| 78 | replacement = torch.full((1, 3, 1), 99.0) |
| 79 | with torch.no_grad(): |
| 80 | active_k, _ = update_kv_cache(cache, 8, replacement, replacement) |
| 81 | assert active_k[0, -3:, 0].tolist() == [99.0, 99.0, 99.0] |
| 82 | |
| 83 | |
| 84 | def test_bounded_anchor_translation_preserves_relative_camera_transform(): |
| 85 | angle = torch.tensor(0.7) |
| 86 | rotation = torch.tensor([ |
| 87 | [torch.cos(angle), 0.0, torch.sin(angle)], |
| 88 | [0.0, 1.0, 0.0], |
| 89 | [-torch.sin(angle), 0.0, torch.cos(angle)], |
| 90 | ]) |
| 91 | cameras = torch.eye(4).repeat(1, 2, 1, 1) |
| 92 | cameras[0, 0, :3, 3] = torch.tensor([3.0, 1.0, -2.0]) |
| 93 | cameras[0, 1, :3, :3] = rotation |
| 94 | cameras[0, 1, :3, 3] = torch.tensor([-1.0, 4.0, 2.0]) |
| 95 | before = cameras[:, 0] @ torch.linalg.inv(cameras[:, 1]) |
| 96 | rebased = rebase_viewmat_translation(cameras, cameras[:, :1]) |
| 97 | after = rebased[:, 0] @ torch.linalg.inv(rebased[:, 1]) |
| 98 | torch.testing.assert_close(before, after, atol=1e-5, rtol=1e-5) |
| 99 | torch.testing.assert_close(rebased[:, 0, :3, 3], torch.zeros(1, 3), atol=1e-6, rtol=0) |
| 100 | |
| 101 | |
| 102 | def test_causal_action_path_keeps_fp32_cameras(): |
| 103 | condition = build_causal_action_condition( |
| 104 | "wj-8", num_frames=9, width=64, height=64, |
| 105 | translation_speed=0.025, rotation_speed_deg=0.6, |
| 106 | pitch_limit_deg=60.0, fov_deg=70.0, |
| 107 | device=torch.device("cpu"), fps=24.0, |
| 108 | ) |
| 109 | assert all(value.dtype == torch.float32 for value in condition.values()) |
| 110 | |
| 111 | |
| 112 | def test_causal_cli_has_no_cfg_and_accepts_both_cache_flag_spellings(): |
| 113 | source = (WM_ROOT / "inference_wm_causal.py").read_text(encoding="utf-8") |
| 114 | assert "negative-prompt" not in source and "video-cfg" not in source and "audio-cfg" not in source |
| 115 | for spelling in ( |
| 116 | "--video-local-attn-size", "--video_local_attn_size", |
| 117 | "--video-sink-size", "--video_sink_size", |
| 118 | "--video-chunk-size", "--video_chunk_size", |
| 119 | ): |
| 120 | assert spelling in source |
| 121 | pipeline = (WM_ROOT / "ltx-pipelines" / "src" / "ltx_pipelines" / "causal_ti2vid.py").read_text() |
| 122 | assert "encode_prompts([prompt]" in pipeline |
| 123 | |
| 124 | |
| 125 | def test_causal_case_runner_dry_run_uses_causal_entrypoint(): |
| 126 | result = subprocess.run( |
| 127 | [sys.executable, str(WM_ROOT / "scripts" / "run_wm_case_causal.py"), |
| 128 | "--case", "examples/wm_causal_cases/0079", "--dry-run"], |
| 129 | cwd=WM_ROOT, check=True, capture_output=True, text=True, |
| 130 | ) |
| 131 | assert "inference_wm_causal.py" in result.stdout |
| 132 | assert "echo-wm-flash.safetensors" in result.stdout |
| 133 | assert "--num-frames 385" in result.stdout |
| 134 | assert "--fov-deg 70.0" in result.stdout |
| 135 | assert "--translation-speed" not in result.stdout |
| 136 | assert "--rotation-speed-deg" not in result.stdout |
| 137 | assert "--pitch-limit-deg" not in result.stdout |
| 138 | assert "--video-local-attn-size" not in result.stdout |
| 139 | assert "--action-overlay" in result.stdout |
| 140 | |
| 141 | no_overlay = subprocess.run( |
| 142 | [sys.executable, str(WM_ROOT / "scripts" / "run_wm_case_causal.py"), |
| 143 | "--case", "examples/wm_causal_cases/0079", "--dry-run", "--no-action-overlay"], |
| 144 | cwd=WM_ROOT, check=True, capture_output=True, text=True, |
| 145 | ) |
| 146 | assert "--no-action-overlay" in no_overlay.stdout |
| 147 | |
| 148 | |
| 149 | def test_causal_multigpu_runner_uses_active_python_environment(): |
| 150 | source = (WM_ROOT / "scripts" / "run_wm_causal_cases_multigpu.sh").read_text() |
| 151 | assert 'python_bin="${PYTHON_BIN:-python}"' in source |
| 152 | assert 'case "${ACTION_OVERLAY-1}" in' in source |
| 153 | |
| 154 | |
| 155 | def test_checked_in_wbench_causal_cases_have_four_4_second_actions(): |
| 156 | expected_actions = { |
| 157 | "0024": "l-96,w-96,l-96,w-96", |
| 158 | "0075": "w-96,w-96,w-96,w-96", |
| 159 | "0079": "l-96,l-96,l-96,l-96", |
| 160 | "0081": "k-96,i-96,s-96,w-96", |
| 161 | "0122": "a-96,d-96,a-96,d-96", |
| 162 | "0170": "w-96,s-96,a-96,l-96", |
| 163 | } |
| 164 | expected_prompts = { |
| 165 | "0024": "Ancient Roman stone ruins on a sunny afternoon. A pale stone road runs between rows of broken Corinthian columns of varying heights. Arched stone doorways and remnants of walls stand among the columns. Scattered stone fragments and rubble lie on the ground. Strong afternoon sunlight from the upper right casts dramatic long shadows of the columns onto the stone floor. Dry scrubland and a deep blue sky form the background. To the right, beyond the nearest row of columns, the base of a partially collapsed temple platform with worn steps leads up to a pair of standing columns still supporting a fragment of entablature. Further along the road ahead, a large fallen capital block lies on its side in the path. Behind the viewpoint, the stone road extends back toward a reconstructed archway with carved Latin inscriptions above the opening. A few scattered wild poppies grow between the rubble to the left. A tourist in a straw hat, light khaki shirt, and brown cargo pants with a camera hanging on a neck strap. Initially stationary. Adjusts camera strap slightly, glances around softly. Third-person view from directly behind the tourist at a mid-height angle, following the figure as the central subject between the ancient columns.", |
| 166 | "0075": "A cherry blossom garden in full spring bloom. A stone path winds between rows of sakura trees with dense pink blossoms. Pink petals drift through the air. A large red torii gate stands ahead framing the path. A stone lantern sits on the left edge. People stroll in the background. Warm afternoon sunlight filters through the canopy. Beyond the torii gate, a small wooden shrine with a sloped roof and offering box is nestled among the cherry trees. A second stone lantern stands on the right side further along the path, and a wooden bench sits beneath a particularly large sakura tree past the gate. First-person viewer. First-person view with both hands holding a digital camera with a large rear LCD screen showing the live viewfinder image of the scene ahead.", |
| 167 | "0079": "An enchanted crystal cave with massive prismatic crystal formations in purple, teal, and pink. Bioluminescent fungi glow on the cave floor and walls. Floating light motes drift through the air. The crystals refract light into rainbow spectra. To the right, a large crystalline cave monster with glowing purple eyes lurks behind tall crystal clusters. Deep cavern atmosphere with ethereal luminescence. Further to the right beyond the monster, a subterranean crystal pool glows with turquoise light, fed by a thin waterfall dripping from a stalactite cluster. The cave opens into a wider chamber with an ancient stone altar covered in glowing runes. First-person viewer. First-person view with the right hand holding a twisted wooden magic wand topped with a bright blue-white crystal orb that radiates light. The wand rotates together with the viewer's perspective when turning.", |
| 168 | "0081": "A sunlit artist's studio with exposed brick walls and wooden shelves holding art supplies. A canvas sits on a wooden easel showing a half-finished landscape painting. Paint tubes, brushes, and palettes are scattered on a worktable. A tall arched window lets in bright natural light, with potted plants and ferns on the sill. Framed sketches hang on the walls. Above, exposed wooden ceiling beams support a hanging pendant lamp and a dried flower wreath. Below the worktable, paint-stained rags, a jar of turpentine, and stacked canvases lean against the wall. Behind the viewer, a cluttered bookshelf holds art reference books, a ceramic coffee mug, and a small plaster bust. First-person viewer. First-person view with the right hand holding a wooden paintbrush tipped with blue paint, extending toward the canvas.", |
| 169 | "0122": "A volcanic crater rim in CG rendered style with dramatic fire lighting. Dark volcanic rock with glowing orange cracks in the foreground, a volcanic crater filled with bright molten lava below, ash clouds and volcanic steam vents hissing above. The atmosphere is intense with fire and heat. To the left along the rim, jagged obsidian spires jut upward and a collapsed lava tube opening is visible. To the right, the rim widens into a rocky plateau with scattered volcanic boulders and a secondary steam vent column rising from a fissure. A small dragon hatchling with red-orange scales and tiny wings, on the volcanic rim. Lifts off with wings flapping, moving through the air. Third-person perspective. The camera is positioned behind and slightly above the dragon hatchling, following it along the crater rim.", |
| 170 | "0170": "A grand mythological hall rendered in oil-painting style with visible brushstroke textures. Tall marble pillars with ornate Corinthian capitals line both sides of a wide corridor. Voluminous golden clouds billow between and beyond the pillars, filling the background with ethereal light. A large luminous archway glows at the far end of the hall. The floor is polished marble, and the overall palette is warm gold, cream, and blue, evoking a classical Renaissance or Baroque painting of Mount Olympus. Behind the viewer, additional marble pillars recede into a second chamber with a vaulted ceiling painted with celestial figures. To the left, the spaces between pillars open onto a cloud-filled void with distant mountain peaks below. To the right, a stone balustrade overlooks a vast golden cloudscape lit by an unseen divine light source. The luminous archway ahead leads deeper into the divine realm. A god-like male figure seen from behind, wearing flowing robes in deep blue and gold that trail behind him with heavy fabric dynamics. He has bare feet and walks forward with a purposeful stride. A faint halo or nimbus of light encircles his head. The brushstroke texture of the oil-painting style is visible on his robes and skin. Third-person rear-follow camera positioned directly behind the figure at mid-torso height. The figure is centered in the lower portion of the frame, walking forward through the hall of pillars toward the glowing archway. The camera tracks at a stable distance, maintaining the painterly composition.", |
| 171 | } |
| 172 | root = WM_ROOT / "examples" / "wm_causal_cases" |
| 173 | for name, action in expected_actions.items(): |
| 174 | case = json.loads((root / name / "case.json").read_text(encoding="utf-8")) |
| 175 | assert set(case) == {"prompt", "action", "fov_deg", "seed"} |
| 176 | assert (root / name / "input.jpg").is_file() |
| 177 | assert case["action"] == action |
| 178 | assert case["seed"] == 42 |
| 179 | assert all(segment.endswith("-96") for segment in action.split(",")) |
| 180 | assert case["prompt"] == expected_prompts[name] |
| 181 | |
| 182 | |
| 183 | def test_causal_config_has_wbench_camera_and_bounded_cache_defaults(): |
| 184 | config = yaml.safe_load((WM_ROOT / "configs" / "inference_wm_causal.yaml").read_text()) |
| 185 | assert config["model"]["checkpoint"] == "checkpoints/echo-wm-flash.safetensors" |
| 186 | assert config["action"] == { |
| 187 | "enabled": True, |
| 188 | "ucpe": True, |
| 189 | "translation_speed": 0.05, |
| 190 | "rotation_speed_deg": 0.4, |
| 191 | "pitch_limit_deg": 40.0, |
| 192 | "fov_deg": 70.0, |
| 193 | } |
| 194 | assert config["causal"] == { |
| 195 | "timesteps": [1000, 750, 500, 250], |
| 196 | "video_local_attn_size": 19, |
| 197 | "video_sink_size": 7, |
| 198 | "video_chunk_size": 3, |
| 199 | } |
| 200 |