| |
| """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", |
| }, |
| }, |
| } |
|
|
| |
| |
| |
| 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() |
|
|
