#!/usr/bin/env python3 """Check the cleaned release layout, public recipes, and source syntax. This check is self-contained. It does not require the historical Git commit used by the earlier readability-only comparison. """ from __future__ import annotations import ast import importlib.util from pathlib import Path import subprocess import sys import yaml ROOT = Path(__file__).resolve().parents[1] sys.dont_write_bytecode = True RELEASES = { "1.3B": { "recipe": "three_encoder.yaml", "required": ( "code/model/generation.py", "code/model/distribution_matching.py", "code/trainer/distribution_matching.py", "code/trainer/feature_gradient_replay.py", "code/losses/multi_encoder_mmd.py", "code/losses/video_nystrom_mmd.py", "code/video_fd/replay/replay_video_encoder_vjp.py", ), "removed": ( "code/losses/replay_video_fd.py", "code/losses/video_fd.py", "code/losses/video_nystrom_reference_bank.py", "code/model/fd.py", "code/model/frame_fd_head.py", "code/trainer/fresh_frame_fd_gradcache.py", "code/trainer/two_level_fidel.py", "code/utils/chunked_vae_decode.py", "code/utils/fd_feature_grad.py", "code/utils/video_chunking.py", "code/utils/video_feature_extract.py", ), "expected": { "total_batch_size": 256, "expected_current_teacher_samples": 128, "fresh_video_local_samples": 64, "fresh_video_generator_grad_local_samples": 32, "triple_dinov3_weight": 10.0, "triple_vjepa2_weight": 1.0, "triple_videomae_v2_weight": 1.0, "triple_aggregation_divisor": 3.0, "max_steps": 150, "generator_lr_cosine_end_step": 125, "distribution_loss": "wan_kv_mmd", }, }, "14B": { "recipe": "train.yaml", "required": ( "code/model/generation_base.py", "code/model/distribution_matching.py", "code/trainer/distribution_matching.py", "code/trainer/feature_gradient_replay.py", "code/losses/multi_encoder_mmd.py", "code/losses/video_nystrom_mmd.py", "code/video_fd/replay/replay_video_encoder_vjp.py", ), "removed": ( "code/losses/replay_video_fd.py", "code/losses/video_fd.py", "code/losses/video_nystrom_reference_bank.py", "code/model/fd.py", "code/model/frame_fd_head.py", "code/tools/precompute_frame_fd_stats.py", "code/tools/precompute_vjepa2_fidel_reference_bank.py", "code/trainer/fresh_frame_fd_gradcache.py", "code/trainer/two_level_fidel.py", "code/utils/chunked_vae_decode.py", "code/utils/fd_feature_grad.py", "code/utils/video_chunking.py", "code/utils/video_feature_extract.py", ), "expected": { "total_batch_size": 256, "expected_current_teacher_samples": 128, "fresh_video_local_samples": 32, "fresh_video_generator_grad_local_samples": 16, "triple_active_branches": [ "dinov3", "vjepa2", "videomae_v2"], "triple_dinov3_weight": 10.0, "triple_vjepa2_weight": 1.0, "triple_videomae_v2_weight": 1.0, "triple_aggregation_divisor": 3.0, "max_steps": 80, "lr_schedule_total_steps": 200, "distribution_loss": "wan_kv_mmd", }, }, } # These names belonged only to the removed auxiliary objectives. Historical # checkpoint names such as fd_state_dict and the active video_fd encoder package # deliberately remain allowed. REMOVED_TOKENS = ( "frame_fd_", "fidel_", "fd_repr_", "fd_guidance_scale", "fresh_frame_fd_gradcache", "losses.replay_video_fd", "losses.video_fd", "from model.fd import", "import model.fd", "model.frame_fd_head", "precompute_frame_fd_stats", "precompute_vjepa2_fidel_reference_bank", "trainer.two_level_fidel", "utils.chunked_vae_decode", "utils.fd_feature_grad", "utils.video_chunking", "utils.video_feature_extract", ) REMOVED_CONFIG_KEYS = { "fd_decode_frames", "fd_ema_beta", "fd_exclude_first_frame", "fd_feature_aggregation", "fd_fid_norm_eps", "fd_guidance_scale", "fd_num_frames", "fd_queue_size", } REMOVED_CONFIG_PREFIXES = ("frame_fd_", "fidel_", "fd_repr_") def load_config_helper(release: str): path = ROOT / release / "code/utils/config.py" spec = importlib.util.spec_from_file_location( f"release_config_{release.replace('.', '_')}", path) module = importlib.util.module_from_spec(spec) assert spec.loader is not None spec.loader.exec_module(module) return module def nested_keys(value): if isinstance(value, dict): for key, child in value.items(): yield str(key) yield from nested_keys(child) elif isinstance(value, list): for child in value: yield from nested_keys(child) def check_release(release: str, settings: dict) -> tuple[int, int]: release_root = ROOT / release for relative in settings["required"]: assert (release_root / relative).is_file(), ( f"missing active release file: {release}/{relative}") for relative in settings["removed"]: assert not (release_root / relative).exists(), ( f"removed auxiliary-loss file is still present: " f"{release}/{relative}") source_paths = sorted((release_root / "code").rglob("*.py")) for path in source_paths: text = path.read_text() ast.parse(text, filename=str(path.relative_to(ROOT))) for token in REMOVED_TOKENS: assert token not in text, ( f"stale auxiliary-loss reference {token!r} in " f"{path.relative_to(ROOT)}") shell_paths = sorted((release_root / "scripts").rglob("*.sh")) for path in shell_paths: subprocess.run(["bash", "-n", str(path)], check=True) yaml_paths = sorted((release_root / "configs").rglob("*.yaml")) yaml_paths += sorted((release_root / "code/configs").rglob("*.yaml")) for path in yaml_paths: payload = yaml.safe_load(path.read_text()) stale = sorted({ key for key in nested_keys(payload) if key in REMOVED_CONFIG_KEYS or key.startswith(REMOVED_CONFIG_PREFIXES) }) assert not stale, ( f"stale auxiliary-loss config keys in {path.relative_to(ROOT)}: " f"{', '.join(stale)}") recipe_path = release_root / "configs" / settings["recipe"] public = yaml.safe_load(recipe_path.read_text()) resolved = load_config_helper(release).resolve_config_aliases(public) for key, expected in settings["expected"].items(): actual = resolved.get(key) assert actual == expected, ( f"{recipe_path.relative_to(ROOT)}: {key}={actual!r}, " f"expected {expected!r}") print( f"PASS {release}: active three-encoder recipe, cleaned layout, " f"{len(source_paths)} Python files and {len(shell_paths)} shell files") return len(source_paths), len(shell_paths) def main() -> None: python_count = 0 shell_count = 0 for release, settings in RELEASES.items(): parsed, checked_shell = check_release(release, settings) python_count += parsed shell_count += checked_shell root_scripts = sorted((ROOT / "scripts").glob("*.py")) for path in root_scripts: ast.parse(path.read_text(), filename=str(path.relative_to(ROOT))) print( f"PASS root scripts: {len(root_scripts)} Python files parse; " f"total release syntax checks {python_count} Python / " f"{shell_count} shell files") if __name__ == "__main__": main()