sd-colab-toolkit.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Linaqruf/sd-notebook-collection/blob/main/sd-colab-toolkit.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# I. Installation"
      ],
      "metadata": {
        "id": "akFjukqz5PbB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 1.1 Install Colab Toolkit\n",
        "#@markdown This will install required Python packages\n",
        "import os\n",
        "import zipfile\n",
        "import shutil\n",
        "import subprocess\n",
        "import requests\n",
        "import re\n",
        "import time\n",
        "from tqdm import tqdm\n",
        "from urllib.parse import urlparse, unquote\n",
        "\n",
        "root_dir    = \"/content\"\n",
        "repo_dir    = os.path.join(root_dir, \"sd-scripts\")\n",
        "models_dir  = os.path.join(root_dir, \"models\")\n",
        "vaes_dir    = os.path.join(root_dir, \"vae\")\n",
        "lora_dir    = os.path.join(root_dir, \"network_weight\")\n",
        "deps_dir    = os.path.join(root_dir, \"deps\")\n",
        "drive_dir   = os.path.join(root_dir, \"drive/MyDrive\")\n",
        "tools_dir   = os.path.join(repo_dir, \"tools\")\n",
        "\n",
        "repo_url    = \"https://github.com/kohya-ss/sd-scripts\"\n",
        "\n",
        "def cprint(*args, color=\"default\", reset=True, tqdm_desc=False):\n",
        "    color_codes = {\n",
        "        \"default\"     : \"\\033[0m\",\n",
        "        \"green\"       : \"\\033[0;32m\",\n",
        "        \"red\"         : \"\\033[0;31m\",\n",
        "        \"bold_green\"  : \"\\033[1;32m\",\n",
        "        \"bold_red\"    : \"\\033[1;31m\",\n",
        "    }\n",
        "    \n",
        "    if color in color_codes:\n",
        "        color_start = color_codes[color]\n",
        "        color_end = \"\"\n",
        "        if reset:\n",
        "            color_end = color_codes[\"default\"]\n",
        "        formatted_text = \" \".join(str(arg) for arg in args)\n",
        "        if tqdm_desc:\n",
        "            color_end = \"\"\n",
        "            return color_start + formatted_text + color_end\n",
        "        else:\n",
        "            print(color_start + formatted_text + color_end)\n",
        "    else:\n",
        "        if tqdm_desc:\n",
        "            return \" \".join(str(arg) for arg in args)\n",
        "        else:\n",
        "            print(*args)\n",
        "\n",
        "def clone_repo(url, dir):\n",
        "    if not os.path.exists(dir):\n",
        "        subprocess.run([\"git\", \"clone\", url, dir], check=True)\n",
        "\n",
        "def ubuntu_deps(url, dst, desc):\n",
        "    os.makedirs(dst, exist_ok=True)\n",
        "    filename  = get_filename(url)\n",
        "    subprocess.run([\"wget\", url], stdout=subprocess.DEVNULL)\n",
        "\n",
        "    with zipfile.ZipFile(filename, \"r\") as deps:\n",
        "        deps.extractall(dst)\n",
        "\n",
        "    for file in tqdm(os.listdir(dst), desc=desc):\n",
        "        if file.endswith(\".deb\"):\n",
        "            subprocess.run([\"dpkg\", \"-i\", os.path.join(dst, file)], stdout=subprocess.DEVNULL)\n",
        "    \n",
        "    os.remove(filename)\n",
        "    shutil.rmtree(dst)\n",
        "\n",
        "def get_filename(url):\n",
        "    response = requests.get(url, stream=True)\n",
        "    response.raise_for_status()\n",
        "\n",
        "    if 'content-disposition' in response.headers:\n",
        "        content_disposition = response.headers['content-disposition']\n",
        "        filename = re.findall('filename=\"?([^\"]+)\"?', content_disposition)[0]\n",
        "    else:\n",
        "        url_path = urlparse(url).path\n",
        "        filename = unquote(os.path.basename(url_path))\n",
        "        \n",
        "    return filename\n",
        "\n",
        "def update_requirements(repo_dir, desired_module, desired_version, filepath):\n",
        "    with open(filepath, \"r\") as f:\n",
        "        lines = f.readlines()\n",
        "\n",
        "    updated_lines = []\n",
        "    for line in lines:\n",
        "        if desired_module in line:\n",
        "            line = f\"{desired_module}=={desired_version}\\n\"\n",
        "        updated_lines.append(line)\n",
        "\n",
        "    with open(filepath, \"w\") as f:\n",
        "      f.writelines(updated_lines)\n",
        "\n",
        "def install_dependencies():\n",
        "    requirements_file = os.path.join(repo_dir, \"requirements.txt\")\n",
        "\n",
        "    ram_patch_url = \"https://huggingface.co/Linaqruf/fast-repo/resolve/main/ram_patch.zip\"\n",
        "\n",
        "    cprint(f\"Installing ubuntu dependencies...\", color=\"green\")\n",
        "    subprocess.run([\"apt\", \"install\", \"aria2\", \"lz4\", \"libunwind8-dev\", \"-y\"], check=True)\n",
        "\n",
        "    ubuntu_deps(ram_patch_url, deps_dir, cprint(\"Installing RAM allocation patch\", color=\"green\", tqdm_desc=True))\n",
        "    \n",
        "    cprint(f\"Installing requirements...\", color=\"green\")\n",
        "    update_requirements(repo_dir, \"requests\", \"2.27.1\", requirements_file)\n",
        "    subprocess.run(['pip', 'install', '--upgrade', '--no-cache-dir', 'gdown'], check=True)\n",
        "    subprocess.run(['pip', 'install', '--upgrade', '-r', requirements_file], cwd=repo_dir, check=True)\n",
        "\n",
        "def calculate_elapsed_time(start_time):\n",
        "    end_time = time.time()\n",
        "    elapsed_time = int(end_time - start_time)\n",
        "\n",
        "    if elapsed_time < 60:\n",
        "        return f\"{elapsed_time} sec\"\n",
        "    else:\n",
        "        mins, secs = divmod(elapsed_time, 60)\n",
        "        return f\"{mins} mins {secs} sec\"\n",
        "\n",
        "def main():\n",
        "    os.chdir(root_dir)\n",
        "    start_time = time.time()\n",
        "\n",
        "    for dir in [models_dir, vaes_dir]:\n",
        "        os.makedirs(dir, exist_ok=True)\n",
        "\n",
        "    cprint(f\"Installing 'kohya-ss/sd-scripts'...\", color=\"green\")\n",
        "    clone_repo(repo_url, repo_dir)\n",
        "    install_dependencies()\n",
        "\n",
        "    elapsed_time = calculate_elapsed_time(start_time)\n",
        "    \n",
        "    cprint(f\"\\nFinished installation. Took {elapsed_time}.\", color=\"green\")\n",
        "    cprint(f\"All is done! Go to the next step.\", color=\"green\")\n",
        "\n",
        "main()\n"
      ],
      "metadata": {
        "id": "0BicRIFqIjg0",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **1.2. Download Model**\n",
        "import os\n",
        "import re\n",
        "import json\n",
        "import glob\n",
        "import gdown\n",
        "import time\n",
        "import requests\n",
        "import subprocess\n",
        "from IPython.utils import capture\n",
        "from urllib.parse import urlparse, unquote\n",
        "from pathlib import Path\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown ### **Download Stable Diffusion Model**\n",
        "\n",
        "model_url = \"https://civitai.com/api/download/models/82446\" #@param [\"Anime Model\", \"Anything V3.1\", \"AnyLoRA\", \"ChilloutMix Ni\", \"Stable Diffusion V1.5\", \"Replicant V3\", \"Illuminati Diffusion V1.1\", \"Waifu Diffusion V1.5 Beta 3\", \"Stable Diffusion V2.1\"] {allow-input: true}\n",
        "# @markdown ### **Download VAE Model**\n",
        "vae_url = \"\" #@param [\"\", \"Anime / Anything VAE\", \"Blessed VAE\", \"Waifu Diffusion VAE\", \"Stable Diffusion VAE\"] {allow-input: true}\n",
        "# @markdown ### **Download LoRA Model**\n",
        "lora_url = \"\" #@param {type: \"string\"}\n",
        "\n",
        "available_models = {\n",
        "    # SDv1.x Pretrained Model\n",
        "    \"Anime Model\"                 : \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/models/animefull-final-pruned.ckpt\",\n",
        "    \"Anything V3.1\"               : \"https://huggingface.co/Linaqruf/anything-v3.0/resolve/main/anything-v3-fp16-pruned.safetensors\",\n",
        "    \"AnyLoRA\"                     : \"https://huggingface.co/Lykon/AnyLoRA/resolve/main/AnyLoRA_noVae_fp16-pruned.ckpt\",\n",
        "    \"ChilloutMix Ni\"              : \"https://huggingface.co/naonovn/chilloutmix_NiPrunedFp32Fix/resolve/main/chilloutmix_NiPrunedFp32Fix.safetensors\",\n",
        "    \"Stable Diffusion V1.5\"       : \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/stable_diffusion_1_5-pruned.safetensors\",\n",
        "    # SDv2.x Pretrained Model\n",
        "    \"Replicant V3\"                : \"https://huggingface.co/gsdf/Replicant-V3.0/resolve/main/Replicant-V3.0_fp16.safetensors\",\n",
        "    \"Illuminati Diffusion V1.1\"   : \"https://huggingface.co/Linaqruf/stolen/resolve/main/pruned-models/illuminatiDiffusionV1_v11.safetensors\",\n",
        "    \"Waifu Diffusion V1.5 Beta 3\" : \"https://huggingface.co/waifu-diffusion/wd-1-5-beta3/resolve/main/wd-beta3-base-fp16.safetensors\",\n",
        "    \"Stable Diffusion V2.1\"       : \"https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.safetensors\",\n",
        "}\n",
        "\n",
        "available_vaes = {\n",
        "    \"Anime / Anything VAE\"        : \"https://huggingface.co/Linaqruf/personal-backup/resolve/main/vae/animevae.pt\",\n",
        "    \"Blessed VAE\"                 : \"https://huggingface.co/NoCrypt/blessed_vae/resolve/main/blessed2.vae.pt\",\n",
        "    \"Waifu Diffusion VAE\"         : \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime.ckpt\",\n",
        "    \"Stable Diffusion VAE\"        : \"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt\",\n",
        "}\n",
        "\n",
        "if model_url is not None:\n",
        "    valid_model_url = model_url\n",
        "    if model_url in available_models:\n",
        "        valid_model_url = available_models[model_url]\n",
        "\n",
        "if vae_url is not None:\n",
        "    valid_vae_url = vae_url\n",
        "    if vae_url in available_vaes:\n",
        "        valid_vae_url = available_vaes[vae_url]\n",
        "\n",
        "def get_supported_extensions():\n",
        "    return tuple([\".ckpt\", \".safetensors\", \".pt\", \".pth\"])\n",
        "\n",
        "def get_filename(url, quiet=True):\n",
        "    extensions = get_supported_extensions()\n",
        "\n",
        "    if url.startswith(\"/content/drive/MyDrive/\") or url.endswith(tuple(extensions)):\n",
        "        filename = os.path.basename(url)\n",
        "    else:\n",
        "        response = requests.get(url, stream=True)\n",
        "        response.raise_for_status()\n",
        "\n",
        "        if 'content-disposition' in response.headers:\n",
        "            content_disposition = response.headers['content-disposition']\n",
        "            filename = re.findall('filename=\"?([^\"]+)\"?', content_disposition)[0]\n",
        "        else:\n",
        "            url_path = urlparse(url).path\n",
        "            filename = unquote(os.path.basename(url_path))\n",
        "\n",
        "    if filename is not None and filename.endswith(tuple(extensions)):\n",
        "        return filename\n",
        "    else:\n",
        "        return None\n",
        "\n",
        "def parse_args(config):\n",
        "    args = []\n",
        "\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args.append(f\"{v}\")\n",
        "        elif isinstance(v, str) and v is not None:\n",
        "            args.append(f'--{k}={v}')\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args.append(f\"--{k}\")\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args.append(f\"--{k}={v}\")\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args.append(f\"--{k}={v}\")\n",
        "\n",
        "    return args\n",
        "\n",
        "def aria2_download(dir, filename, url):\n",
        "    hf_token    = \"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\"\n",
        "    user_header = f\"Authorization: Bearer {hf_token}\"\n",
        "\n",
        "    aria2_config = {\n",
        "        \"console-log-level\"         : \"error\",\n",
        "        \"summary-interval\"          : 10,\n",
        "        \"header\"                    : user_header if \"huggingface.co\" in url else None,\n",
        "        \"continue\"                  : True,\n",
        "        \"max-connection-per-server\" : 16,\n",
        "        \"min-split-size\"            : \"1M\",\n",
        "        \"split\"                     : 16,\n",
        "        \"dir\"                       : dir,\n",
        "        \"out\"                       : filename,\n",
        "        \"_url\"                      : url,\n",
        "    }\n",
        "    aria2_args = parse_args(aria2_config)\n",
        "    subprocess.run([\"aria2c\", *aria2_args])\n",
        "\n",
        "def gdown_download(url, dst):\n",
        "    if \"/uc?id/\" in url:\n",
        "        return gdown.download(url, dst + \"/\", quiet=False)\n",
        "    elif \"/file/d/\" in url:\n",
        "        return gdown.download(url, dst + \"/\", quiet=False, fuzzy=True)\n",
        "    elif \"/drive/folders/\" in url:\n",
        "        os.chdir(dst)\n",
        "        return gdown.download_folder(url, quiet=True, use_cookies=False)\n",
        "\n",
        "def download(urls, dst, target):\n",
        "    for url in tqdm(urls, desc=cprint(f\"Downloading {target} from url\", color=\"green\", tqdm_desc=True)):\n",
        "        with capture.capture_output() as cap:\n",
        "            url = url.replace(\" \", \"\")\n",
        "            try:\n",
        "                filename = get_filename(url, quiet=False)\n",
        "            except Exception:\n",
        "                filename = None\n",
        "                continue\n",
        "\n",
        "            if \"drive.google.com\" in url:\n",
        "                try:\n",
        "                    gdown = gdown_download(url, dst)\n",
        "                except Exception as e:\n",
        "                    print(f\"Error occurred: {str(e)}\")\n",
        "            elif url.startswith(\"/content/drive/MyDrive/\"):\n",
        "                filepath = os.path.join(dst, filename)\n",
        "                Path(filepath).write_bytes(Path(url).read_bytes())\n",
        "            else:\n",
        "                if \"huggingface.co\" in url:\n",
        "                    if \"/blob/\" in url:\n",
        "                        url = url.replace(\"/blob/\", \"/resolve/\")\n",
        "                aria2_download(dst, filename, url)\n",
        "        \n",
        "def main():\n",
        "    global model_path, vae_path, lora_path\n",
        "\n",
        "    start_time = time.time()\n",
        "    model_path = vae_path = lora_path = None\n",
        "\n",
        "    download_targets = {\n",
        "        \"model\": (valid_model_url.split(','), models_dir),\n",
        "        \"vae\": (valid_vae_url.split(','), vaes_dir),\n",
        "        \"lora\": (lora_url.split(','), lora_dir),\n",
        "    }\n",
        "\n",
        "    for target, (urls, dst) in download_targets.items():\n",
        "        if urls and urls != \"PASTE {} URL OR GDRIVE PATH HERE\".format(target.upper()):\n",
        "            initial_files = glob.glob(os.path.join(dst, \"*\"))\n",
        "            download(urls, dst, target)\n",
        "\n",
        "            downloaded_files = []\n",
        "            for filename in initial_files:\n",
        "                filepath = os.path.join(dst, filename)\n",
        "                if os.path.exists(filepath):\n",
        "                    downloaded_files.append(filepath)\n",
        "\n",
        "            if len(downloaded_files) == 0:\n",
        "                downloaded_files = sorted(\n",
        "                    glob.glob(os.path.join(dst, \"*\")), key=os.path.getmtime, reverse=True\n",
        "                )\n",
        "\n",
        "            cprint(f\"Downloaded files for {target}:\", color=\"green\")\n",
        "            for file in downloaded_files:\n",
        "                cprint(\"  - \", os.path.basename(file), color=\"green\")\n",
        "\n",
        "    elapsed_time = calculate_elapsed_time(start_time)\n",
        "    \n",
        "    cprint(f\"\\nFinished installation. Took {elapsed_time}.\", color=\"green\")\n",
        "    cprint(f\"All is done! Go to the next step.\", color=\"green\")\n",
        "\n",
        "main()\n"
      ],
      "metadata": {
        "id": "A2NC0tm8vU-K",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 2. Model Conversion"
      ],
      "metadata": {
        "id": "WDPfF4uc5pd1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "%store -r\n",
        "#@title ## 7.2. Model Pruner\n",
        "\n",
        "os.chdir(tools_dir)\n",
        "\n",
        "if not os.path.exists('prune.py'):\n",
        "    !wget https://raw.githubusercontent.com/lopho/stable-diffusion-prune/main/prune.py\n",
        "\n",
        "#@markdown Convert to Float16\n",
        "fp16 = True #@param {'type':'boolean'}\n",
        "#@markdown Use EMA for weights\n",
        "ema = False #@param {'type':'boolean'}\n",
        "#@markdown Strip CLIP weights\n",
        "no_clip = False #@param {'type':'boolean'}\n",
        "#@markdown Strip VAE weights\n",
        "no_vae = False #@param {'type':'boolean'}\n",
        "#@markdown Strip depth model weights\n",
        "no_depth = False #@param {'type':'boolean'}\n",
        "#@markdown Strip UNet weights\n",
        "no_unet = False #@param {'type':'boolean'}\n",
        "\n",
        "model_path = \"/content/models/Counterfeit-V3.0_fp16.safetensors\" #@param {'type' : 'string'}\n",
        "\n",
        "config = {\n",
        "    \"fp16\": fp16,\n",
        "    \"ema\": ema,\n",
        "    \"no_clip\": no_clip,\n",
        "    \"no_vae\": no_vae,\n",
        "    \"no_depth\": no_depth,\n",
        "    \"no_unet\": no_unet,\n",
        "}\n",
        "\n",
        "suffixes = {\n",
        "    \"fp16\": \"-fp16\",\n",
        "    \"ema\": \"-ema\",\n",
        "    \"no_clip\": \"-no-clip\",\n",
        "    \"no_vae\": \"-no-vae\",\n",
        "    \"no_depth\": \"-no-depth\",\n",
        "    \"no_unet\": \"-no-unet\",\n",
        "}\n",
        "\n",
        "print(f\"Loading model from {model_path}\")\n",
        "\n",
        "dir_name = os.path.dirname(model_path)\n",
        "base_name = os.path.basename(model_path)\n",
        "output_name = base_name.split('.')[0]\n",
        "\n",
        "for option, suffix in suffixes.items():\n",
        "    if config[option]:\n",
        "        print(f\"Applying option {option}\")\n",
        "        output_name += suffix\n",
        "        \n",
        "output_name += '-pruned'\n",
        "output_path = os.path.join(dir_name, output_name + ('.ckpt' if model_path.endswith(\".ckpt\") else \".safetensors\"))\n",
        "\n",
        "args = \"\"\n",
        "for k, v in config.items():\n",
        "    if k.startswith(\"_\"):\n",
        "        args += f'\"{v}\" '\n",
        "    elif isinstance(v, str):\n",
        "        args += f'--{k}=\"{v}\" '\n",
        "    elif isinstance(v, bool) and v:\n",
        "        args += f\"--{k} \"\n",
        "    elif isinstance(v, float) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "    elif isinstance(v, int) and not isinstance(v, bool):\n",
        "        args += f\"--{k}={v} \"\n",
        "\n",
        "final_args = f\"python3 prune.py {model_path} {output_path} {args}\"\n",
        "!{final_args}\n",
        "\n",
        "print(f\"Saving pruned model to {output_path}\")"
      ],
      "metadata": {
        "cellView": "form",
        "id": "fb3rxuCaSYta"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## 7.1. Convert Diffusers to Checkpoint\n",
        "import os\n",
        "%store -r\n",
        "\n",
        "os.chdir(tools_dir)\n",
        "\n",
        "#@markdown ### Conversion Config\n",
        "model_to_load = \"\" #@param {'type': 'string'}\n",
        "model_to_save = os.path.splitext(model_to_load)[0]\n",
        "convert = \"checkpoint_to_diffusers\" #@param [\"diffusers_to_checkpoint\", \"checkpoint_to_diffusers\"] {'allow-input': false}\n",
        "v2 = True #@param {type:'boolean'}\n",
        "global_step = 0 #@param {'type': 'number'}\n",
        "epoch = 0 #@param {'type': 'number'}\n",
        "use_safetensors = True #@param {'type': 'boolean'}\n",
        "save_precision_as = \"--float\" #@param [\"--fp16\",\"--bf16\",\"--float\"] {'allow-input': false}\n",
        "\n",
        "#@markdown Additional option for diffusers\n",
        "feature_extractor = True #@param {'type': 'boolean'}\n",
        "safety_checker = True #@param {'type': 'boolean'}\n",
        "\n",
        "reference_model = \"stabilityai/stable-diffusion-2-1\" if v2 else \"runwayml/stable-diffusion-v1-5\" \n",
        "model_output = f\"{model_to_save}.safetensors\" if use_safetensors else f\"{model_to_save}.ckpt\"\n",
        "\n",
        "urls = [\n",
        "    (\"preprocessor_config.json\", \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/preprocessor_config.json\"),\n",
        "    (\"config.json\", \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/config.json\"),\n",
        "    (\"pytorch_model.bin\", \"https://huggingface.co/CompVis/stable-diffusion-safety-checker/resolve/main/pytorch_model.bin\"),\n",
        "]\n",
        "\n",
        "diffusers_to_sd_dict = {\n",
        "    \"_model_to_load\": model_to_load,\n",
        "    \"_model_to_save\": model_output,\n",
        "    \"global_step\": global_step,\n",
        "    \"epoch\": epoch,\n",
        "    \"save_precision_as\": save_precision_as,\n",
        "}\n",
        "\n",
        "sd_to_diffusers_dict = {\n",
        "    \"_model_to_load\": model_to_load,\n",
        "    \"_model_to_save\": model_to_save,\n",
        "    \"v2\": True if v2 else False,\n",
        "    \"v1\": True if not v2 else False,\n",
        "    \"global_step\": global_step,\n",
        "    \"epoch\": epoch,\n",
        "    \"fp16\": True if save_precision_as == \"fp16\" else False,\n",
        "    \"use_safetensors\": use_safetensors,\n",
        "    \"reference_model\": reference_model\n",
        "}\n",
        "\n",
        "def convert_dict(config):\n",
        "    args = \"\"\n",
        "    for k, v in config.items():\n",
        "        if k.startswith(\"_\"):\n",
        "            args += f'\"{v}\" '\n",
        "        elif isinstance(v, str):\n",
        "            args += f'--{k}=\"{v}\" '\n",
        "        elif isinstance(v, bool) and v:\n",
        "            args += f\"--{k} \"\n",
        "        elif isinstance(v, float) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "        elif isinstance(v, int) and not isinstance(v, bool):\n",
        "            args += f\"--{k}={v} \"\n",
        "\n",
        "    return args\n",
        "\n",
        "def run_script(script_name, script_args):\n",
        "    !python {script_name} {script_args}\n",
        "\n",
        "def download(output, url, save_dir):\n",
        "    !aria2c --console-log-level=error --summary-interval=10 -c -x 16 -k 1M -s 16 -d '{save_dir}' -o '{output}' {url}\n",
        "\n",
        "diffusers_to_sd_args = convert_dict(diffusers_to_sd_dict)\n",
        "sd_to_diffusers_args = convert_dict(sd_to_diffusers_dict)\n",
        "\n",
        "if convert == \"diffusers_to_checkpoint\":\n",
        "    if model_to_load.endswith((\"ckpt\",\"safetensors\")):\n",
        "        print(f\"{os.path.basename(model_to_load)} is not in diffusers format\")\n",
        "    else:\n",
        "        run_script(\"convert_diffusers20_original_sd.py\", diffusers_to_sd_args)\n",
        "else:\n",
        "    if not model_to_load.endswith((\"ckpt\",\"safetensors\")):\n",
        "        print(f\"{os.path.basename(model_to_load)} is not in ckpt/safetensors format\")\n",
        "    else:     \n",
        "        run_script(\"convert_diffusers20_original_sd.py\", sd_to_diffusers_args)\n",
        "\n",
        "        if feature_extractor:\n",
        "            save_dir = os.path.join(model_to_save, \"feature_extractor\")\n",
        "            os.makedirs(save_dir, exist_ok=True)\n",
        "            output, url = urls[0]\n",
        "            download(output, url, save_dir)\n",
        "            \n",
        "        if safety_checker:\n",
        "            save_dir = os.path.join(model_to_save, \"safety_checker\")\n",
        "            os.makedirs(save_dir, exist_ok=True)\n",
        "            for output, url in urls[1:]:\n",
        "                download(output, url, save_dir)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "TdCb8_dSSzzr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 2.3. Replace VAE of Existing Model \n",
        "\n",
        "os.chdir(tools_dir)\n",
        "if not os.path.exists('merge_vae.py'):\n",
        "  !wget https://raw.githubusercontent.com/Linaqruf/kohya-trainer/main/tools/merge_vae.py\n",
        "\n",
        "#@markdown You need to input model ends with `.ckpt`, because `.safetensors` model won't work.\n",
        "\n",
        "target_model = \"\" #@param {'type': 'string'}\n",
        "target_vae = \"/content/vae/anime.vae.pt\" #@param {'type': 'string'}\n",
        "use_safetensors = False #@param {type:'boolean'}\n",
        "# get the base file name and directory\n",
        "base_name = os.path.basename(target_model)\n",
        "base_dir = os.path.dirname(target_model)\n",
        "\n",
        "# get the file name without extension\n",
        "file_name = os.path.splitext(base_name)[0]\n",
        "\n",
        "# create the new file name\n",
        "new_file_name = file_name + \"-vae-swapped\"\n",
        "\n",
        "# get the file extension\n",
        "file_ext = os.path.splitext(base_name)[1]\n",
        "\n",
        "# create the output file path\n",
        "output_model = os.path.join(base_dir, new_file_name + file_ext)\n",
        "\n",
        "!python merge_vae.py \\\n",
        "    {target_model} \\\n",
        "    {target_vae} \\\n",
        "    {output_model}\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "g2QfhhlfbGsf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 2.4. Convert CKPT-2-Safetensors\n",
        "\n",
        "import os\n",
        "import torch\n",
        "from safetensors.torch import load_file, save_file\n",
        "from torch import load, save\n",
        "\n",
        "model_path = \"\" #@param {type: 'string'}\n",
        "\n",
        "def is_safetensors(path):\n",
        "  return os.path.splitext(path)[1].lower() == '.safetensors'\n",
        "\n",
        "def convert(model_path):\n",
        "  print(\"Loading model:\", os.path.basename(model_path))\n",
        "  \n",
        "  try:\n",
        "      with torch.no_grad():\n",
        "          print(\"Conversion in progress, please wait...\")\n",
        "          if is_safetensors(model_path):\n",
        "            model = load_file(model_path, device=\"cpu\")\n",
        "          else:\n",
        "            model = load(model_path, map_location=\"cpu\")\n",
        "          \n",
        "          if 'state_dict' in model:\n",
        "            sd = model['state_dict']\n",
        "          else:\n",
        "            sd = model\n",
        "\n",
        "          save_to = \".ckpt\" if is_safetensors(model_path) else \".safetensors\"\n",
        "          output = os.path.splitext(model_path)[0] + save_to\n",
        "\n",
        "          if is_safetensors(model_path):\n",
        "            save(sd, output)\n",
        "          else:\n",
        "            save_file(sd, output)\n",
        "\n",
        "      print(f'Successfully converted {os.path.basename(model_path)} to {os.path.basename(output)}')\n",
        "      print(f'located in this path : {output}')\n",
        "  except Exception as ex:\n",
        "      print(f'ERROR converting {os.path.basename(model_path)}: {ex}')\n",
        "\n",
        "  print('Done!')\n",
        "\n",
        "def main():\n",
        "  convert(model_path)\n",
        "main()\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "L8WI17pmnlVM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# VIII. Deployment"
      ],
      "metadata": {
        "id": "nyIl9BhNXKUq"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## 7.1. Upload Config\n",
        "from huggingface_hub import login\n",
        "from huggingface_hub import HfApi\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "\n",
        "# @markdown Login to Huggingface Hub\n",
        "# @markdown > Get **your** huggingface `WRITE` token [here](https://huggingface.co/settings/tokens)\n",
        "write_token = \"\"  # @param {type:\"string\"}\n",
        "# @markdown Fill this if you want to upload to your organization, or just leave it empty.\n",
        "orgs_name = \"\"  # @param{type:\"string\"}\n",
        "# @markdown If your model/dataset repo does not exist, it will automatically create it.\n",
        "repo_name = \"stolen\"  # @param{type:\"string\"}\n",
        "make_private = False  # @param{type:\"boolean\"}\n",
        "\n",
        "def authenticate(write_token):\n",
        "    login(write_token, add_to_git_credential=True)\n",
        "    api = HfApi()\n",
        "    return api.whoami(write_token), api\n",
        "\n",
        "def create_repo(api, user, orgs_name, repo_name, make_private=False):\n",
        "    if orgs_name == \"\":\n",
        "        repo_id = user[\"name\"] + \"/\" + repo_name.strip()\n",
        "    else:\n",
        "        repo_id = orgs_name + \"/\" + repo_name.strip()\n",
        "\n",
        "    try:\n",
        "        validate_repo_id(repo_id)\n",
        "        api.create_repo(repo_id=repo_id, private=make_private)\n",
        "        print(f\"Model repo '{repo_id}' didn't exist, creating repo\")\n",
        "    except HfHubHTTPError as e:\n",
        "        print(f\"Model repo '{repo_id}' exists, skipping create repo\")\n",
        "    \n",
        "    print(f\"Model repo '{repo_id}' link: https://huggingface.co/{repo_id}\\n\")\n",
        "    return repo_id\n",
        "\n",
        "user, api = authenticate(write_token)\n",
        "\n",
        "model_repo = create_repo(api, user, orgs_name, repo_name, make_private)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "QTXsM170GUpk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 8.2. Upload with Huggingface Hub"
      ],
      "metadata": {
        "id": "Fuxghk8MnG6j"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ### 8.2.1. Upload Model\n",
        "from huggingface_hub import HfApi\n",
        "from huggingface_hub.utils import validate_repo_id, HfHubHTTPError\n",
        "from pathlib import Path\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "#@markdown This will be uploaded to model repo\n",
        "\n",
        "model_path = \"/content/models\" #@param {type :\"string\"}\n",
        "path_in_repo = \"fp16\" #@param {type :\"string\"}\n",
        "revision = \"\" #@param {type :\"string\"}\n",
        "if revision:\n",
        "  api.create_branch(repo_id=model_repo, \n",
        "                branch=revision, \n",
        "                exist_ok=True)\n",
        "else:\n",
        "  revision = \"main\"\n",
        "project_name = os.path.basename(model_path)\n",
        "if project_name in [\".safetensors\", \"ckpt\", \"pt\"]:\n",
        "  project_name = os.path.split(model_path)[0]\n",
        "# @markdown Other Information\n",
        "commit_message = \"\"  # @param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "    commit_message = \"feat: upload \" + project_name + \" checkpoint\"\n",
        "\n",
        "if os.path.exists(model_path):\n",
        "  vae_exists = os.path.exists(os.path.join(model_path, 'vae'))\n",
        "  unet_exists = os.path.exists(os.path.join(model_path, 'unet'))\n",
        "  text_encoder_exists = os.path.exists(os.path.join(model_path, 'text_encoder'))\n",
        "    \n",
        "def upload_model(model_paths, is_folder :bool, commit_message):\n",
        "  path_obj = Path(model_paths)\n",
        "  trained_model = path_obj.parts[-1]\n",
        "  \n",
        "  if path_in_repo:\n",
        "    trained_model = path_in_repo\n",
        "    \n",
        "  if is_folder == True:\n",
        "    print(f\"Uploading {trained_model} to https://huggingface.co/\"+model_repo)\n",
        "    print(f\"Please wait...\")\n",
        "\n",
        "    if vae_exists and unet_exists and text_encoder_exists:\n",
        "      if not commit_message:\n",
        "        commit_message = f\"feat: upload diffusers version of {trained_model}\"\n",
        "\n",
        "      api.upload_folder(\n",
        "          folder_path=model_paths,\n",
        "          repo_id=model_repo,\n",
        "          revision=revision,\n",
        "          commit_message=commit_message,\n",
        "          ignore_patterns=\".ipynb_checkpoints\"\n",
        "          )\n",
        "    \n",
        "    else:\n",
        "      if not commit_message:\n",
        "        commit_message = f\"feat: upload {trained_model} checkpoint folder\"\n",
        "\n",
        "      api.upload_folder(\n",
        "          folder_path=model_paths,\n",
        "          path_in_repo=trained_model,\n",
        "          repo_id=model_repo,\n",
        "          revision=revision,\n",
        "          commit_message=commit_message,\n",
        "          ignore_patterns=\".ipynb_checkpoints\"\n",
        "          )\n",
        "    print(f\"Upload success, located at https://huggingface.co/\"+model_repo+\"/tree/main\\n\")\n",
        "  else: \n",
        "    print(f\"Uploading {trained_model} to https://huggingface.co/\"+model_repo)\n",
        "    print(f\"Please wait...\")\n",
        "    if not commit_message:\n",
        "      if model_paths.endswith(\".safetensors\"):\n",
        "        commit_message = f\"feat: upload safetensors version of {trained_model} \"\n",
        "      else:\n",
        "        commit_message = f\"feat: upload {trained_model} checkpoint\"\n",
        "            \n",
        "    api.upload_file(\n",
        "        path_or_fileobj=model_paths,\n",
        "        path_in_repo=trained_model,\n",
        "        repo_id=model_repo,\n",
        "        revision=revision,\n",
        "        commit_message=commit_message,\n",
        "        )\n",
        "        \n",
        "    print(f\"Upload success, located at https://huggingface.co/\"+model_repo+\"/blob/main/\"+trained_model+\"\\n\")\n",
        "      \n",
        "def upload():\n",
        "    if model_path.endswith((\".ckpt\", \".safetensors\", \".pt\")):\n",
        "      upload_model(model_path, False, commit_message)\n",
        "    else:\n",
        "      upload_model(model_path, True, commit_message)\n",
        "\n",
        "upload()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "CIeoJA-eO-8t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 8.3. Upload with GIT (Alternative)"
      ],
      "metadata": {
        "id": "CKZpg4keWS5c"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ### 8.3.1. Clone Repository\n",
        "%cd /content/\n",
        "clone_model = True #@param {'type': 'boolean'}\n",
        "\n",
        "!git lfs install --skip-smudge\n",
        "!export GIT_LFS_SKIP_SMUDGE=1\n",
        "\n",
        "if clone_model:\n",
        "  !git clone https://huggingface.co/{model_repo} /content/{repo_name}"
      ],
      "metadata": {
        "cellView": "form",
        "id": "6nBlrOrytO9F"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ### 8.3.2. Commit using Git \n",
        "import os\n",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "#@markdown Choose which repo you want to commit\n",
        "commit_model = True #@param {'type': 'boolean'}\n",
        "#@markdown #### Other Information\n",
        "commit_message = \"\" #@param {type :\"string\"}\n",
        "\n",
        "if not commit_message:\n",
        "  commit_message = f\"feat: upload {repo_name}\"\n",
        "\n",
        "!git config --global user.email \"example@mail.com\"\n",
        "!git config --global user.name \"example\"\n",
        "\n",
        "def commit(repo_folder, commit_message):\n",
        "  os.chdir(os.path.join(root_dir, repo_folder))\n",
        "  !git lfs install\n",
        "  !huggingface-cli lfs-enable-largefiles .\n",
        "  !git add .\n",
        "  !git commit -m \"{commit_message}\"\n",
        "  !git push\n",
        "\n",
        "commit(repo_name, commit_message)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "7bJev4PzOFFB"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}