Updating the jupyter notebooks to remove the manual pruning step.

Author: djbielejeski

dreambooth_colab_joepenna.ipynb

@@ -214,17 +214,15 @@
   {
    "cell_type": "code",
    "source": [
-    "#@title # Prune the Model\n",
+    "#@title # Copy and name the checkpoint file\n",
     "\n",
     "directory_paths = !ls -d logs/*\n",
     "last_checkpoint_file = directory_paths[-1] + \"/checkpoints/last.ckpt\"\n",
-    "!python \"prune_ckpt.py\" --ckpt {last_checkpoint_file}\n",
-    "last_checkpoint_file_pruned = directory_paths[-1] + \"/checkpoints/last-pruned.ckpt\"\n",
     "training_images = !find training_images/*\n",
     "date_string = !date +\"%Y-%m-%dT%H-%M-%S\"\n",
     "file_name = date_string[-1] + \"_\" + project_name + \"_\" + str(len(training_images)) + \"_training_images_\" +  str(max_training_steps) + \"_max_training_steps_\" + token + \"_token_\" + class_word + \"_class_word.ckpt\"\n",
     "!mkdir -p trained_models\n",
-    "!mv {last_checkpoint_file_pruned} trained_models/{file_name}\n",
+    "!mv {last_checkpoint_file} trained_models/{file_name}\n",
     "\n",
     "print(\"Download your trained model file from trained_models/\" + file_name + \" and use in your favorite Stable Diffusion repo!\")"
    ],

dreambooth_runpod_joepenna.ipynb

@@ -364,8 +364,7 @@
    "id": "dc49d0bd",
    "metadata": {},
    "source": [
-    "## Pruning (12GB to 2GB)\n",
-    "We are working on having this happen automatically (TODO: PR's welcome)"
+    "## Copy and name the checkpoint file"
    ]
   },
   {
@@ -376,17 +375,15 @@
    },
    "outputs": [],
    "source": [
-    "# Prune the Model\n",
+    "# Copy the checkpoint into our `trained_models` folder\n",
     "\n",
     "directory_paths = !ls -d logs/*\n",
     "last_checkpoint_file = directory_paths[-1] + \"/checkpoints/last.ckpt\"\n",
-    "!python \"prune_ckpt.py\" --ckpt {last_checkpoint_file}\n",
-    "last_checkpoint_file_pruned = directory_paths[-1] + \"/checkpoints/last-pruned.ckpt\"\n",
     "training_images = !find training_images/*\n",
     "date_string = !date +\"%Y-%m-%dT%H-%M-%S\"\n",
     "file_name = date_string[-1] + \"_\" + project_name + \"_\" + str(len(training_images)) + \"_training_images_\" +  str(max_training_steps) + \"_max_training_steps_\" + token + \"_token_\" + class_word + \"_class_word.ckpt\"\n",
     "!mkdir -p trained_models\n",
-    "!mv {last_checkpoint_file_pruned} trained_models/{file_name}\n",
+    "!mv {last_checkpoint_file} trained_models/{file_name}\n",
     "\n",
     "print(\"Download your trained model file from trained_models/\" + file_name + \" and use in your favorite Stable Diffusion repo!\")"
    ]