Files
mask-ddpm/example/run_pipeline.py

51 lines
1.3 KiB
Python

#!/usr/bin/env python3
"""One-click pipeline: prepare -> train -> export -> evaluate -> plot loss."""
import argparse
import subprocess
import sys
from pathlib import Path
def run(cmd):
print("running:", " ".join(cmd))
subprocess.run(cmd, check=True)
def parse_args():
parser = argparse.ArgumentParser(description="Run full HAI pipeline.")
base_dir = Path(__file__).resolve().parent
parser.add_argument(
"--config",
default=str(base_dir / "config.json"),
help="Path to training config JSON",
)
parser.add_argument(
"--device",
default="auto",
help="cpu, cuda, or auto (used for export_samples.py)",
)
return parser.parse_args()
def main():
args = parse_args()
base_dir = Path(__file__).resolve().parent
run([sys.executable, str(base_dir / "prepare_data.py")])
run([sys.executable, str(base_dir / "train.py"), "--config", args.config])
run(
[
sys.executable,
str(base_dir / "export_samples.py"),
"--include-time",
"--device",
args.device,
]
)
run([sys.executable, str(base_dir / "evaluate_generated.py")])
run([sys.executable, str(base_dir / "plot_loss.py")])
if __name__ == "__main__":
main()