diff --git a/examples/controlnet/train_controlnet.py b/examples/controlnet/train_controlnet.py index 2b972d1db7..22af44445e 100644 --- a/examples/controlnet/train_controlnet.py +++ b/examples/controlnet/train_controlnet.py @@ -785,16 +785,17 @@ def main(args): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - i = len(weights) - 1 + if accelerator.is_main_process: + i = len(weights) - 1 - while len(weights) > 0: - weights.pop() - model = models[i] + while len(weights) > 0: + weights.pop() + model = models[i] - sub_dir = "controlnet" - model.save_pretrained(os.path.join(output_dir, sub_dir)) + sub_dir = "controlnet" + model.save_pretrained(os.path.join(output_dir, sub_dir)) - i -= 1 + i -= 1 def load_model_hook(models, input_dir): while len(models) > 0: diff --git a/examples/controlnet/train_controlnet_sdxl.py b/examples/controlnet/train_controlnet_sdxl.py index 40c77ca56d..54a9ddd23d 100644 --- a/examples/controlnet/train_controlnet_sdxl.py +++ b/examples/controlnet/train_controlnet_sdxl.py @@ -840,16 +840,17 @@ def main(args): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - i = len(weights) - 1 + if accelerator.is_main_process: + i = len(weights) - 1 - while len(weights) > 0: - weights.pop() - model = models[i] + while len(weights) > 0: + weights.pop() + model = models[i] - sub_dir = "controlnet" - model.save_pretrained(os.path.join(output_dir, sub_dir)) + sub_dir = "controlnet" + model.save_pretrained(os.path.join(output_dir, sub_dir)) - i -= 1 + i -= 1 def load_model_hook(models, input_dir): while len(models) > 0: diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index fef162a83f..2e3a5baabf 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -920,12 +920,13 @@ def main(args): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - for model in models: - sub_dir = "unet" if isinstance(model, type(accelerator.unwrap_model(unet))) else "text_encoder" - model.save_pretrained(os.path.join(output_dir, sub_dir)) + if accelerator.is_main_process: + for model in models: + sub_dir = "unet" if isinstance(model, type(accelerator.unwrap_model(unet))) else "text_encoder" + model.save_pretrained(os.path.join(output_dir, sub_dir)) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): while len(models) > 0: diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py index 8b46d2143b..394f19ed65 100644 --- a/examples/dreambooth/train_dreambooth_lora.py +++ b/examples/dreambooth/train_dreambooth_lora.py @@ -894,27 +894,28 @@ def main(args): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - # there are only two options here. Either are just the unet attn processor layers - # or there are the unet and text encoder atten layers - unet_lora_layers_to_save = None - text_encoder_lora_layers_to_save = None + if accelerator.is_main_process: + # there are only two options here. Either are just the unet attn processor layers + # or there are the unet and text encoder atten layers + unet_lora_layers_to_save = None + text_encoder_lora_layers_to_save = None - for model in models: - if isinstance(model, type(accelerator.unwrap_model(unet))): - unet_lora_layers_to_save = unet_attn_processors_state_dict(model) - elif isinstance(model, type(accelerator.unwrap_model(text_encoder))): - text_encoder_lora_layers_to_save = text_encoder_lora_state_dict(model) - else: - raise ValueError(f"unexpected save model: {model.__class__}") + for model in models: + if isinstance(model, type(accelerator.unwrap_model(unet))): + unet_lora_layers_to_save = unet_attn_processors_state_dict(model) + elif isinstance(model, type(accelerator.unwrap_model(text_encoder))): + text_encoder_lora_layers_to_save = text_encoder_lora_state_dict(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() - LoraLoaderMixin.save_lora_weights( - output_dir, - unet_lora_layers=unet_lora_layers_to_save, - text_encoder_lora_layers=text_encoder_lora_layers_to_save, - ) + LoraLoaderMixin.save_lora_weights( + output_dir, + unet_lora_layers=unet_lora_layers_to_save, + text_encoder_lora_layers=text_encoder_lora_layers_to_save, + ) def load_model_hook(models, input_dir): unet_ = None diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py index 247d111c06..88ae8e4359 100644 --- a/examples/dreambooth/train_dreambooth_lora_sdxl.py +++ b/examples/dreambooth/train_dreambooth_lora_sdxl.py @@ -798,31 +798,32 @@ def main(args): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - # there are only two options here. Either are just the unet attn processor layers - # or there are the unet and text encoder atten layers - unet_lora_layers_to_save = None - text_encoder_one_lora_layers_to_save = None - text_encoder_two_lora_layers_to_save = None + if accelerator.is_main_process: + # there are only two options here. Either are just the unet attn processor layers + # or there are the unet and text encoder atten layers + unet_lora_layers_to_save = None + text_encoder_one_lora_layers_to_save = None + text_encoder_two_lora_layers_to_save = None - for model in models: - if isinstance(model, type(accelerator.unwrap_model(unet))): - unet_lora_layers_to_save = unet_attn_processors_state_dict(model) - elif isinstance(model, type(accelerator.unwrap_model(text_encoder_one))): - text_encoder_one_lora_layers_to_save = text_encoder_lora_state_dict(model) - elif isinstance(model, type(accelerator.unwrap_model(text_encoder_two))): - text_encoder_two_lora_layers_to_save = text_encoder_lora_state_dict(model) - else: - raise ValueError(f"unexpected save model: {model.__class__}") + for model in models: + if isinstance(model, type(accelerator.unwrap_model(unet))): + unet_lora_layers_to_save = unet_attn_processors_state_dict(model) + elif isinstance(model, type(accelerator.unwrap_model(text_encoder_one))): + text_encoder_one_lora_layers_to_save = text_encoder_lora_state_dict(model) + elif isinstance(model, type(accelerator.unwrap_model(text_encoder_two))): + text_encoder_two_lora_layers_to_save = text_encoder_lora_state_dict(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() - StableDiffusionXLPipeline.save_lora_weights( - output_dir, - unet_lora_layers=unet_lora_layers_to_save, - text_encoder_lora_layers=text_encoder_one_lora_layers_to_save, - text_encoder_2_lora_layers=text_encoder_two_lora_layers_to_save, - ) + StableDiffusionXLPipeline.save_lora_weights( + output_dir, + unet_lora_layers=unet_lora_layers_to_save, + text_encoder_lora_layers=text_encoder_one_lora_layers_to_save, + text_encoder_2_lora_layers=text_encoder_two_lora_layers_to_save, + ) def load_model_hook(models, input_dir): unet_ = None diff --git a/examples/instruct_pix2pix/train_instruct_pix2pix.py b/examples/instruct_pix2pix/train_instruct_pix2pix.py index fc0a347e37..afd5bd355e 100644 --- a/examples/instruct_pix2pix/train_instruct_pix2pix.py +++ b/examples/instruct_pix2pix/train_instruct_pix2pix.py @@ -485,14 +485,15 @@ def main(): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - if args.use_ema: - ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) + if accelerator.is_main_process: + if args.use_ema: + ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "unet")) + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): if args.use_ema: diff --git a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py index b7ac105c24..2c7d839e43 100644 --- a/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py +++ b/examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py @@ -528,14 +528,15 @@ def main(): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - if args.use_ema: - ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) + if accelerator.is_main_process: + if args.use_ema: + ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "unet")) + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): if args.use_ema: diff --git a/examples/research_projects/controlnet/train_controlnet_webdataset.py b/examples/research_projects/controlnet/train_controlnet_webdataset.py index 9d732abfc4..3122a3952b 100644 --- a/examples/research_projects/controlnet/train_controlnet_webdataset.py +++ b/examples/research_projects/controlnet/train_controlnet_webdataset.py @@ -1010,16 +1010,17 @@ def main(args): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - i = len(weights) - 1 + if accelerator.is_main_process: + i = len(weights) - 1 - while len(weights) > 0: - weights.pop() - model = models[i] + while len(weights) > 0: + weights.pop() + model = models[i] - sub_dir = "controlnet" - model.save_pretrained(os.path.join(output_dir, sub_dir)) + sub_dir = "controlnet" + model.save_pretrained(os.path.join(output_dir, sub_dir)) - i -= 1 + i -= 1 def load_model_hook(models, input_dir): while len(models) > 0: diff --git a/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py b/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py index 0cd915b423..2548c3a286 100644 --- a/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py +++ b/examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py @@ -552,14 +552,15 @@ def main(): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - if args.use_ema: - ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) + if accelerator.is_main_process: + if args.use_ema: + ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "unet")) + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): if args.use_ema: diff --git a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py index 12ff40bbd6..ba5ccd238f 100644 --- a/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py +++ b/examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py @@ -313,14 +313,15 @@ def main(args): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - if args.use_ema: - ema_model.save_pretrained(os.path.join(output_dir, "unet_ema")) + if accelerator.is_main_process: + if args.use_ema: + ema_model.save_pretrained(os.path.join(output_dir, "unet_ema")) - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "unet")) + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): if args.use_ema: diff --git a/examples/text_to_image/train_text_to_image.py b/examples/text_to_image/train_text_to_image.py index 96b5014760..542ee61de2 100644 --- a/examples/text_to_image/train_text_to_image.py +++ b/examples/text_to_image/train_text_to_image.py @@ -629,14 +629,15 @@ def main(): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - if args.use_ema: - ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) + if accelerator.is_main_process: + if args.use_ema: + ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "unet")) + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): if args.use_ema: diff --git a/examples/text_to_image/train_text_to_image_lora_sdxl.py b/examples/text_to_image/train_text_to_image_lora_sdxl.py index aa69c822ca..a01e68340b 100644 --- a/examples/text_to_image/train_text_to_image_lora_sdxl.py +++ b/examples/text_to_image/train_text_to_image_lora_sdxl.py @@ -669,31 +669,32 @@ def main(args): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - # there are only two options here. Either are just the unet attn processor layers - # or there are the unet and text encoder atten layers - unet_lora_layers_to_save = None - text_encoder_one_lora_layers_to_save = None - text_encoder_two_lora_layers_to_save = None + if accelerator.is_main_process: + # there are only two options here. Either are just the unet attn processor layers + # or there are the unet and text encoder atten layers + unet_lora_layers_to_save = None + text_encoder_one_lora_layers_to_save = None + text_encoder_two_lora_layers_to_save = None - for model in models: - if isinstance(model, type(accelerator.unwrap_model(unet))): - unet_lora_layers_to_save = unet_attn_processors_state_dict(model) - elif isinstance(model, type(accelerator.unwrap_model(text_encoder_one))): - text_encoder_one_lora_layers_to_save = text_encoder_lora_state_dict(model) - elif isinstance(model, type(accelerator.unwrap_model(text_encoder_two))): - text_encoder_two_lora_layers_to_save = text_encoder_lora_state_dict(model) - else: - raise ValueError(f"unexpected save model: {model.__class__}") + for model in models: + if isinstance(model, type(accelerator.unwrap_model(unet))): + unet_lora_layers_to_save = unet_attn_processors_state_dict(model) + elif isinstance(model, type(accelerator.unwrap_model(text_encoder_one))): + text_encoder_one_lora_layers_to_save = text_encoder_lora_state_dict(model) + elif isinstance(model, type(accelerator.unwrap_model(text_encoder_two))): + text_encoder_two_lora_layers_to_save = text_encoder_lora_state_dict(model) + else: + raise ValueError(f"unexpected save model: {model.__class__}") - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() - StableDiffusionXLPipeline.save_lora_weights( - output_dir, - unet_lora_layers=unet_lora_layers_to_save, - text_encoder_lora_layers=text_encoder_one_lora_layers_to_save, - text_encoder_2_lora_layers=text_encoder_two_lora_layers_to_save, - ) + StableDiffusionXLPipeline.save_lora_weights( + output_dir, + unet_lora_layers=unet_lora_layers_to_save, + text_encoder_lora_layers=text_encoder_one_lora_layers_to_save, + text_encoder_2_lora_layers=text_encoder_two_lora_layers_to_save, + ) def load_model_hook(models, input_dir): unet_ = None diff --git a/examples/text_to_image/train_text_to_image_sdxl.py b/examples/text_to_image/train_text_to_image_sdxl.py index d698243f74..c6a1b907a3 100644 --- a/examples/text_to_image/train_text_to_image_sdxl.py +++ b/examples/text_to_image/train_text_to_image_sdxl.py @@ -651,14 +651,15 @@ def main(args): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - if args.use_ema: - ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) + if accelerator.is_main_process: + if args.use_ema: + ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema")) - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "unet")) + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): if args.use_ema: diff --git a/examples/unconditional_image_generation/train_unconditional.py b/examples/unconditional_image_generation/train_unconditional.py index bfa4826902..40e71e7186 100644 --- a/examples/unconditional_image_generation/train_unconditional.py +++ b/examples/unconditional_image_generation/train_unconditional.py @@ -309,14 +309,15 @@ def main(args): if version.parse(accelerate.__version__) >= version.parse("0.16.0"): # create custom saving & loading hooks so that `accelerator.save_state(...)` serializes in a nice format def save_model_hook(models, weights, output_dir): - if args.use_ema: - ema_model.save_pretrained(os.path.join(output_dir, "unet_ema")) + if accelerator.is_main_process: + if args.use_ema: + ema_model.save_pretrained(os.path.join(output_dir, "unet_ema")) - for i, model in enumerate(models): - model.save_pretrained(os.path.join(output_dir, "unet")) + for i, model in enumerate(models): + model.save_pretrained(os.path.join(output_dir, "unet")) - # make sure to pop weight so that corresponding model is not saved again - weights.pop() + # make sure to pop weight so that corresponding model is not saved again + weights.pop() def load_model_hook(models, input_dir): if args.use_ema: