Better checkpoint file names.

JoePenna · May 9, 2023 · 8f8a49c · 8f8a49c
1 parent cdf58e4
commit 8f8a49c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/dreambooth_helpers/joepenna_dreambooth_config.py b/dreambooth_helpers/joepenna_dreambooth_config.py
@@ -198,10 +198,10 @@ def saturate_from_file(
     def toJSON(self):
         return json.dumps(self, default=lambda o: o.__dict__, sort_keys=True, indent=4)
 
-    def create_checkpoint_file_name(self, steps):
+    def create_checkpoint_file_name(self, steps: str):
         date_string = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H-%M-%S")
         return f"{date_string}_{self.project_name}_" \
-               f"{steps}_steps_" \
+               f"{int(steps):05d}_steps_" \
                f"{self.training_images_count}_training_images_" \
                f"{self.token}_token_" \
                f"{self.class_word}_class_word.ckpt".replace(" ", "_")