add add mps support

2025-05-04 20:03:57 +08:00
parent 6cb7c1d894
commit d77c5f42f1
6 changed files with 56 additions and 57 deletions
@@ -4,7 +4,7 @@ parser.add_argument("--checkpoint_path", type=str, default="")
 parser.add_argument("--port", type=int, default=7865)
 parser.add_argument("--device_id", type=int, default=0)
 parser.add_argument("--share", action='store_true', default=False)
-parser.add_argument("--bf16", action='store_true', default=True)
+parser.add_argument("--bf16", action='store_true', default=False)

 args = parser.parse_args()

@@ -19,7 +19,7 @@ from data_sampler import DataSampler


 def main(args):
- 
+
    model_demo = ACEStepPipeline(
        checkpoint_dir=args.checkpoint_path,
        dtype="bfloat16" if args.bf16 else "float32"