sdxl_model_converter_duskfallcrew_jupyter.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dY-0HMB7VyNu"
   },
   "source": [
    "# **SDXL Model Converter Jupyter Edition**\n",
    "A Colab Notebook To Convert SDXL Checkpoint to Diffusers format\n",
    "Originally coded by [Linaqruf](https://github.com/Linaqruf?tab=repositories)\n",
    "\n",
    "Ported from Linaqruf's SDXL one, and added extra functionality from the SD 1.5 converter.\n",
    "\n",
    "#### Summary:\n",
    "\n",
    "This notebook is a comprehensive tool for converting and uploading Diffusers models to the Hugging Face Hub. It provides a user-friendly interface for:\n",
    "\n",
    "- Converting Diffusers models to the Hugging Face format\n",
    "- Selecting and uploading folders containing the converted models to the Hugging Face Hub\n",
    "- Authenticating with the Hugging Face Hub using the notebook_login() function\n",
    "- Specifying repository details, including the Hugging Face username, repository name, and organization name (if applicable)\n",
    "- Choosing upload options, such as making the repository private or not\n",
    "- Monitoring upload progress and verifying the success of the upload\n",
    "- Cleaning folders\n",
    "- Fixing SDXL Keys.\n",
    "\n",
    "This is the JUPYTER edition that i'm working on, please note this will be largely incompatible with the colab runtime. \n",
    "Instructions are missing at the moment because i can't brain, and nor can Llama3b.\n",
    "\n",
    "----------------\n",
    "### **Important Note:**\n",
    "\n",
    "\n",
    "Before diving in, ensure you create a Hugging Face token with write permissions. Follow this link for instructions on token creation.\n",
    "\n",
    "You need to create a huggingface token, go to [this link](https://huggingface.co/settings/tokens), then `create new token` or copy available token with the `Write` role.\n",
    "\n",
    "-----------------\n",
    "\n",
    "\n",
    "| Link Name| Description | Link |\n",
    "| --- | --- | --- |\n",
    "| [Huggingface Backup](https://colab.research.google.com/github/kieranxsomer/HuggingFace_Backup/blob/main/HuggingFace_Backup.ipynb) | backup checkpoints! | [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/github/kieranxsomer/HuggingFace_Backup/blob/main/HuggingFace_Backup.ipynb)\n",
    "| [SD 1.5 Conversion to Diffusers](https://colab.research.google.com/drive/1zAzdsaa2KQcF6W0V4eCLZ6eUO8hsDJTo?usp=drive_link)| Convert SD 1.5 to Diffusers| [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/drive/1zAzdsaa2KQcF6W0V4eCLZ6eUO8hsDJTo?usp=drive_link)\n",
    "| [SDXL Conversion to Diffusers](https://colab.research.google.com/drive/1CcSCmUB_UkT-8TlUkwDDKnHB4T7nti01?usp=drive_link)| Convert SDXL to Diffusers| [![](https://img.shields.io/static/v1?message=Open%20in%20Colab&logo=googlecolab&labelColor=5c5c5c&color=0f80c1&label=%20&style=flat)](https://colab.research.google.com/drive/1CcSCmUB_UkT-8TlUkwDDKnHB4T7nti01?usp=drive_link)\n",
    "|Discord| E&D Discord |[Invite](https://discord.gg/5t2kYxt7An)\n",
    "|CivitAi| Duskfallcrew @ Civitai |[Duskfallcrew](https://civitai.com/user/duskfallcrew/)\n",
    "|Huggingface| E&D Huggingface |[Earth & Dusk](https://huggingface.co/EarthnDusk)\n",
    "|Ko-Fi| Kofi Support |[![ko-fi](https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat)](https://ko-fi.com/Z8Z8L4EO)\n",
    "|Github| Duskfallcrew Github |[Duskfallcrew](https://github.com/duskfallcrew)\n",
    "| Youtube: | Duskfall Music|[Duskfall Music & More](https://www.youtube.com/channel/UCk7MGP7nrJz5awBSP75xmVw)\n",
    "| Spotify: | E&D Royalty Free| [PLAYLIST](https://open.spotify.com/playlist/00R8x00YktB4u541imdSSf?si=57a8f0f0fe87434e)\n",
    "|DA Group | AI Group| [DeviantArt Group](https://www.deviantart.com/diffusionai)\n",
    "| Reddit | Earth & Dusk| [Subreddit](https://www.reddit.com/r/earthndusk/)\n",
    "\n",
    "\n",
    "\n",
    "> ## Collaboration\n",
    "\n",
    "\n",
    "I am NOT A programmer by nature, I patch with what little knowledge I have. I Failed programming several times over the years, so if something needs cleaning up and you want to patch it - pull request it!\n",
    "\n",
    "\n",
    ">## About\n",
    "\n",
    "\n",
    "We are a system of over 300 alters, proudly navigating life with Dissociative Identity Disorder, ADHD, Autism, and CPTSD. We believe in the potential of AI to break down barriers and enhance aspects of mental health, even as it presents challenges. Our creative journey is an ongoing exploration of identity and expression, and we invite you to join us in this adventure.\n",
    "\n",
    "\n",
    "\n",
    ">Future ideas:\n",
    "\n",
    "- Looking to port PT to safetensors into the same notebook.\n",
    "- Looking to figure out how to manage to get the conversions to roll on google drive for more space options.\n",
    "- Porting it to VastAi/Runpod\n",
    "- Looking to port inference for testing into the same notebook. - Not sure why you'd need this unless you're training, I think this was in the SDXL one from Linaqruf, but this isn't required on this one.\n",
    "- Looking to figure out how to convert a different vae. - I could NOT get this to work no matter what I did. I have some THEORIES, but i'm not a programmer. However due to the fact I have a SCRIPT for this now - this may be an option.\n",
    "\n",
    "\n",
    "\n",
    "\n",
    ">## Credits:\n",
    "\n",
    "\n",
    "| Patched Origin | Description | Link |\n",
    "| --- | --- | --- |\n",
    "|Patched from| ARCHIVED |[SDXL - Linaqruf](https://colab.research.google.com/github/Linaqruf/sdxl-model-converter/blob/main/sdxl_model_converter.ipynb)\n",
    "|***Linaqruf @ Github***: |https://github.com/Linaqruf\n",
    "|Linaqruf Ko-Fi | [![](https://dcbadge.vercel.app/api/shield/850007095775723532?style=flat)](https://lookup.guru/850007095775723532) [![ko-fi](https://img.shields.io/badge/Support%20me%20on%20Ko--fi-F16061?logo=ko-fi&logoColor=white&style=flat)](https://ko-fi.com/linaqruf)\n",
    "| Linaqruf Saweria |<a href=\"https://saweria.co/linaqruf\"><img alt=\"Saweria\" src=\"https://img.shields.io/badge/Saweria-7B3F00?style=flat&logo=ko-fi&logoColor=white\"/></a>\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  ♻ **Install Kohya Script**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "R3hfwYaXKn_V"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from subprocess import getoutput\n",
    "\n",
    "# Directories\n",
    "root_dir = \"/path/to/your/workspace\"  # Update this to your workspace directory\n",
    "models_dir = os.path.join(root_dir, \"models\")\n",
    "repo_dir = os.path.join(models_dir, \"kohya-trainer\")\n",
    "tools_dir = os.path.join(repo_dir, \"tools\")\n",
    "vae_dir = os.path.join(root_dir, \"vae\")\n",
    "\n",
    "# Repository details\n",
    "repo_url = \"https://github.com/kohya-ss/sd-scripts\" #@param {type:\"string\"}\n",
    "branch = \"sdxl\" #@param {type:\"string\"}\n",
    "\n",
    "def clone_repo(url, dir, branch):\n",
    "    if not os.path.exists(dir):\n",
    "       !git clone -b {branch} {url} {dir}\n",
    "\n",
    "def setup_directories(dirs):\n",
    "    for dir in dirs:\n",
    "        os.makedirs(dir, exist_ok=True)\n",
    "\n",
    "def install_dependencies():\n",
    "    t4_xformers_wheel = \"https://github.com/Linaqruf/colab-xformers/releases/download/0.0.20/xformers-0.0.20+1d635e1.d20230519-cp310-cp310-linux_x86_64.whl\"\n",
    "    gpu_info = getoutput('nvidia-smi')\n",
    "    !apt install aria2\n",
    "    !pip install -q --upgrade diffusers[torch]==0.18.2 transformers==4.30.2 einops==0.6.0 open-clip-torch==2.20.0 invisible-watermark -e .\n",
    "    !pip install -q xformers==0.0.20\n",
    "    !pip install \"jax[cuda12_pip]==0.4.23\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
    "    !pip install glob2\n",
    "    !pip install --force-reinstall -qqqq huggingface_hub\n",
    "    !pip install transformers\n",
    "    !pip install ipywidgets\n",
    "\n",
    "def prepare_environment():\n",
    "    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n",
    "    os.environ[\"SAFETENSORS_FAST_GPU\"] = \"1\"\n",
    "    os.environ[\"PYTHONWARNINGS\"] = \"ignore\"\n",
    "    os.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"garbage_collection_threshold:0.9,max_split_size_mb:512\"\n",
    "    os.environ[\"TCMALLOC_AGGRESSIVE_DECOMMIT\"] = \"t\"\n",
    "    os.environ[\"CUDA_MODULE_LOADING\"] = \"LAZY\"\n",
    "\n",
    "\n",
    "def main():\n",
    "    os.chdir(models_dir)  # Change working directory to models_dir\n",
    "    clone_repo(repo_url, repo_dir, branch)\n",
    "    setup_directories([repo_dir, tools_dir, vae_dir])\n",
    "    install_dependencies()\n",
    "    prepare_environment()\n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ♻ **Clean Folder**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "sRiX3jnynw3P"
   },
   "outputs": [],
   "source": [
    "from ipywidgets import Text, Button, Output\n",
    "from IPython.display import display, Markdown\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "# Function to clear and delete a folder\n",
    "def clear_and_delete_folder(colab_folder_path):\n",
    "    try:\n",
    "        # Use shutil.rmtree to remove all files and subdirectories\n",
    "        shutil.rmtree(colab_folder_path)\n",
    "        display(Markdown(f\"Deleted all contents in folder: `{colab_folder_path}`\"))\n",
    "    except Exception as e:\n",
    "        display(Markdown(f\"Error deleting folder `{colab_folder_path}`: {e}\"))\n",
    "\n",
    "# Create a text input widget for the folder path\n",
    "folder_path_input = Text(value=\"\", placeholder=\"Enter folder path\", description=\"Folder Path:\")\n",
    "\n",
    "# Create a button widget to trigger the deletion\n",
    "delete_button = Button(description=\"Delete Folder\")\n",
    "\n",
    "# Create an output widget to display the result\n",
    "output_widget = Output()\n",
    "\n",
    "# Define a function to handle the button click event\n",
    "def delete_folder(b):\n",
    "    colab_folder_path = folder_path_input.value\n",
    "    with output_widget:\n",
    "        clear_and_delete_folder(colab_folder_path)\n",
    "\n",
    "# Link the button click event to the delete_folder function\n",
    "delete_button.on_click(delete_folder)\n",
    "\n",
    "# Display the widgets\n",
    "display(folder_path_input)\n",
    "display(delete_button)\n",
    "display(output_widget)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ♻ **Fix Before Converting (Optional, may not work)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "NiWfPWulD6rW"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import requests\n",
    "import importlib.util\n",
    "\n",
    "# Note that this is meant to work IN TANDEM with the colab notebook, you'll need the installation of kohya-ss FIRST. in tandem meaning this is just the PY file that works seperate to the colab, and this works just as wel las the other.\n",
    "\n",
    "def check_and_download_script(script_name, script_url):\n",
    "    if not os.path.exists(script_name):\n",
    "        response = requests.get(script_url)\n",
    "        with open(script_name, 'wb') as f:\n",
    "            f.write(response.content)\n",
    "\n",
    "def fix_sdxl_model_keys(model_data):\n",
    "    text_encoder1, text_encoder2, vae, unet = model_data\n",
    "\n",
    "    # Fix keys in text_encoder1\n",
    "    fixed_text_encoder1 = {}\n",
    "    for k, v in text_encoder1.state_dict().items():\n",
    "        new_k = k.replace('first_stage_model.', '')\n",
    "        fixed_text_encoder1[new_k] = v\n",
    "\n",
    "    # Fix keys in text_encoder2\n",
    "    fixed_text_encoder2 = {}\n",
    "    for k, v in text_encoder2.state_dict().items():\n",
    "        new_k = k.replace('first_stage_model.', '')\n",
    "        fixed_text_encoder2[new_k] = v\n",
    "\n",
    "    # Fix keys in vae\n",
    "    fixed_vae = {}\n",
    "    for k, v in vae.state_dict().items():\n",
    "        new_k = k.replace('first_stage_model.', '')\n",
    "        fixed_vae[new_k] = v\n",
    "\n",
    "    # Fix keys in unet\n",
    "    fixed_unet = {}\n",
    "    for k, v in unet.state_dict().items():\n",
    "        new_k = k.replace('first_stage_model.', '')\n",
    "        fixed_unet[new_k] = v\n",
    "\n",
    "    return fixed_text_encoder1, fixed_text_encoder2, fixed_vae, fixed_unet\n",
    "\n",
    "def main():\n",
    "    script_name = \"convert_sdxl_to_diffusers.py\"\n",
    "    script_url = \"https://raw.githubusercontent.com/duskfallcrew/sdxl-model-converter/main/convert_sdxl_to_diffusers.py\"\n",
    "    check_and_download_script(script_name, script_url)\n",
    "\n",
    "    # Load the downloaded script\n",
    "    spec = importlib.util.spec_from_file_location(\"convert_sdxl_to_diffusers\", script_name)\n",
    "    convert_sdxl_to_diffusers = importlib.util.module_from_spec(spec)\n",
    "    spec.loader.exec_module(convert_sdxl_to_diffusers)\n",
    "\n",
    "    # Modify the convert_model function to use fix_sdxl_model_keys\n",
    "    def convert_model(args):\n",
    "        load_dtype = torch.float16 if args.fp16 else None\n",
    "        save_dtype = get_save_dtype(args)\n",
    "\n",
    "        is_load_checkpoint = determine_load_checkpoint(args.model_to_load)\n",
    "        is_save_checkpoint = not is_load_checkpoint  # reverse of load model\n",
    "\n",
    "        loaded_model_data = load_sdxl_model(args, is_load_checkpoint, load_dtype)\n",
    "        fixed_model_data = fix_sdxl_model_keys(loaded_model_data)\n",
    "        convert_and_save_sdxl_model(args, is_save_checkpoint, fixed_model_data, save_dtype)\n",
    "\n",
    "    # Call the modified convert_model function\n",
    "    convert_sdxl_to_diffusers.convert_model(args)  # Replace args with the required arguments\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  ♻  **Download SDXL**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "_nmd5ciZKv3Z"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import re\n",
    "import json\n",
    "import glob\n",
    "import gdown\n",
    "import requests\n",
    "import subprocess\n",
    "from urllib.parse import urlparse, unquote\n",
    "from pathlib import Path\n",
    "from ipywidgets import Text, Button, Output\n",
    "\n",
    "os.chdir(root_dir)\n",
    "\n",
    "# Create text widgets for input\n",
    "huggingface_token_widget = Text(value=\"\", placeholder=\"Huggingface Read Token\", description=\"Huggingface Token:\")\n",
    "sdxl_model_url_widget = Text(value=\"\", placeholder=\"SDXL Model URL\", description=\"SDXL Model URL:\")\n",
    "\n",
    "# Create a button to trigger the download\n",
    "download_button = Button(description=\"Download SDXL Model\")\n",
    "\n",
    "# Create an output widget to display the results\n",
    "output_widget = Output()\n",
    "\n",
    "def get_supported_extensions():\n",
    "    return tuple([\".ckpt\", \".safetensors\", \".pt\", \".pth\"])\n",
    "\n",
    "\n",
    "def get_filename(url):\n",
    "    extensions = get_supported_extensions()\n",
    "\n",
    "    if 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.endswith(tuple(get_supported_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",
    "    user_header = f\"Authorization: Bearer {HUGGINGFACE_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, filepath):\n",
    "    if \"/uc?id/\" in url or \"/file/d/\" in url:\n",
    "        return gdown.download(url, filepath, quiet=False)\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(url, dst):\n",
    "    filename = get_filename(url)\n",
    "    filepath = os.path.join(dst, filename)\n",
    "\n",
    "    if \"drive.google.com\" in url:\n",
    "        gdown = gdown_download(url, dst, filepath)\n",
    "    elif url.startswith(\"/content/drive/MyDrive/\"):\n",
    "        return url\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",
    "    huggingface_token = huggingface_token_widget.value\n",
    "    sdxl_model_url = sdxl_model_url_widget.value\n",
    "    os.chdir(root_dir)\n",
    "    model_path = sdxl_model_url\n",
    "    download(model_path, models_dir)\n",
    "    print(f\"Model downloaded at: {model_path}\")\n",
    "\n",
    "#def main(b):\n",
    "   # huggingface_token = huggingface_token_widget.value\n",
    "    #sdxl_model_url = sdxl_model_url_widget.value\n",
    "    #os.chdir(root_dir)  # <--- Add this line\n",
    "    #model_path = sdxl_model_url\n",
    "   # download(model_path, models_dir)\n",
    "   # with output_widget:\n",
    "       # print(f\"Model downloaded at: {model_path}\")\n",
    "#Uncomment if the other one without #B doesn't work\n",
    "# main(b) \n",
    "# same thing with the above one, if \"B\" is require it's becaues of the buttons. \n",
    "\n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  ♻ **Convert SDXL to Diffusers**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "uSEEJiDaK3Qw"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import urllib.request\n",
    "from ipywidgets import Text, Button, Output\n",
    "\n",
    "os.chdir(root_dir)\n",
    "\n",
    "# Create text widgets for input\n",
    "model_to_load_widget = Text(value=\"\", placeholder=\"Model to load\", description=\"Model to load:\")\n",
    "save_precision_as_widget = Text(value=\"fp16\", placeholder=\"Save precision as\", description=\"Save precision as:\")\n",
    "reference_model_widget = Text(value=\"stabilityai/stable-diffusion-xl-base-1.0\", placeholder=\"Reference model\", description=\"Reference model:\")\n",
    "\n",
    "# Create a button to trigger the conversion\n",
    "convert_button = Button(description=\"Convert SDXL to Diffusers\")\n",
    "\n",
    "# Create an output widget to display the results\n",
    "output_widget = Output()\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",
    "    return args\n",
    "\n",
    "def check_and_download_script(script_name, script_url):\n",
    "    if not os.path.exists(script_name):\n",
    "        print(f\"{script_name} not found, downloading...\")\n",
    "        urllib.request.urlretrieve(script_url, script_name)\n",
    "\n",
    "def run_script(script_name, script_args):\n",
    "   !python {script_name} {script_args}\n",
    "\n",
    "def main(b):\n",
    "    model_to_load = model_to_load_widget.value\n",
    "    save_precision_as = save_precision_as_widget.value\n",
    "    reference_model = reference_model_widget.value\n",
    "\n",
    "    script_name = \"convert_sdxl_to_diffusers.py\"\n",
    "    script_url = \"https://raw.githubusercontent.com/duskfallcrew/sdxl-model-converter/main/convert_sdxl_to_diffusers.py\"\n",
    "    check_and_download_script(script_name, script_url)\n",
    "\n",
    "    config = {\n",
    "        \"model_to_load\": model_to_load,\n",
    "        \"save_precision_as\": save_precision_as,\n",
    "        \"epoch\": 0,\n",
    "        \"global_step\": 0,\n",
    "        \"reference_model\": reference_model,\n",
    "    }\n",
    "    run_script(script_name, convert_dict(config))\n",
    "\n",
    "    with output_widget:\n",
    "        print(\"Conversion complete!\")\n",
    "\n",
    "# Link the button to the main function\n",
    "convert_button.on_click(main)\n",
    "\n",
    "# Display the widgets\n",
    "display(model_to_load_widget)\n",
    "display(save_precision_as_widget)\n",
    "display(reference_model_widget)\n",
    "display(convert_button)\n",
    "display(output_widget)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  ♻  Diffusers Deploy & Upload "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "KsVa2oLvRwL0"
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "from huggingface_hub import notebook_login, HfApi\n",
    "from ipywidgets import *\n",
    "from IPython.display import display\n",
    "from pathlib import Path\n",
    "\n",
    "# Authenticate with Hugging Face Hub\n",
    "notebook_login()\n",
    "\n",
    "# Define widgets for user inputs\n",
    "drive_path = Text(value=os.getcwd(), placeholder='Path to upload from', description='Path:')\n",
    "hfuser = Text(placeholder='YourUSERHERE', description='HF User:')\n",
    "hfrepo = Text(placeholder='ModelNameHere', description='HF Repo:')\n",
    "orgs_name = Text(placeholder='Organization name (optional)', description='Organization name:')\n",
    "make_private = Checkbox(value=False, description='Make repository private')\n",
    "\n",
    "# Initialize Hugging Face API\n",
    "api = HfApi()\n",
    "\n",
    "# Define function to update folder list\n",
    "def update_folders(_):\n",
    "    global folder_picker\n",
    "    path = Path(drive_path.value)\n",
    "    if path.exists() and path.is_dir():\n",
    "        all_folders = [str(f) for f in path.glob(\"*\") if f.is_dir()]\n",
    "        folder_picker.options = all_folders\n",
    "    else:\n",
    "        with out:\n",
    "            print(\"Invalid directory. Please check the path and try again.\")\n",
    "\n",
    "# Define upload function\n",
    "def upload_folders(_):\n",
    "    repo_id = f\"{hfuser.value or hfuser.placeholder}/{hfrepo.value or hfrepo.placeholder}\"\n",
    "    if orgs_name.value:\n",
    "        repo_id = f\"{orgs_name.value}/{hfrepo.value or hfrepo.placeholder}\"\n",
    "    try:\n",
    "        validate_repo_id(repo_id)\n",
    "        api.create_repo(repo_id=repo_id, repo_type=\"model\", private=make_private.value)\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",
    "    with out:\n",
    "        if folder_picker is None or len(folder_picker.value) < 1:\n",
    "            print(\"Nothing selected for upload, make sure to select one of the folders in the list, or verify there are folders in the specified directory.\")\n",
    "        for folder in folder_picker.value:\n",
    "            print(f\"Uploading to HF: huggingface.co/{repo_id}/{os.path.basename(folder)}\")\n",
    "            api.upload_folder(\n",
    "                folder_path=folder,\n",
    "                path_in_repo=os.path.basename(folder),\n",
    "                repo_id=repo_id,\n",
    "                commit_message=\"Upload with Earth & Dusk Huggingface 🤗 Backup\"\n",
    "            )\n",
    "        print(\"DONE\")\n",
    "        print(\"Go to your repo and accept the PRs created to see your files.\")\n",
    "\n",
    "# Bind the functions to buttons\n",
    "update_btn = Button(description='Update Folders')\n",
    "upload_btn = Button(description='Upload')\n",
    "out = Output()\n",
    "\n",
    "update_btn.on_click(update_folders)\n",
    "upload_btn.on_click(upload_folders)\n",
    "\n",
    "# Create and display the widget layout\n",
    "folder_list = [f for f in os.listdir(drive_path.value) if os.path.isdir(os.path.join(drive_path.value, f))]\n",
    "folder_picker = SelectMultiple(options=folder_list, layout=Layout(width=\"600px\"))\n",
    "box = VBox([\n",
    "    drive_path,\n",
    "    HBox([hfuser, hfrepo]),\n",
    "    HBox([orgs_name, make_private]),\n",
    "    update_btn,\n",
    "    folder_picker,\n",
    "    upload_btn,\n",
    "    out\n",
    "])\n",
    "\n",
    "display(box)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  ♻ Extra: Vae PT to Safetensors (Should work but hey I dunno!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import io\n",
    "import os\n",
    "import requests\n",
    "import torch\n",
    "import yaml\n",
    "from diffusers import AutoencoderKL\n",
    "from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (\n",
    "    assign_to_checkpoint,\n",
    "    conv_attn_to_linear,\n",
    "    create_vae_diffusers_config,\n",
    "    renew_vae_attention_paths,\n",
    "    renew_vae_resnet_paths,\n",
    ")\n",
    "import safetensors\n",
    "import json\n",
    "\n",
    "def custom_convert_ldm_vae_checkpoint(checkpoint, config):\n",
    "    vae_state_dict = checkpoint\n",
    "\n",
    "    new_checkpoint = {}\n",
    "\n",
    "    new_checkpoint[\"encoder.conv_in.weight\"] = vae_state_dict[\"encoder.conv_in.weight\"]\n",
    "    new_checkpoint[\"encoder.conv_in.bias\"] = vae_state_dict[\"encoder.conv_in.bias\"]\n",
    "    new_checkpoint[\"encoder.conv_out.weight\"] = vae_state_dict[\"encoder.conv_out.weight\"]\n",
    "    new_checkpoint[\"encoder.conv_out.bias\"] = vae_state_dict[\"encoder.conv_out.bias\"]\n",
    "    new_checkpoint[\"encoder.conv_norm_out.weight\"] = vae_state_dict[\"encoder.norm_out.weight\"]\n",
    "    new_checkpoint[\"encoder.conv_norm_out.bias\"] = vae_state_dict[\"encoder.norm_out.bias\"]\n",
    "\n",
    "    new_checkpoint[\"decoder.conv_in.weight\"] = vae_state_dict[\"decoder.conv_in.weight\"]\n",
    "    new_checkpoint[\"decoder.conv_in.bias\"] = vae_state_dict[\"decoder.conv_in.bias\"]\n",
    "    new_checkpoint[\"decoder.conv_out.weight\"] = vae_state_dict[\"decoder.conv_out.weight\"]\n",
    "    new_checkpoint[\"decoder.conv_out.bias\"] = vae_state_dict[\"decoder.conv_out.bias\"]\n",
    "    new_checkpoint[\"decoder.conv_norm_out.weight\"] = vae_state_dict[\"decoder.norm_out.weight\"]\n",
    "    new_checkpoint[\"decoder.conv_norm_out.bias\"] = vae_state_dict[\"decoder.norm_out.bias\"]\n",
    "\n",
    "    new_checkpoint[\"quant_conv.weight\"] = vae_state_dict[\"quant_conv.weight\"]\n",
    "    new_checkpoint[\"quant_conv.bias\"] = vae_state_dict[\"quant_conv.bias\"]\n",
    "    new_checkpoint[\"post_quant_conv.weight\"] = vae_state_dict[\"post_quant_conv.weight\"]\n",
    "    new_checkpoint[\"post_quant_conv.bias\"] = vae_state_dict[\"post_quant_conv.bias\"]\n",
    "\n",
    "    # Retrieves the keys for the encoder down blocks only\n",
    "    num_down_blocks = len({\".\".join(layer.split(\".\")[:3]) for layer in vae_state_dict if \"encoder.down\" in layer})\n",
    "    down_blocks = {\n",
    "        layer_id: [key for key in vae_state_dict if f\"down.{layer_id}\" in key] for layer_id in range(num_down_blocks)\n",
    "    }\n",
    "\n",
    "    # Retrieves the keys for the decoder up blocks only\n",
    "    num_up_blocks = len({\".\".join(layer.split(\".\")[:3]) for layer in vae_state_dict if \"decoder.up\" in layer})\n",
    "    up_blocks = {\n",
    "        layer_id: [key for key in vae_state_dict if f\"up.{layer_id}\" in key] for layer_id in range(num_up_blocks)\n",
    "    }\n",
    "\n",
    "    for i in range(num_down_blocks):\n",
    "        resnets = [key for key in down_blocks[i] if f\"down.{i}\" in key and f\"down.{i}.downsample\" not in key]\n",
    "\n",
    "        if f\"encoder.down.{i}.downsample.conv.weight\" in vae_state_dict:\n",
    "            new_checkpoint[f\"encoder.down_blocks.{i}.downsamplers.0.conv.weight\"] = vae_state_dict.pop(\n",
    "                f\"encoder.down.{i}.downsample.conv.weight\"\n",
    "            )\n",
    "            new_checkpoint[f\"encoder.down_blocks.{i}.downsamplers.0.conv.bias\"] = vae_state_dict.pop(\n",
    "                f\"encoder.down.{i}.downsample.conv.bias\"\n",
    "            )\n",
    "\n",
    "        paths = renew_vae_resnet_paths(resnets)\n",
    "        meta_path = {\"old\": f\"down.{i}.block\", \"new\": f\"down_blocks.{i}.resnets\"}\n",
    "        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n",
    "\n",
    "    mid_resnets = [key for key in vae_state_dict if \"encoder.mid.block\" in key]\n",
    "    num_mid_res_blocks = 2\n",
    "    for i in range(1, num_mid_res_blocks + 1):\n",
    "        resnets = [key for key in mid_resnets if f\"encoder.mid.block_{i}\" in key]\n",
    "\n",
    "        paths = renew_vae_resnet_paths(resnets)\n",
    "        meta_path = {\"old\": f\"mid.block_{i}\", \"new\": f\"mid_block.resnets.{i - 1}\"}\n",
    "        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n",
    "\n",
    "    mid_attentions = [key for key in vae_state_dict if \"encoder.mid.attn\" in key]\n",
    "    paths = renew_vae_attention_paths(mid_attentions)\n",
    "    meta_path = {\"old\": \"mid.attn_1\", \"new\": \"mid_block.attentions.0\"}\n",
    "    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n",
    "    conv_attn_to_linear(new_checkpoint)\n",
    "\n",
    "    for i in range(num_up_blocks):\n",
    "        block_id = num_up_blocks - 1 - i\n",
    "        resnets = [\n",
    "            key for key in up_blocks[block_id] if f\"up.{block_id}\" in key and f\"up.{block_id}.upsample\" not in key\n",
    "        ]\n",
    "\n",
    "        if f\"decoder.up.{block_id}.upsample.conv.weight\" in vae_state_dict:\n",
    "            new_checkpoint[f\"decoder.up_blocks.{i}.upsamplers.0.conv.weight\"] = vae_state_dict[\n",
    "                f\"decoder.up.{block_id}.upsample.conv.weight\"\n",
    "            ]\n",
    "            new_checkpoint[f\"decoder.up_blocks.{i}.upsamplers.0.conv.bias\"] = vae_state_dict[\n",
    "                f\"decoder.up.{block_id}.upsample.conv.bias\"\n",
    "            ]\n",
    "\n",
    "        paths = renew_vae_resnet_paths(resnets)\n",
    "        meta_path = {\"old\": f\"up.{block_id}.block\", \"new\": f\"up_blocks.{i}.resnets\"}\n",
    "        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n",
    "\n",
    "    mid_resnets = [key for key in vae_state_dict if \"decoder.mid.block\" in key]\n",
    "    num_mid_res_blocks = 2\n",
    "    for i in range(1, num_mid_res_blocks + 1):\n",
    "        resnets = [key for key in mid_resnets if f\"decoder.mid.block_{i}\" in key]\n",
    "\n",
    "        paths = renew_vae_resnet_paths(resnets)\n",
    "        meta_path = {\"old\": f\"mid.block_{i}\", \"new\": f\"mid_block.resnets.{i - 1}\"}\n",
    "        assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n",
    "\n",
    "    mid_attentions = [key for key in vae_state_dict if \"decoder.mid.attn\" in key]\n",
    "    paths = renew_vae_attention_paths(mid_attentions)\n",
    "    meta_path = {\"old\": \"mid.attn_1\", \"new\": \"mid_block.attentions.0\"}\n",
    "    assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)\n",
    "    conv_attn_to_linear(new_checkpoint)\n",
    "    return new_checkpoint\n",
    "\n",
    "\n",
    "def vae_pt_to_vae_diffuser(\n",
    "    checkpoint_path: str,\n",
    "    output_path: str,\n",
    "    config_json: str,\n",
    "):\n",
    "    # Load the config JSON\n",
    "    with open(config_json, \"r\") as f:\n",
    "        vae_config = json.load(f)\n",
    "\n",
    "    # Load the checkpoint\n",
    "    if checkpoint_path.endswith(\"safetensors\"):\n",
    "        from safetensors import safe_open\n",
    "\n",
    "        checkpoint = {}\n",
    "        with safe_open(checkpoint_path, framework=\"pt\", device=\"cpu\") as f:\n",
    "            for key in f.keys():\n",
    "                checkpoint[key] = f.get_tensor(key)\n",
    "    else:\n",
    "        checkpoint = torch.load(checkpoint_path, map_location=\"cpu\")[\"state_dict\"]\n",
    "\n",
    "    # Convert the VAE model\n",
    "    converted_vae_checkpoint = custom_convert_ldm_vae_checkpoint(checkpoint, vae_config)\n",
    "\n",
    "    vae = AutoencoderKL(**vae_config)\n",
    "    vae.load_state_dict(converted_vae_checkpoint)\n",
    "\n",
    "    # Save the converted model to a SafeTensors file\n",
    "    with safetensors.safe_open(output_path, \"wb\", framework=\"pt\") as f:\n",
    "        for key, tensor in converted_vae_checkpoint.items():\n",
    "            f.set_tensor(key, tensor.cpu().numpy())\n",
    "\n",
    "    print(f\"Model converted to SafeTensors format and saved to {output_path}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "\n",
    "    parser.add_argument(\"--vae_pt_path\", default=None, type=str, required=True, help=\"Path to the VAE.pt to convert.\")\n",
    "    parser.add_argument(\"--dump_path\", default=None, type=str, required=True, help=\"Path to the VAE.pt to convert.\")\n",
    "    parser.add_argument(\"--config_json\", default=None, type=str, required=True, help=\"Path to the config JSON file.\")\n",
    "\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path, args.config_json)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "gpuType": "V28",
   "private_outputs": true,
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}