edit readme

2025-05-05 18:56:00 +08:00
parent 1fb1d20567
commit 4c71c1e8b4
2 changed files with 10 additions and 5 deletions
@@ -1,11 +1,12 @@
 import argparse
 parser = argparse.ArgumentParser()
 parser.add_argument("--checkpoint_path", type=str, default="")
+parser.add_argument("--server_name", type=str, default="127.0.0.1")
 parser.add_argument("--port", type=int, default=7865)
 parser.add_argument("--device_id", type=int, default=0)
 parser.add_argument("--share", type=bool, default=False)
 parser.add_argument("--bf16", type=bool, default=True)
-
+parser.add_argument("--torch_compile", type=bool, default=False)
 args = parser.parse_args()

 import os
@@ -22,7 +23,8 @@ def main(args):

    model_demo = ACEStepPipeline(
        checkpoint_dir=args.checkpoint_path,
-        dtype="bfloat16" if args.bf16 else "float32"
+        dtype="bfloat16" if args.bf16 else "float32",
+        torch_compile=args.torch_compile
    )
    data_sampler = DataSampler()

@@ -31,6 +33,7 @@ def main(args):
        sample_data_func=data_sampler.sample,
    )
    demo.launch(
+        server_name=args.server_name,
        server_port=args.port,
        share=args.share
    )