KUI71/ACE-Step

1import argparse2from ui.components import create_main_demo_ui3from pipeline_ace_step import ACEStepPipeline4from data_sampler import DataSampler5import os6 7 8parser = argparse.ArgumentParser()9parser.add_argument("--checkpoint_path", type=str, default=None)10parser.add_argument("--server_name", type=str, default="0.0.0.0")11parser.add_argument("--port", type=int, default=7860)12parser.add_argument("--device_id", type=int, default=0)13parser.add_argument("--share", action='store_true', default=False)14parser.add_argument("--bf16", action='store_true', default=True)15parser.add_argument("--torch_compile", type=bool, default=False)16 17args = parser.parse_args()18os.environ["CUDA_VISIBLE_DEVICES"] = str(args.device_id)19 20 21persistent_storage_path = "/data"22 23 24def main(args):25 26    model_demo = ACEStepPipeline(27        checkpoint_dir=args.checkpoint_path,28        dtype="bfloat16" if args.bf16 else "float32",29        persistent_storage_path=persistent_storage_path,30        torch_compile=args.torch_compile31    )32    data_sampler = DataSampler()33 34    demo = create_main_demo_ui(35        text2music_process_func=model_demo.__call__,36        sample_data_func=data_sampler.sample,37        load_data_func=data_sampler.load_json,38    )39    demo.queue(default_concurrency_limit=8).launch(40        41    )42 43 44if __name__ == "__main__":45    main(args)46