Files
mask-ddpm/example/run_all.py
2026-01-24 00:20:50 +08:00

82 lines
2.6 KiB
Python

#!/usr/bin/env python3
"""One-command pipeline runner with config-driven paths."""
import argparse
import json
import subprocess
import sys
from pathlib import Path
from platform_utils import safe_path, is_windows, resolve_path
def run(cmd):
print("running:", " ".join(cmd))
cmd = [safe_path(arg) for arg in cmd]
if is_windows():
subprocess.run(cmd, check=True, shell=False)
else:
subprocess.run(cmd, check=True)
def parse_args():
parser = argparse.ArgumentParser(description="Run prepare -> train -> export -> evaluate.")
base_dir = Path(__file__).resolve().parent
parser.add_argument("--config", default=str(base_dir / "config.json"))
parser.add_argument("--device", default="auto", help="cpu, cuda, or auto")
parser.add_argument("--skip-prepare", action="store_true")
parser.add_argument("--skip-train", action="store_true")
parser.add_argument("--skip-export", action="store_true")
parser.add_argument("--skip-eval", action="store_true")
return parser.parse_args()
def main():
args = parse_args()
base_dir = Path(__file__).resolve().parent
config_path = Path(args.config)
with open(config_path, "r", encoding="utf-8") as f:
cfg = json.load(f)
config_path = resolve_path(base_dir, config_path)
timesteps = cfg.get("timesteps", 200)
seq_len = cfg.get("sample_seq_len", cfg.get("seq_len", 64))
batch_size = cfg.get("sample_batch_size", cfg.get("batch_size", 2))
clip_k = cfg.get("clip_k", 5.0)
if not args.skip_prepare:
run([sys.executable, str(base_dir / "prepare_data.py")])
if not args.skip_train:
run([sys.executable, str(base_dir / "train.py"), "--config", str(config_path), "--device", args.device])
if not args.skip_export:
run(
[
sys.executable,
str(base_dir / "export_samples.py"),
"--include-time",
"--device",
args.device,
"--config",
str(config_path),
"--timesteps",
str(timesteps),
"--seq-len",
str(seq_len),
"--batch-size",
str(batch_size),
"--clip-k",
str(clip_k),
"--use-ema",
]
)
if not args.skip_eval:
ref = cfg.get("data_glob") or cfg.get("data_path") or ""
if ref:
run([sys.executable, str(base_dir / "evaluate_generated.py"), "--reference", str(ref)])
else:
run([sys.executable, str(base_dir / "evaluate_generated.py")])
if __name__ == "__main__":
main()