chore: import upstream snapshot with attribution

train.py

@@ -0,0 +1,47 @@
+# Adopted from https://github.com/guandeh17/Self-Forcing
+# SPDX-License-Identifier: Apache-2.0
+import argparse
+import os
+from omegaconf import OmegaConf
+import wandb
+
+from trainer import ScoreDistillationTrainer, DiffusionTrainer
+from utils.config import normalize_config
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config_path", type=str, required=True)
+    parser.add_argument("--no_save", action="store_true")
+    parser.add_argument("--no_visualize", action="store_true")
+    parser.add_argument("--logdir", type=str, default="", help="Path to the directory to save logs")
+    parser.add_argument("--wandb-save-dir", type=str, default="", help="Path to the directory to save wandb logs")
+    parser.add_argument("--disable-wandb", action="store_true")
+    parser.add_argument("--no-auto-resume", action="store_true", help="Disable auto resume from latest checkpoint in logdir")
+    parser.add_argument("--generate-before-train", action="store_true", help="Run one evaluation inference before training starts")
+
+    args = parser.parse_args()
+
+    config = normalize_config(OmegaConf.load(args.config_path))
+    config.no_save = args.no_save
+    config.no_visualize = args.no_visualize
+
+    config_name = os.path.splitext(os.path.basename(args.config_path))[0]
+    config.config_name = config_name
+    config.logdir = args.logdir
+    config.wandb_save_dir = args.wandb_save_dir
+    config.disable_wandb = args.disable_wandb
+    config.auto_resume = not args.no_auto_resume  # Default to True unless --no-auto-resume is specified
+    config.generate_before_train = args.generate_before_train
+
+    if config.trainer == "score_distillation":
+        trainer = ScoreDistillationTrainer(config)
+    elif config.trainer == "diffusion":
+        trainer = DiffusionTrainer(config)
+    trainer.train()
+
+    wandb.finish()
+
+
+if __name__ == "__main__":
+    main()