Update app.py

app.py

@@ -4,9 +4,12 @@ from pipeline_ace_step import ACEStepPipeline
 from data_sampler import DataSampler
 import os
 
+# 获取当前脚本的绝对路径，用于构建默认的存储路径
+APP_ROOT = os.path.dirname(os.path.abspath(__file__))
 
 parser = argparse.ArgumentParser()
-parser.add_argument("--checkpoint_path", type=str, default=None)
+# 将 checkpoint_path 的默认值改为应用程序根目录下的 'checkpoints' 文件夹
+parser.add_argument("--checkpoint_path", type=str, default=os.path.join(APP_ROOT, "checkpoints"))
 parser.add_argument("--server_name", type=str, default="0.0.0.0")
 parser.add_argument("--port", type=int, default=7860)
 parser.add_argument("--device_id", type=int, default=0)
@@ -17,16 +20,18 @@ parser.add_argument("--torch_compile", type=bool, default=False)
 args = parser.parse_args()
 os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)
 
-
-persistent_storage_path = "/data"
+# 将 persistent_storage_path 的默认值改为应用程序根目录下的 'persistent_data' 文件夹
+persistent_storage_path = os.path.join(APP_ROOT, "persistent_data")
 
 
 def main(args):
+    print(f"Using checkpoint path: {args.checkpoint_path}")
+    print(f"Using persistent storage path: {persistent_storage_path}")
 
     model_demo = ACEStepPipeline(
         checkpoint_dir=args.checkpoint_path,
         dtype="bfloat16" if args.bf16 else "float32",
-        persistent_storage_path=persistent_storage_path,
+        persistent_storage_path=persistent_storage_path, # 传递修改后的路径
         torch_compile=args.torch_compile
     )
     data_sampler = DataSampler()
@@ -37,9 +42,11 @@ def main(args):
         load_data_func=data_sampler.load_json,
     )
     demo.queue(default_concurrency_limit=8).launch(
-        
+        server_name=args.server_name, # 添加这一行以使用命令行参数
+        server_port=args.port,       # 添加这一行以使用命令行参数
+        share=args.share             # 添加这一行以使用命令行参数
     )
 
 
 if __name__ == "__main__":
-    main(args)
+    main(args)