guard save model hooks to only execute on main process (#4929)

examples/controlnet/train_controlnet.py

@@ -785,6 +785,7 @@ 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 accelerator.is_main_process:
                 i = len(weights) - 1
 
                 while len(weights) > 0:

examples/controlnet/train_controlnet_sdxl.py

@@ -840,6 +840,7 @@ 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 accelerator.is_main_process:
                 i = len(weights) - 1
 
                 while len(weights) > 0:

examples/dreambooth/train_dreambooth.py

@@ -920,6 +920,7 @@ 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):
+        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))

examples/dreambooth/train_dreambooth_lora.py

@@ -894,6 +894,7 @@ 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):
+        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

examples/dreambooth/train_dreambooth_lora_sdxl.py

@@ -798,6 +798,7 @@ 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):
+        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

examples/instruct_pix2pix/train_instruct_pix2pix.py

@@ -485,6 +485,7 @@ 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 accelerator.is_main_process:
                 if args.use_ema:
                     ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
 

examples/instruct_pix2pix/train_instruct_pix2pix_sdxl.py

@@ -528,6 +528,7 @@ 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 accelerator.is_main_process:
                 if args.use_ema:
                     ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
 

examples/research_projects/controlnet/train_controlnet_webdataset.py

@@ -1010,6 +1010,7 @@ 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 accelerator.is_main_process:
                 i = len(weights) - 1
 
                 while len(weights) > 0:

examples/research_projects/onnxruntime/text_to_image/train_text_to_image.py

@@ -552,6 +552,7 @@ 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 accelerator.is_main_process:
                 if args.use_ema:
                     ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
 

examples/research_projects/onnxruntime/unconditional_image_generation/train_unconditional.py

@@ -313,6 +313,7 @@ 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 accelerator.is_main_process:
                 if args.use_ema:
                     ema_model.save_pretrained(os.path.join(output_dir, "unet_ema"))
 

examples/text_to_image/train_text_to_image.py

@@ -629,6 +629,7 @@ 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 accelerator.is_main_process:
                 if args.use_ema:
                     ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
 

examples/text_to_image/train_text_to_image_lora_sdxl.py

@@ -669,6 +669,7 @@ 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):
+        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

examples/text_to_image/train_text_to_image_sdxl.py

@@ -651,6 +651,7 @@ 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 accelerator.is_main_process:
                 if args.use_ema:
                     ema_unet.save_pretrained(os.path.join(output_dir, "unet_ema"))
 

examples/unconditional_image_generation/train_unconditional.py

@@ -309,6 +309,7 @@ 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 accelerator.is_main_process:
                 if args.use_ema:
                     ema_model.save_pretrained(os.path.join(output_dir, "unet_ema"))