Moar tweaks

Author: lewtun

hackathon/dreambooth.ipynb

@@ -161,7 +161,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "%pip install -qqU diffusers transformers bitsandbytes accelerate ftfy datasets ipywidgets"
+    "%pip install -qqU diffusers transformers bitsandbytes accelerate ftfy datasets"
    ]
   },
   {
@@ -887,7 +887,6 @@
    "source": [
     "import math\n",
     "\n",
-    "import bitsandbytes as bnb\n",
     "import torch.nn.functional as F\n",
     "from accelerate import Accelerator\n",
     "from accelerate.utils import set_seed\n",
@@ -910,6 +909,7 @@
     "\n",
     "    # Use 8-bit Adam for lower memory usage or to fine-tune the model in 16GB GPUs\n",
     "    if args.use_8bit_adam:\n",
+    "        import bitsandbytes as bnb\n",
     "        optimizer_class = bnb.optim.AdamW8bit\n",
     "    else:\n",
     "        optimizer_class = torch.optim.AdamW\n",
@@ -1019,6 +1019,7 @@
     "\n",
     "    # Create the pipeline using using the trained modules and save it.\n",
     "    if accelerator.is_main_process:\n",
+    "        print(f\"Loading pipeline and saving to {args.output_dir}...\")\n",
     "        scheduler = PNDMScheduler(\n",
     "            beta_start=0.00085,\n",
     "            beta_end=0.012,\n",
@@ -1185,8 +1186,10 @@
     }
    ],
    "source": [
-    "# Pick a funny prompt here and it will be saved as the default for widget on the Hub!\n",
+    "# Pick a funny prompt here and it will be used as the widget's default \n",
+    "# when we push to the Hub in the next section\n",
     "prompt = f\"a photo of {name_of_your_concept} {type_of_thing} in the Acropolis\"\n",
+    "\n",
     "# Tune the guidance to control how closely the generations follow the prompt.\n",
     "# Values between 7-11 usually work best\n",
     "guidance_scale = 7\n",