83a7bf3519
* Use magic trailing commas instead of disabling formatting to keep args on separate lines * Scope `fmt: off`s better where that's not possible
33 lines
1009 B
Python
33 lines
1009 B
Python
from mflux.dreambooth.dreambooth import DreamBooth
|
|
from mflux.dreambooth.dreambooth_initializer import DreamBoothInitializer
|
|
from mflux.error.exceptions import StopTrainingException
|
|
from mflux.ui.cli.parsers import CommandLineParser
|
|
|
|
|
|
def main():
|
|
parser = CommandLineParser(description="Finetune a LoRA adapter")
|
|
parser.add_general_arguments()
|
|
parser.add_model_arguments(require_model_arg=False)
|
|
parser.add_training_arguments()
|
|
args = parser.parse_args()
|
|
|
|
flux, runtime_config, training_spec, training_state = DreamBoothInitializer.initialize(
|
|
config_path=args.train_config,
|
|
checkpoint_path=args.train_checkpoint,
|
|
)
|
|
|
|
try:
|
|
DreamBooth.train(
|
|
flux=flux,
|
|
runtime_config=runtime_config,
|
|
training_spec=training_spec,
|
|
training_state=training_state,
|
|
)
|
|
except StopTrainingException as stop_exc:
|
|
training_state.save(training_spec)
|
|
print(stop_exc)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|