support flux training

2024-09-06 10:37:28 +08:00
parent 4654aa0cab
commit 416b73b8c0
5 changed files with 218 additions and 46 deletions
--- a/diffsynth/trainers/text_to_image.py
+++ b/diffsynth/trainers/text_to_image.py
@@ -152,7 +152,7 @@ def add_general_parsers(parser):
        "--precision",
        type=str,
        default="16-mixed",
-        choices=["32", "16", "16-mixed"],
+        choices=["32", "16", "16-mixed", "bf16"],
        help="Training precision",
    )
    parser.add_argument(