Fix issues for saving checkpointing steps

leisuzz · Nov 13, 2024 · 5d07df8 · 5d07df8
1 parent 486a115
commit 5d07df8
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 4 deletions.
diff --git a/examples/dreambooth/train_dreambooth_lora.py b/examples/dreambooth/train_dreambooth_lora.py
@@ -955,7 +955,8 @@ def save_model_hook(models, weights, output_dir):
                     raise ValueError(f"unexpected save model: {model.__class__}")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
             StableDiffusionLoraLoaderMixin.save_lora_weights(
                 output_dir,

diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py
@@ -1302,7 +1302,8 @@ def save_model_hook(models, weights, output_dir):
                     raise ValueError(f"unexpected save model: {model.__class__}")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
             StableDiffusion3Pipeline.save_lora_weights(
                 output_dir,

diff --git a/examples/dreambooth/train_dreambooth_lora_sdxl.py b/examples/dreambooth/train_dreambooth_lora_sdxl.py
@@ -1235,7 +1235,8 @@ def save_model_hook(models, weights, output_dir):
                     raise ValueError(f"unexpected save model: {model.__class__}")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
             StableDiffusionXLPipeline.save_lora_weights(
                 output_dir,

diff --git a/examples/dreambooth/train_dreambooth_sd3.py b/examples/dreambooth/train_dreambooth_sd3.py
@@ -1192,7 +1192,8 @@ def save_model_hook(models, weights, output_dir):
                     raise ValueError(f"Wrong model supplied: {type(model)=}.")
 
                 # make sure to pop weight so that corresponding model is not saved again
-                weights.pop()
+                if weights:
+                    weights.pop()
 
     def load_model_hook(models, input_dir):
         for _ in range(len(models)):