support reference image

2026-03-18 22:08:13 +00:00 · 2025-04-11 16:23:14 +08:00
parent a3b4f235a0
commit 7a06a58f49
1 changed files with 1 additions and 0 deletions
--- a/train_flux_reference.py
+++ b/train_flux_reference.py
@@ -212,6 +212,7 @@ if __name__ == '__main__':
            ),
        ],
        dataset_weight=(4, 2, 2, 1),
+        steps_per_epoch=args.steps_per_epoch,
    )
    train_loader = torch.utils.data.DataLoader(
        dataset,