Local_fast_DreamBooth-Win.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qEsNHTtVlbkV"
      },
      "source": [
        "# **Local-DreamBooth colab From https://github.com/TheLastBen/fast-stable-diffusion, if you face any issues, feel free to discuss them.** \n",
        "Keep your notebook updated for best experience.\n",
        "\n",
        "This is a windows version, no support for bitsandbytes or DeepSpeed, it will require more than 18GB of VRAM. A Linux version is in progress ...\n",
        "To connect a local runtime, run these commands in Anaconda :"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "llIqV5ZkH4qa"
      },
      "outputs": [],
      "source": [
        "#@markdown #Dependencies, run only once, make sure you have all A1111 dependencies installed before running this cell, including xformers.\n",
        "\n",
        "#1- pip install jupyter_http_over_ws\n",
        "#2- jupyter serverextension enable --py jupyter_http_over_ws\n",
        "#3- jupyter notebook --NotebookApp.allow_origin='https://colab.research.google.com'   --port=8888  --NotebookApp.port_retries=0 --no-browser\n",
        "#4- Use the link given : \"https://localhost:8888/?token=xxxxxx\" as the local server in google colab\n",
        "\n",
        "from subprocess import getoutput\n",
        "import os\n",
        "\n",
        "!git clone --depth 1 https://github.com/TheLastBen/diffusers\n",
        "!pip uninstall -q -y diffusers\n",
        "!pip uninstall -q -y transformers\n",
        "!pip install -q git+https://github.com/TheLastBen/diffusers\n",
        "!pip uninstall -y -q accelerate\n",
        "!pip install -q accelerate==0.12.0\n",
        "!pip install -q OmegaConf\n",
        "!pip install -q -U wandb\n",
        "!pip uninstall -y -q wget\n",
        "!pip install -q wget==3.2\n",
        "!pip install py7zr\n",
        "import wget\n",
        "wget.download('https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dreambooth/Deps')\n",
        "!ren Deps Deps.7z\n",
        "!py7zr x Deps.7z\n",
        "s=getoutput('pip show wget')\n",
        "s=s[213:259]\n",
        "!xcopy /E /Q /H /Y \"$MAIN_DIR/usr/local/lib/python3.7/dist-packages\" \"$s\"\n",
        "os.remove('Deps.7z')\n",
        "!rd /S /Q \"usr\"\n",
        "print('\u001b[1;31mDone!')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0tN76Cj5P3RL"
      },
      "source": [
        "# Dreambooth"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "A1B299g-_VJo"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from IPython.display import clear_output\n",
        "from IPython.utils import capture\n",
        "import wget\n",
        "import time\n",
        "\n",
        "#@markdown #Create/Load Session\n",
        "\n",
        "#@markdown ---\n",
        "\n",
        "try:\n",
        "   if Session_Name:\n",
        "     pass\n",
        "except:   \n",
        "   MAIN_DIR=os.getcwd()\n",
        "\n",
        "Session_Name = \"\" #@param{type: 'string'}\n",
        "while Session_Name==\"\":\n",
        "  print('\u001b[1;31mInput the Session Name:')\n",
        "  Session_Name=input('')\n",
        "INSTANCE_NAME=Session_Name\n",
        "\n",
        "\n",
        "#@markdown - To resume a previous session, just enter its name, it if it exists, it will load it, otherwise a new session will be created.\n",
        "\n",
        "WORKSPACE=MAIN_DIR+'/Fast-Dreambooth'\n",
        "MODEL_NAME=MAIN_DIR+'/stable-diffusion-v1-5'\n",
        "OUTPUT_DIR=MAIN_DIR+'/models/'+Session_Name\n",
        "SESSION_DIR=WORKSPACE+'/Sessions/'+Session_Name\n",
        "INSTANCE_DIR=WORKSPACE+'/Sessions/'+Session_Name+'/instance_images'\n",
        "MDLPTH=str(SESSION_DIR+'/'+Session_Name+'.ckpt')\n",
        "\n",
        "if os.path.exists(str(SESSION_DIR)) and not os.path.exists(str(SESSION_DIR+\"/\"+Session_Name+'.ckpt')):\n",
        "  print('\u001b[1;32mSession found, Loading session....')\n",
        "  while not os.path.exists(MODEL_NAME+'/unet/diffusion_pytorch_model.bin'):\n",
        "      print('\u001b[1;31mNo model found, make sure you put the diffusers model in the right folder')\n",
        "      time.sleep(5)\n",
        "  print('\u001b[1;32mSession loaded with no previous CKPT.')\n",
        "\n",
        "elif os.path.exists(str(SESSION_DIR+\"/\"+Session_Name+'.ckpt')):\n",
        "  print('\u001b[1;32mSession found, loading the model, this might take a few minutes...')\n",
        "  if not os.path.exists(str(OUTPUT_DIR)):\n",
        "    %mkdir \"$OUTPUT_DIR\"\n",
        "  with capture.capture_output() as cap:\n",
        "    %cd $MAIN_DIR\n",
        "    wget.download('https://raw.githubusercontent.com/TheLastBen/fast-stable-diffusion/main/Dreambooth/convertodiffloc.py', MAIN_DIR)\n",
        "  !python convertodiffloc.py --checkpoint_path \"$MDLPTH\" --dump_path \"$OUTPUT_DIR\"\n",
        "  if os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "    resume=True\n",
        "    os.remove('convertodiffloc.py')\n",
        "    os.remove('v1-inference.yaml')\n",
        "    clear_output()\n",
        "    print('\u001b[1;32mSession loaded with the trained model')\n",
        "  else:\n",
        "    os.remove('convertodiffloc.py')\n",
        "    os.remove('v1-inference.yaml')\n",
        "    while not os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "      print('\u001b[1;31mConversion error, it appears the the CKPT from this session is incompatible or corrupt, remove it to continue')\n",
        "      time.sleep(5)\n",
        "\n",
        "\n",
        "elif not os.path.exists(str(SESSION_DIR)):\n",
        "    %mkdir \"$INSTANCE_DIR\"\n",
        "    print('\u001b[1;32mCreating session...')\n",
        "    while not os.path.exists(MODEL_NAME+'/unet/diffusion_pytorch_model.bin'):\n",
        "        print('\u001b[1;31mNo model found, make sure you put the diffusers model in the right folder')\n",
        "        time.sleep(5)\n",
        "    print('\u001b[1;32mSession Created.')\n",
        "    \n",
        "#@markdown \n",
        "\n",
        "#@markdown # The most importent step is to rename the instance picture to the same instance unique identifier for each subject, example :\n",
        "#@markdown - If you have 30 pictures of yourself, simply select them all and rename only one to the chosen identifier for example : phtmejhn, the files would be : phtmejhn (1).jpg, phtmejhn (2).png ....etc then upload them, do the same for other people or objects with a different identifier, and that's it.\n",
        "#@markdown - Check out this example : https://i.imgur.com/d2lD3rz.jpeg\n",
        "      \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "LC4ukG60fgMy"
      },
      "outputs": [],
      "source": [
        "#@markdown #Instance Images\n",
        "#@markdown ----\n",
        "\n",
        "#@markdown\n",
        "#@markdown - Run the cell to Upload the instance pictures.\n",
        "\n",
        "\n",
        "\n",
        "Remove_existing_instance_images= True #@param{type: 'boolean'}\n",
        "#@markdown - This will remove the previous instance images, uncheck to add the new isntance pictures to the existing ones (if any).\n",
        "\n",
        "\n",
        "if Remove_existing_instance_images:\n",
        "  if os.path.exists(str(INSTANCE_DIR)):\n",
        "    !rd /S /Q \"$INSTANCE_DIR\"\n",
        "\n",
        "if not os.path.exists(str(INSTANCE_DIR)):\n",
        "  %mkdir \"$INSTANCE_DIR\"\n",
        "\n",
        "IMAGES__FOLDER=\"\" #@param{type: 'string'}\n",
        "IMAGES_FOLDER=MAIN_DIR+'/'+IMAGES__FOLDER\n",
        "\n",
        "#@markdown - Enter the path of the folder containing your instance images\n",
        "\n",
        "while IMAGES__FOLDER !=\"\" and not os.path.exists(str(IMAGES_FOLDER)):\n",
        "  print('\u001b[1;31mThe image folder specified does not exist, use the colab file explorer to copy the path')\n",
        "  time.sleep(5)\n",
        "\n",
        "\n",
        "if IMAGES__FOLDER!=\"\":\n",
        "  with capture.capture_output() as cap:\n",
        "    !xcopy /Y \"$IMAGES_FOLDER\" \"$INSTANCE_DIR\" \n",
        "    if os.path.exists(INSTANCE_DIR+\"/.ipynb_checkpoints\"):\n",
        "      !rd /S /Q  INSTANCE_DIR+\"/.ipynb_checkpoints\"      \n",
        "  print('\u001b[1;32mDone, proceed to the training cell')\n",
        "elif IMAGES__FOLDER==\"\":\n",
        "  print(('\u001b[1;31mProceeding without uploading instance images.'))\n",
        "\n",
        "%cd \"$INSTANCE_DIR\"\n",
        "[os.rename(f, f.replace(' ', '_')) for f in os.listdir('.') if not f.startswith('.')]\n",
        "%cd $MAIN_DIR"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnmQYfZilzY6"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "1-9QbkfAVYYU"
      },
      "outputs": [],
      "source": [
        "#@markdown ---\n",
        "#@markdown #Start DreamBooth\n",
        "#@markdown ---\n",
        "\n",
        "Resume_Training = False #@param {type:\"boolean\"}\n",
        "\n",
        "while not Resume_Training and not os.path.exists(MODEL_NAME+'/unet/diffusion_pytorch_model.bin'):\n",
        "   print('\u001b[1;31mNo model found, make sure you put the diffusers model in the right folder')\n",
        "   time.sleep(5)\n",
        "\n",
        "#@markdown  - If you're not satisfied with the result, check this box, run again the cell and it will continue training the current model.\n",
        "\n",
        "MODELT_NAME=MODEL_NAME\n",
        "\n",
        "Training_Steps=3000 #@param{type: 'number'}\n",
        "#@markdown - Total Steps = Number of Instance images * 100, if you use 30 images, use 3000 steps, if you're not satisfied with the result, resume training for another 500 steps, and so on ...\n",
        "\n",
        "Seed=96576 #@param{type: 'number'}\n",
        "\n",
        "\n",
        "try:\n",
        "   resume\n",
        "   if resume and not Resume_Training:\n",
        "     print('\u001b[1;31mOverwrite your previously trained model ?, answering \"yes\" will train a new model, answering \"no\" will resume the training of the previous model?  yes or no ?\u001b[0m')\n",
        "     while True:\n",
        "        ansres=input('')\n",
        "        if ansres=='no':\n",
        "          Resume_Training = True\n",
        "          del ansres\n",
        "          break\n",
        "        elif ansres=='yes':\n",
        "          Resume_Training = False\n",
        "          resume= False\n",
        "          break\n",
        "except:\n",
        "  pass\n",
        "\n",
        "if Resume_Training and os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "  MODELT_NAME=OUTPUT_DIR\n",
        "  print('\u001b[1;32mResuming Training...\u001b[0m')\n",
        "elif Resume_Training and not os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "  MODELT_NAME=MODEL_NAME  \n",
        "  print('\u001b[1;31mPrevious model not found, training a new model...\u001b[0m') \n",
        "\n",
        "\n",
        "\n",
        "!accelerate launch $MAIN_DIR/diffusers/examples/dreambooth/train_dreambooth.py \\\n",
        "  --image_captions_filename \\\n",
        "  --train_text_encoder \\\n",
        "  --pretrained_model_name_or_path=\"$MODELT_NAME\" \\\n",
        "  --instance_data_dir=\"$INSTANCE_DIR\" \\\n",
        "  --output_dir=\"$OUTPUT_DIR\" \\\n",
        "  --instance_prompt=\"\" \\\n",
        "  --seed=$Seed \\\n",
        "  --resolution=512 \\\n",
        "  --mixed_precision=\"fp16\" \\\n",
        "  --sample_batch_size=1 \\\n",
        "  --train_batch_size=1 \\\n",
        "  --gradient_accumulation_steps=1 --gradient_checkpointing \\\n",
        "  --learning_rate=2e-6 \\\n",
        "  --lr_scheduler=\"polynomial\" \\\n",
        "  --center_crop \\\n",
        "  --lr_warmup_steps=0 \\\n",
        "  --max_train_steps=$Training_Steps\n",
        "\n",
        "if os.path.exists(OUTPUT_DIR+'/unet/diffusion_pytorch_model.bin'):\n",
        "  print(\"Almost done ...\")\n",
        "  %cd \"$\n",
        "  !wget -O convertosd.py https://github.com/TheLastBen/fast-stable-diffusion/raw/main/Dreambooth/convertosd.py\n",
        "  clear_output()\n",
        "  !sed -i '201s@.*@    model_path = \"{OUTPUT_DIR}\"@' $MAIN_DIR/convertosd.py\n",
        "  !sed -i '202s@.*@    checkpoint_path= \"{SESSION_DIR}/{Session_Name}.ckpt\"@' $MAIN_DIR/convertosd.py\n",
        "  !python $MAIN_DIR/convertosd.py\n",
        "  clear_output()\n",
        "  if os.path.exists(SESSION_DIR+\"/\"+INSTANCE_NAME+'.ckpt'):\n",
        "    print(\"\u001b[1;32mDONE, the CKPT model is in the sessions folder\")\n",
        "  else:\n",
        "    print(\"\u001b[1;31mSomething went wrong\")\n",
        "    \n",
        "else:\n",
        "  print(\"\u001b[1;31mSomething went wrong\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "provenance": []
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}