Upload train.py with huggingface_hub

train.py

@@ -463,7 +463,7 @@ def main():
     """Main training function"""
     parser = argparse.ArgumentParser(description="Train Byte Dream diffusion model")
     parser.add_argument("--config", type=str, default="config.yaml", help="Path to config file")
-    parser.add_argument("--train_data", type=str, required=True, help="Path to training data")
+    parser.add_argument("--train_data", type=str, default="./dataset", help="Path to training data (default: ./dataset)")
     parser.add_argument("--output_dir", type=str, default="./models/bytedream", help="Output directory")
     parser.add_argument("--resume", type=str, default=None, help="Resume from checkpoint")
     parser.add_argument("--device", type=str, default="cpu", help="Device to train on")