SDfu_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "BCyeZMHNT2E-",
        "yWpWFeyO8APF"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.9"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xzA1-mt88AO_"
      },
      "source": [
        "# Stable Diffusion on diffusers\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7IFfx8GQIAQm"
      },
      "source": [
        "\n",
        "**Run this cell after each session restart**"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## General Setup"
      ],
      "metadata": {
        "id": "JVZD4V5yb1P3"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tkbcraCUaPEy",
        "cellView": "form"
      },
      "source": [
        "!pip install gputil\n",
        "\n",
        "import warnings\n",
        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
        "import logging\n",
        "logging.getLogger('diffusers.models.modeling_utils').setLevel(logging.CRITICAL)\n",
        "\n",
        "from IPython.display import HTML, Image, display, clear_output\n",
        "from moviepy.editor import ImageSequenceClip, ipython_display\n",
        "import ipywidgets as widgets\n",
        "import os\n",
        "import numpy as np\n",
        "\n",
        "# !apt-get -qq install ffmpeg\n",
        "!pip install ninja\n",
        "from google.colab import drive\n",
        "drive.mount('/G', force_remount=True)\n",
        "gdir = '/G/MyDrive/'\n",
        "%cd $gdir\n",
        "\n",
        "#@markdown Setting up Stable Diffusion code in the directory below on your Google drive (creating it, if it doesn't exist):\n",
        "work_dir = 'sdfu' #@param {type:\"string\"}\n",
        "#@markdown NB: Avoid connecting Google drive manually via the icon in Files section on the left. Doing so may break further operations.\n",
        "\n",
        "work_dir = gdir + work_dir + '/'\n",
        "if not os.path.isdir(work_dir):\n",
        "  !git clone https://github.com/eps696/SDfu $work_dir\n",
        "%cd $work_dir\n",
        "# the order is important! basicsr (for controlnet) needs pt latest\n",
        "!pip install --no-deps xformers==0.0.29 # for torch 2.5.1+cu124\n",
        "!pip install --no-deps optimum-quanto # 0.2.4 for torch 2.4\n",
        "\n",
        "!pip install --no-deps basicsr\n",
        "# workaround for making outdated basicsr work with modern torchvision\n",
        "import sys, types\n",
        "from torchvision.transforms.functional import rgb_to_grayscale\n",
        "# Create a module for `torchvision.transforms.functional_tensor`\n",
        "functional_tensor = types.ModuleType(\"torchvision.transforms.functional_tensor\")\n",
        "functional_tensor.rgb_to_grayscale = rgb_to_grayscale\n",
        "# Add this module to sys.modules so other imports can access it\n",
        "sys.modules[\"torchvision.transforms.functional_tensor\"] = functional_tensor\n",
        "\n",
        "!pip install --no-deps git+https://github.com/openai/CLIP\n",
        "!pip install --no-deps torchmetrics omegaconf trampoline torchdiffeq torchsde jsonmerge clean-fid resize_right av\n",
        "# !pip install -r requirements.txt\n",
        "\n",
        "from src.core.text import txt_clean\n",
        "from src.core.utils import basename\n",
        "from download import get_model\n",
        "maindir = '/content/models'\n",
        "command = ''\n",
        "args = ''\n",
        "\n",
        "clear_output()\n",
        "\n",
        "# Hardware check\n",
        "!ln -sf /opt/bin/nvidia-smi /usr/bin/nvidia-smi\n",
        "import GPUtil as GPU\n",
        "gpu = GPU.getGPUs()[0]\n",
        "!nvidia-smi -L\n",
        "print(\"GPU RAM {0:.0f}MB | Free {1:.0f}MB)\".format(gpu.memoryTotal, gpu.memoryFree))\n",
        "print('\\nDone!')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Njelbgu8APJ"
      },
      "source": [
        "## Images\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load base model\n",
        "\n",
        "model = \"1.5-dreamlike\" #@param ['1.5','1.5-dreamlike','LCM','2.1','2-inpainting','Kandinsky 2.2','SDXL','SDXL Lightning 4 steps','SDXL Lightning 8 steps']\n",
        "\n",
        "#@markdown NB: Options and methods below work mostly with SD 1.x models.\n",
        "#@markdown Some can work with SDXL. Models 2.x and Kandinsky don't support them, and are given only for reference.\n",
        "\n",
        "if '1.5' in model and not os.path.exists('/content/models/v1'):\n",
        "  get_model(\"https://www.dropbox.com/s/9wuhum8w0iqs4o7/sdfu-v15-full-fp16.zip?dl=1\", '/content/models/v1')\n",
        "  get_model(\"https://www.dropbox.com/s/z9uycihl6tybx9y/sdfu-v1-vaes-fp16.zip?dl=1\", '/content/models/v1')\n",
        "if model == '1.5-dreamlike' and not os.path.exists('/content/models/v1/unet15drm'):\n",
        "  get_model(\"https://www.dropbox.com/s/7hwh27xfzcy7a2g/sdfu-v15drm-unetx-fp16.zip?dl=1\", '/content/models/v1')\n",
        "if model == 'LCM' and not os.path.exists('/content/models/lcm'):\n",
        "  get_model(\"https://www.dropbox.com/s/3sql6gsjmtvw2zo/sdfu-lcm-fp16.zip?dl=1\", '/content/models')\n",
        "if model in ['2.1', '2-inpainting'] and not os.path.exists('/content/models/v2'):\n",
        "  get_model(\"https://www.dropbox.com/s/38876tjuklvwq82/sdfu-v21-full-fp16.zip?dl=1\", '/content/models/v2')\n",
        "if model == '2-inpainting' and not os.path.exists('/content/models/v2/unet2i'):\n",
        "  get_model(\"https://www.dropbox.com/s/r5fa1mdxpw9e8k2/sdfu-v2i-unet-fp16.zip?dl=1\", '/content/models/v2')\n",
        "\n",
        "# clipseg for text masking\n",
        "if not os.path.exists('/content/models/xtra/clipseg'):\n",
        "  get_model(\"https://www.dropbox.com/scl/fi/blwzoyu9abp4q6wfq4hoj/rd64-uni.pth?rlkey=mhfp4jlles5oio1sn1eneyx4l&dl=1\", '/content/models/xtra/clipseg', unzip=False)\n",
        "\n",
        "modict = {'1.5':'15', '1.5-dreamlike':'15drm', 'LCM':'lcm', '2-inpainting':'2i', '2.1':'21', 'Kandinsky 2.2':'kand', 'SDXL':'xl', 'SDXL Lightning 4 steps':'xl4', 'SDXL Lightning 8 steps':'xl8'}\n",
        "modname = modict[model]\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "A9a_9LXrV2dR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Inputs"
      ],
      "metadata": {
        "id": "Ud1WQ0EqWLcw"
      }
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GwoOBOcR8APK",
        "cellView": "form"
      },
      "source": [
        "#@markdown Run this for every new generation!\n",
        "\n",
        "#@markdown All paths below are relative to the work directory on G drive (which is set during General setup above).\n",
        "\n",
        "#@markdown Specify a text string or path to a text file to use **txt2img**:\n",
        "prompt = '' #@param {type:\"string\"}\n",
        "#@markdown Specify path to a reference image or directory to use it as a **visual prompt** with IP adapter:\n",
        "imgref = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Specify path to an image or directory to use **img2img**:\n",
        "images = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Specify mask as a text description or a path to an image or directory, to use **inpainting**:\n",
        "mask = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Load **finetuned** files if needed (ensure that you selected correct base model!):\n",
        "method = \"custom\" #@param ['text','lora','custom']\n",
        "load_file = '' #@param {type:\"string\"}\n",
        "\n",
        "maindir = '/content/models'\n",
        "%cd $work_dir\n",
        "command = ' -md %s ' % maindir\n",
        "workname = ''\n",
        "\n",
        "if len(prompt) > 0:\n",
        "  if os.path.exists(os.path.join(gdir, prompt)):\n",
        "    prompt = os.path.join(gdir, prompt)\n",
        "  if os.path.exists(prompt):\n",
        "    print('found', prompt)\n",
        "    command += ' -t %s ' % prompt\n",
        "  else:\n",
        "    command += ' -t \"%s\" ' % prompt\n",
        "  workname = txt_clean(basename(prompt))[:44]\n",
        "\n",
        "if len(images) > 0:\n",
        "  if not os.path.exists(images):\n",
        "    images = os.path.join(gdir, images)\n",
        "  if os.path.exists(images):\n",
        "    print('found', images)\n",
        "    command += ' -im %s ' % images\n",
        "    if len(workname) == 0:\n",
        "      workname = txt_clean(basename(images))\n",
        "else:\n",
        "  command += ' -f 1 '\n",
        "  print('Strength is set to 1 for txt2img method')\n",
        "\n",
        "if len(imgref) > 0:\n",
        "  if not os.path.exists(imgref):\n",
        "    imgref = os.path.join(gdir, imgref)\n",
        "  if os.path.exists(imgref):\n",
        "    if not os.path.exists('/content/models/image'):\n",
        "      print('downloading IP adapter')\n",
        "      get_model(\"https://www.dropbox.com/scl/fi/zcvg092n1n4aqpmoufd70/sdfu-imageref-fp16.zip?rlkey=ompjbfu2um90hyid2rlglv4d8&dl=1\", maindir)\n",
        "    print('using IP adapter with', imgref)\n",
        "    command += ' -imr %s ' % imgref\n",
        "    workname += '-' + txt_clean(basename(imgref))\n",
        "\n",
        "if len(mask) > 0:\n",
        "  if not os.path.exists(mask):\n",
        "    mask = '\"%s\"' % mask\n",
        "  command += ' -M %s ' % mask\n",
        "\n",
        "if len(load_file) > 0:\n",
        "  if not os.path.exists(load_file):\n",
        "    load_file = os.path.join(gdir, load_file)\n",
        "  if os.path.exists(load_file):\n",
        "    print('found', load_file)\n",
        "    cmd = 'rt' if method=='text' else 'rl' if method=='lora' else 'rd'\n",
        "    command += ' -%s %s ' % (cmd, load_file)\n",
        "\n",
        "command += ' --model %s ' % modname\n",
        "\n",
        "if 'Lightning' in model:\n",
        "  command += ' --lightning '\n",
        "\n",
        "# !echo $command $args\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Add reference images to control output by **Controlnet** if needed:\n",
        "cnet_images = '' #@param {type:\"string\"}\n",
        "cnet_mode = \"depth\" #@param ['depth any 2','depth','canny','pose']\n",
        "cnet_strength = 0.7 #@param {type:\"number\"}\n",
        "if cnet_mode == 'depth any 2': cnet_mode = 'deptha'\n",
        "\n",
        "if len(cnet_images) > 0 and cnet_strength > 0:\n",
        "    if not os.path.exists(cnet_images):\n",
        "      cnet_images = os.path.join(gdir, cnet_images)\n",
        "    if os.path.exists(cnet_images):\n",
        "      if not os.path.exists('/content/models/control'):\n",
        "        print('downloading Controlnet model')\n",
        "        get_model(\"https://www.dropbox.com/s/qhe1zpbubjr3t75/sdfu-controlnet.zip?dl=1\", '/content/models')\n",
        "      print('using Controlnet with', cnet_images)\n",
        "      predir = '/content/%s/%s' % (basename(cnet_images), cnet_mode)\n",
        "      %run src/preproc.py -i $cnet_images --type $cnet_mode -o $predir -md /content/models/control\n",
        "      command += ' -cmod %s -cnimg %s -cts %f ' % (cnet_mode, predir, cnet_strength)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "nrQM2TOu6vnU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Other settings [optional]"
      ],
      "metadata": {
        "id": "BCyeZMHNT2E-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Set low `steps` (2-8) if using TCD sampler or LCM model."
      ],
      "metadata": {
        "id": "CjO9vBt2feiD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Run this cell to override default settings, if needed\n",
        "\n",
        "out_dir = '_out' #@param {type:\"string\"}\n",
        "sizeX = 768 #@param {type:\"integer\"}\n",
        "sizeY = 768 #@param {type:\"integer\"}\n",
        "\n",
        "steps = 33 #@param {type:\"integer\"}\n",
        "cfg_scale = 7.5 #@param {type:\"number\"}\n",
        "strength = 1. #@param {type:\"number\"}\n",
        "imgref_weight = 0.5 #@param {type:\"number\"}\n",
        "ip_adapter = 'orig' #@param ['orig','plus','face-full']\n",
        "\n",
        "sampler = 'ddim' #@param ['ddim','ddpm','pndm','lms','euler', 'euler_a','uni','dpm','TCD']\n",
        "VAE = 'ema' #@param ['original', 'ema', 'mse']\n",
        "batch = 1 #@param {type:\"integer\"}\n",
        "seed = 696 #@param {type:\"integer\"}\n",
        "\n",
        "eta = 0. #@ param {type:\"number\"}\n",
        "\n",
        "unprompt = 'low quality, poorly drawn, out of focus, blurry, tiled, segmented, oversaturated' #@param {type:\"string\"}\n",
        "verbose = True #@param {type:\"boolean\"}\n",
        "\n",
        "args = ''\n",
        "args += ' -o %s ' % out_dir\n",
        "args += ' -sz %d-%d ' % (sizeX, sizeY)\n",
        "\n",
        "steps = 4 if modname=='xl4' else 8 if modname=='xl8' else steps\n",
        "args += ' -s %d ' % steps\n",
        "\n",
        "args += ' -C %g ' % cfg_scale\n",
        "args += ' -f %g ' % strength\n",
        "args += ' -imw %g ' % imgref_weight\n",
        "if ip_adapter != 'orig':\n",
        "  args += ' -ip %s ' % ip_adapter\n",
        "\n",
        "if sampler=='TCD' and modname=='xl':\n",
        "  args += ' -sm TCD --load_lora h1t/TCD-SDXL-LoRA -C 1 '\n",
        "elif sampler=='TCD' and modname[0]=='1':\n",
        "  args += ' -sm tcd --load_lora h1t/TCD-SD15-LoRA -C 1 '\n",
        "elif sampler=='TCD':\n",
        "  print('TCD sampler works with SDXL or 1.x models!')\n",
        "else:\n",
        "  args += ' -sm %s ' % sampler\n",
        "\n",
        "args += ' --vae %s ' % VAE\n",
        "args += ' -b %d ' % batch\n",
        "args += ' -S %d ' % seed\n",
        "args += ' --eta %g ' % eta\n",
        "args += ' -un \"%s\" ' % unprompt\n",
        "if verbose:\n",
        "  args += ' -v '\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "B1B_woHjPvQg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Generate"
      ],
      "metadata": {
        "id": "E60K0IQLdCMV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Separate frames\n",
        "\n",
        "if \"Kandinsky\" in model:\n",
        "  %run src/kand.py $args $command\n",
        "elif \"SDXL\" in model:\n",
        "  %run src/sdxl.py $args $command\n",
        "else:\n",
        "  %run src/gen.py $args $command"
      ],
      "metadata": {
        "cellView": "form",
        "id": "8TDpplliq8Yo"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Interpolations\n",
        "\n",
        "latent_blending = 0. #@param {type:\"number\"}\n",
        "frame_step = 10 #@param {type:\"integer\"}\n",
        "num_repeat = 3 #@param {type:\"integer\"}\n",
        "\n",
        "command += ' -fs %d -n %d ' % (frame_step, num_repeat)\n",
        "\n",
        "if \"Kandinsky\" in model:\n",
        "  %run src/kand.py $args $command\n",
        "elif \"SDXL\" in model:\n",
        "  %run src/sdxl.py $args $command\n",
        "else:\n",
        "  if latent_blending > 0:\n",
        "    command += ' --latblend %f' % latent_blending\n",
        "  %run src/latwalk.py $args $command\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "8KKZ7C_Iogpe"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "set `latent_blending` if you need smooth transitions (except Kandinsky, SDXL and LCM models). value range is 0~1; 0.7 is a good start.  \n",
        "`frame_step` = length of the transition between prompts or images (in frames).  \n",
        "`num_repeat` = repeating inputs to make animation e.g. from a single prompt or image *(ignored when interpolating between images)*."
      ],
      "metadata": {
        "id": "PybkSeM0goXa"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Edit image sequence with Tokenflow\n",
        "\n",
        "!pip install av\n",
        "clear_output()\n",
        "\n",
        "images = '' #@param {type:\"string\"}\n",
        "\n",
        "if len(prompt) > 0 and len(images) > 0 and os.path.exists(images):\n",
        "  %run src/tokenflow.py -md $maindir -im $images -t \"$prompt\" --batch_size 4 --batch_pivot --cpu\n",
        "else:\n",
        "  print(\"Set text prompt and path to the images in the Inputs above!\")\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "dNGCtR9sibfG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Video"
      ],
      "metadata": {
        "id": "VqKR1u-c-_YH"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown Load AnimateDiff motion adapter (for SD 1.x base models), CogX or Zeroscope (lo-res) model\n",
        "\n",
        "animodel = \"AnimateDiff\" #@param ['AnimateDiff','Zeroscope']\n",
        "# 'CogX' > 20gb RAM\n",
        "\n",
        "if 'AnimateDiff' in animodel and not os.path.exists('/content/models/anima'):\n",
        "  assert modname[0]=='1' and os.path.exists('/content/models/v1'), \"AnimateDiff works only with 1.x base SD models\"\n",
        "  get_model(\"https://www.dropbox.com/scl/fi/4gorn5lf9owygizhgwuy6/sdfu-animatediff-fp16.zip?rlkey=mh54vl9ngre9n898pcgsr8fry&dl=1\", '/content/models')\n",
        "elif 'Zeroscope' in animodel and not os.path.exists('/content/models/v2/unetvzs'):\n",
        "  if not os.path.exists('/content/models/v2'):\n",
        "    get_model(\"https://www.dropbox.com/s/38876tjuklvwq82/sdfu-v21-full-fp16.zip?dl=1\", '/content/models/v2')\n",
        "  if not os.path.exists('/content/models/v2/unetvzs'):\n",
        "    get_model(\"https://www.dropbox.com/s/uyaidznqjaot7hw/sdfu-video-unet-fp16.zip?dl=1\", '/content/models/v2')\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "6FNkXo-wW_Od"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@markdown **Set the inputs and run the process**\n",
        "\n",
        "#@markdown Specify a text string:\n",
        "prompt = '' #@param {type:\"string\"}\n",
        "#@markdown Specify reference image(s) to use it as a **visual prompt**:\n",
        "imgref = '' #@param {type:\"string\"}\n",
        "imgref_weight = 0.3 #@param {type:\"number\"}\n",
        "use_all_at_once = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown Specify path to the input video (or a frame sequence), if needed:\n",
        "video = '' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown Specify video length in frames (leave 0 to process complete input video):\n",
        "frames = 0 #@param {type:\"integer\"}\n",
        "\n",
        "#@markdown Other parameters:\n",
        "sampler = 'euler' #@param ['ddim','ddpm','euler','uni']\n",
        "\n",
        "cmds = {'AnimateDiff':'anima', 'CogX':'cogx', 'Zeroscope':'vid'}\n",
        "command = 'src/%s.py ' % cmds[animodel]\n",
        "modpath = ' -ad anima ' if animodel == 'AnimateDiff' else ''\n",
        "\n",
        "if animodel == 'AnimateDiff':\n",
        "  dirname = ''.join(e for e in ''.join(prompt.split(' ')[:3]) if (e.isalnum()))\n",
        "  print('saving as subdir', dirname)\n",
        "  command += '-o _out/%s ' % dirname\n",
        "\n",
        "command += '-sm %s ' % sampler\n",
        "if frames > 0:\n",
        "  command += '-vf %d ' % frames\n",
        "command += ' -md /content/models ' + modpath\n",
        "# %cd $work_dir\n",
        "\n",
        "if len(prompt) > 0:\n",
        "  command += ' -t \"%s\" ' % prompt\n",
        "if len(imgref) > 0 and imgref_weight != 0:\n",
        "  if not os.path.exists(imgref):\n",
        "    imgref = os.path.join(gdir, imgref)\n",
        "  if os.path.exists(imgref):\n",
        "    if animodel == 'AnimateDiff':\n",
        "      if not os.path.exists('/content/models/image'):\n",
        "        print('downloading IP adapter')\n",
        "        get_model(\"https://www.dropbox.com/scl/fi/zcvg092n1n4aqpmoufd70/sdfu-imageref-fp16.zip?rlkey=ompjbfu2um90hyid2rlglv4d8&dl=1\", maindir)\n",
        "      print('using IP adapter with', imgref)\n",
        "      command += ' -imr %s -imw %f ' % (imgref, imgref_weight)\n",
        "      if use_all_at_once:\n",
        "        command += ' --allref '\n",
        "    elif animodel == 'CogX':\n",
        "      command += ' -im %s ' % (imgref)\n",
        "if len(video) > 0:\n",
        "  if not os.path.exists(video):\n",
        "    video = os.path.join(gdir, video)\n",
        "  if os.path.exists(video):\n",
        "    command += ' -iv %s ' % video\n",
        "\n",
        "if len(prompt) > 0 or (len(video) > 0 and os.path.exists(video)):\n",
        "  %run $command\n",
        "else:\n",
        "  print('Not enough inputs yet to run the generation')\n",
        "\n",
        "print('done')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "bEtTjbGPyxS7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yWpWFeyO8APF"
      },
      "source": [
        "## Fine-tune"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load base model\n",
        "\n",
        "model = \"1.5-dreamlike\" #@param ['1.5','1.5-dreamlike','2.1']\n",
        "\n",
        "if '1.5' in model and not os.path.exists('/content/models/v1'):\n",
        "  get_model(\"https://www.dropbox.com/s/9wuhum8w0iqs4o7/sdfu-v15-full-fp16.zip?dl=1\", '/content/models/v1')\n",
        "  get_model(\"https://www.dropbox.com/s/z9uycihl6tybx9y/sdfu-v1-vaes-fp16.zip?dl=1\", '/content/models/v1')\n",
        "if model == '1.5-dreamlike' and not os.path.exists('/content/models/v1/unet15drm'):\n",
        "  get_model(\"https://www.dropbox.com/s/7hwh27xfzcy7a2g/sdfu-v15drm-unetx-fp16.zip?dl=1\", '/content/models/v1')\n",
        "if model in ['2.1', '2-inpainting'] and not os.path.exists('/content/models/v2'):\n",
        "  get_model(\"https://www.dropbox.com/s/38876tjuklvwq82/sdfu-v21-full-fp16.zip?dl=1\", '/content/models/v2')\n",
        "if model == '2-inpainting' and not os.path.exists('/content/models/v2/unet2i'):\n",
        "  get_model(\"https://www.dropbox.com/s/r5fa1mdxpw9e8k2/sdfu-v2i-unet-fp16.zip?dl=1\", '/content/models/v2')\n",
        "\n",
        "modict = {'1.5':'15', '1.5-dreamlike':'15drm', '2.1':'21'}\n",
        "modname = modict[model]\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "YwU2XJz23uay"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tEmNzQmm0t_o",
        "cellView": "form"
      },
      "source": [
        "#@title Data setup\n",
        "#@markdown Put your target images as zip-archive onto Google drive and copy its path below (relative to G-drive root).\n",
        "target_data = 'tgt.zip' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown For Text Inversion or Custom Diffusion, provide a token to use in the prompts to summon your target imagery, e.g. as `<mycat1>`\n",
        "new_token = 'mycat1' #@param {type:\"string\"}\n",
        "\n",
        "#@markdown For Custom diffusion, prepare also a bunch of generic reference images of similar class, to start from.\n",
        "ref_class = 'cat' #@param {type:\"string\"}\n",
        "ref_data = 'ref.zip' #@param {type:\"string\"}\n",
        "\n",
        "data_dir = os.path.join('/content/data/', new_token)\n",
        "!rm -rf $data_dir\n",
        "os.makedirs(data_dir)\n",
        "%cd $data_dir\n",
        "\n",
        "tgt_path = os.path.join(gdir, target_data)\n",
        "!unzip -j -o -q $tgt_path -d tgt\n",
        "tgt_dir = os.path.join(data_dir, 'tgt')\n",
        "\n",
        "ref_path = os.path.join(gdir, ref_data)\n",
        "if len(ref_data) > 0 and os.path.isfile(ref_path):\n",
        "  !unzip -j -o -q $ref_path -d ref\n",
        "  ref_dir = os.path.join(data_dir, 'ref')\n",
        "  with_prior = True\n",
        "else:\n",
        "  with_prior = False\n",
        "\n",
        "%cd $work_dir\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Use either of methods:\n",
        "* [Textual inversion](https://textual-inversion.github.io) = adds new token to the text encoder. Generic but stable. Trained embeddings can be combined together on load.\n",
        "* [LoRA](https://github.com/cloneofsimo/lora) = partially finetunes low-rank add-ons, injected to Unet attention layers. Universal method, **industry standard**, precise, may interfere with wide spectrum of topics.\n",
        "* [Custom diffusion](https://github.com/adobe-research/custom-diffusion) = similar to LoRA (in a way). Can achieve impressive reproduction quality (including faces) with simple prompts, but may lose the point with too complex ones. To train it, you'll need to specify above both **target** reference images and **generic reference** ones (more random, of similar subjects). Apparently, you can generate the latter with SD itself.  \n",
        "\n",
        "Mark `style` if you're training for a style, rather than an object.  \n",
        "Mark `low_mem` if you get OOM.  "
      ],
      "metadata": {
        "id": "CsNRhrxr4aKN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run training\n",
        "method = \"text\" #@param ['text','lora','custom']\n",
        "style = False #@param  {type:\"boolean\"}\n",
        "low_mem = False #@param  {type:\"boolean\"}\n",
        "train_steps = 3000 #@param  {type:\"integer\"}\n",
        "save_steps = 500 #@param  {type:\"integer\"}\n",
        "batch = 1 #@param  {type:\"integer\"}\n",
        "\n",
        "command = ' --data %s ' % tgt_dir\n",
        "if method in ['text','custom']:\n",
        "  command += ' -t %s --term %s ' % (new_token, ref_class)\n",
        "if method in ['custom']:\n",
        "  command += ' --term_data %s ' % ref_dir\n",
        "\n",
        "command += ' -ts %d --save_step %d -b %d ' %(train_steps, save_steps, batch)\n",
        "command += ' -val -md %s ' % maindir\n",
        "command += ' --model %s ' % modname\n",
        "\n",
        "if low_mem:\n",
        "  command += ' --low_mem '\n",
        "if style:\n",
        "  command += ' --style '\n",
        "\n",
        "if method == 'text':\n",
        "  %run src/train.py $command -lr 0.001 --type text\n",
        "elif method == 'lora':\n",
        "  %run src/train.py $command -lr 0.0002 --type lora\n",
        "elif method == 'custom':\n",
        "  %run src/train.py $command --type custom\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "AW5--vadmcYb"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}