Added an option to pass enable_single_replica_ckpt_restoring from configs (command line) to enable restoring with single replica then broadcasting. Changed StandardSave and StandardRestore in orbax with PyTreeSave and PyTreeRestore to enable the feature above.
Removed also unnecessary tests (controlnet and dreambooth) covered in other test cases.
Added an option to pass
enable_single_replica_ckpt_restoringfrom configs (command line) to enable restoring with single replica then broadcasting. Changed StandardSave and StandardRestore in orbax with PyTreeSave and PyTreeRestore to enable the feature above.Removed also unnecessary tests (controlnet and dreambooth) covered in other test cases.