sdxl_model_converter_duskfallcrew_colab.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "V28",
      "private_outputs": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "# **SDXL Model Converter**\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",
        "Still to come -\n",
        "\n",
        "Chat GPT up instructions..\n",
        "add more links.\n",
        "\n",
        "Fixed Jax Error, Using TPUV2 Runtime.\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",
        "\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",
        "\n",
        "Cancelled Ideas that don't make sense for this product\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.\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"
      ],
      "metadata": {
        "id": "dY-0HMB7VyNu"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "R3hfwYaXKn_V"
      },
      "outputs": [],
      "source": [
        "#@title ## **Install Kohya Script**\n",
        "import os\n",
        "from subprocess import getoutput\n",
        "\n",
        "# Directories\n",
        "root_dir = \"/content\"\n",
        "repo_dir = os.path.join(root_dir, \"kohya-trainer\")\n",
        "models_dir = os.path.join(root_dir, \"models\")\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",
        "\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",
        "def main():\n",
        "    os.chdir(root_dir)\n",
        "    clone_repo(repo_url, repo_dir, branch)\n",
        "    os.chdir(repo_dir)\n",
        "    setup_directories([repo_dir, models_dir, tools_dir, vae_dir])\n",
        "    install_dependencies()\n",
        "    prepare_environment()\n",
        "\n",
        "main()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ### ♻ **Clean Folder**\n",
        "\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",
        "# @markdown ### Folder Path for Deletion\n",
        "\n",
        "colab_folder_path = \"\" # @param {type: \"string\"}\n",
        "\n",
        "# Call the function to clear and delete the folder\n",
        "clear_and_delete_folder(colab_folder_path)\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "sRiX3jnynw3P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## ♻ **Fix Before Converting (Optional, may not work)**\n",
        "#@markdown This MAY not work for SDXL, if you require key fixings and this works let me know!\n",
        "from safetensors import safe_open\n",
        "from safetensors.torch import save_file\n",
        "\n",
        "def fix_diffusers_model_conversion(load_path: str = None, save_path: str = None):\n",
        "    # Use default paths if not provided\n",
        "    load_path = load_path or \"input\"  # @param {type: \"string\"}\n",
        "    save_path = save_path or \"output\"  # @param {type: \"string\"}\n",
        "\n",
        "    # load original\n",
        "    tensors = {}\n",
        "    with safe_open(load_path, framework=\"pt\") as f:\n",
        "        for key in f.keys():\n",
        "            tensors[key] = f.get_tensor(key)\n",
        "\n",
        "    # migrate\n",
        "    new_tensors = {}\n",
        "    for k, v in tensors.items():\n",
        "        new_key = k\n",
        "        # only fix the vae\n",
        "        if 'first_stage_model.' in k:\n",
        "            # migrate q, k, v keys\n",
        "            new_key = new_key.replace('.to_q.weight', '.q.weight')\n",
        "            new_key = new_key.replace('.to_q.bias', '.q.bias')\n",
        "            new_key = new_key.replace('.to_k.weight', '.k.weight')\n",
        "            new_key = new_key.replace('.to_k.bias', '.k.bias')\n",
        "            new_key = new_key.replace('.to_v.weight', '.v.weight')\n",
        "            new_key = new_key.replace('.to_v.bias', '.v.bias')\n",
        "        new_tensors[new_key] = v\n",
        "\n",
        "    # save\n",
        "    save_file(new_tensors, save_path)\n",
        "\n",
        "# Example usage\n",
        "fix_diffusers_model_conversion(load_path='/content/drive/path/to/your/original_file.pth',\n",
        "                               save_path='/content/drive/path/to/save/fixed_file.pth')\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "NiWfPWulD6rW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## **Download SDXL**\n",
        "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",
        "\n",
        "os.chdir(root_dir)\n",
        "\n",
        "# @markdown Place your Huggingface [Read Token](https://huggingface.co/settings/tokens) Here.\n",
        "\n",
        "HUGGINGFACE_TOKEN = \"\"#@param {type: \"string\"}\n",
        "SDXL_MODEL_URL = \"\" #@param {type: \"string\"}\n",
        "\n",
        "def get_supported_extensions():\n",
        "    return tuple([\".ckpt\", \".safetensors\", \".pt\", \".pth\"])\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",
        "    model_path = SDXL_MODEL_URL\n",
        "    download(model_path, models_dir)\n",
        "    print(f\"Model downloaded at: {model_path}\")\n",
        "\n",
        "main()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "_nmd5ciZKv3Z"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ## ♻ **Fix Before Converting**\n",
        "from safetensors import safe_open\n",
        "from safetensors.torch import save_file\n",
        "\n",
        "def fix_diffusers_model_conversion(load_path: str = None, save_path: str = None):\n",
        "    # Use default paths if not provided\n",
        "    load_path = load_path or \"input\"  # @param {type: \"string\"}\n",
        "    save_path = save_path or \"output\"  # @param {type: \"string\"}\n",
        "\n",
        "    # load original\n",
        "    tensors = {}\n",
        "    with safe_open(load_path, framework=\"pt\") as f:\n",
        "        for key in f.keys():\n",
        "            tensors[key] = f.get_tensor(key)\n",
        "\n",
        "    # migrate\n",
        "    new_tensors = {}\n",
        "    for k, v in tensors.items():\n",
        "        new_key = k\n",
        "        # only fix the vae\n",
        "        if 'first_stage_model.' in k:\n",
        "            # migrate q, k, v keys\n",
        "            new_key = new_key.replace('.to_q.weight', '.q.weight')\n",
        "            new_key = new_key.replace('.to_q.bias', '.q.bias')\n",
        "            new_key = new_key.replace('.to_k.weight', '.k.weight')\n",
        "            new_key = new_key.replace('.to_k.bias', '.k.bias')\n",
        "            new_key = new_key.replace('.to_v.weight', '.v.weight')\n",
        "            new_key = new_key.replace('.to_v.bias', '.v.bias')\n",
        "        new_tensors[new_key] = v\n",
        "\n",
        "    # save\n",
        "    save_file(new_tensors, save_path)\n",
        "\n",
        "# Example usage\n",
        "fix_diffusers_model_conversion(load_path='/content/drive/path/to/your/original_file.pth',\n",
        "                               save_path='/content/drive/path/to/save/fixed_file.pth')\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "1TbD-nJX5TCQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title ## **Convert SDXL to Diffusers**\n",
        "import os\n",
        "import urllib.request\n",
        "\n",
        "os.chdir(root_dir)\n",
        "#@markdown ### **Conversion Config**\n",
        "model_to_load = \"\" #@param {'type': 'string'}\n",
        "global_step = 0\n",
        "epoch = 0\n",
        "save_precision_as = \"fp16\" #@param [\"fp16\",\"bf16\",\"float\"] {'allow-input': false}\n",
        "reference_model = \"stabilityai/stable-diffusion-xl-base-1.0\"\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():\n",
        "    script_name = \"convert_sdxl_to_diffusers.py\"\n",
        "    script_url = \"https://raw.githubusercontent.com/Linaqruf/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\": epoch,\n",
        "        \"global_step\": global_step,\n",
        "        \"reference_model\": reference_model,\n",
        "    }\n",
        "    run_script(script_name, convert_dict(config))\n",
        "\n",
        "main()\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "uSEEJiDaK3Qw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Deploy**"
      ],
      "metadata": {
        "id": "sjtc_ZNaTpUl"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ### **Huggingface Hub config**\n",
        "from huggingface_hub import login, 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 repo does not exist, it will automatically create it.\n",
        "model_name = \"\"  # @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_model_repo(api, user, orgs_name, model_name, make_private=False):\n",
        "    if orgs_name == \"\":\n",
        "        repo_id = user[\"name\"] + \"/\" + model_name.strip()\n",
        "    else:\n",
        "        repo_id = orgs_name + \"/\" + model_name.strip()\n",
        "\n",
        "    try:\n",
        "        validate_repo_id(repo_id)\n",
        "        api.create_repo(repo_id=repo_id, repo_type=\"model\", 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",
        "\n",
        "    return repo_id\n",
        "\n",
        "user, api = authenticate(write_token)\n",
        "\n",
        "if model_name:\n",
        "    model_repo = create_model_repo(api, user, orgs_name, model_name, make_private)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "pU3wMdOeXE4G"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title ### ♻ **Upload to Huggingface**\n",
        "from huggingface_hub import HfApi\n",
        "from pathlib import Path\n",
        "import os\n",
        "\n",
        "api = HfApi()\n",
        "\n",
        "# @markdown This will be uploaded to model repo\n",
        "model_path = \"\"  # @param {type :\"string\"}\n",
        "path_in_repo = \"\"  # @param {type :\"string\"}\n",
        "project_name = os.path.basename(model_path)\n",
        "\n",
        "# @markdown Other Information\n",
        "commit_message = \"Upload with 🚀🤗 E&D SDXL Diffusers notebook\"  # @param {type :\"string\"}\n",
        "\n",
        "def is_diffusers_model(model_path):\n",
        "    required_folders = {\"unet\", \"text_encoder\", \"text_encoder_2\", \"tokenizer\", \"tokenizer_2\", \"scheduler\", \"vae\"}\n",
        "    return required_folders.issubset(set(os.listdir(model_path))) and os.path.isfile(os.path.join(model_path, \"model_index.json\"))\n",
        "\n",
        "def upload_model(model_paths, is_folder: bool):\n",
        "    path_obj = Path(model_paths)\n",
        "    trained_model = path_obj.parts[-1]\n",
        "\n",
        "    path_in_repo_local = path_in_repo if path_in_repo and not is_diffusers_model(model_paths) else \"\"\n",
        "\n",
        "    notification = f\"Uploading {trained_model} from {model_paths} to https://huggingface.co/{model_repo}\"\n",
        "    print(notification)\n",
        "\n",
        "    if is_folder:\n",
        "        if is_diffusers_model(model_paths):\n",
        "            commit_message = f\"Upload diffusers format: {trained_model}\"\n",
        "            print(\"Detected diffusers model. Adjusting upload parameters.\")\n",
        "        else:\n",
        "            commit_message = f\"Upload checkpoint: {trained_model}\"\n",
        "            print(\"Detected regular model. Adjusting upload parameters.\")\n",
        "\n",
        "        api.upload_folder(\n",
        "            folder_path=model_paths,\n",
        "            path_in_repo=path_in_repo_local,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "            ignore_patterns=\".ipynb_checkpoints\",\n",
        "        )\n",
        "    else:\n",
        "        commit_message = f\"Upload file: {trained_model}\"\n",
        "        api.upload_file(\n",
        "            path_or_fileobj=model_paths,\n",
        "            path_in_repo=path_in_repo_local,\n",
        "            repo_id=model_repo,\n",
        "            commit_message=commit_message,\n",
        "        )\n",
        "\n",
        "    success_notification = f\"♻ Model upload complete, care to try again? Thanks for flying Stable Diffusion Airlines, you owe me five dollars.. Check the model at https://huggingface.co/{model_repo}/tree/main\"\n",
        "    print(success_notification)\n",
        "\n",
        "def upload():\n",
        "    if model_path.endswith((\".ckpt\", \".safetensors\", \".pt\")):\n",
        "        upload_model(model_path, False)\n",
        "    else:\n",
        "        upload_model(model_path, True)\n",
        "\n",
        "upload()\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "YPzfmu_4DtWG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Optional. Folder/Singular File Uploads\n",
        "#@markdown Please ensure that you have noted where you are uploading your files in the \"REPO\" section. Follow these steps to create a new repository on your Hugging Face account:\n",
        "#@markdown [Navigate to Your Hugging Face Account](https://huggingface.co/)\n",
        "\n",
        "#@markdown This is not proven to work with diffusers just yet, if this doesn't work fall back on above.\n",
        "\n",
        "#@markdown Visit the Hugging Face website and log in to your account via the above link\n",
        "\n",
        "#@markdown Go to your user profile page.\n",
        "#@markdown [Click on the \"New  Model\" button.](https://huggingface.co/new)\n",
        "\n",
        "#@markdown Enter a name for your model repository.\n",
        "#@markdown Choose the appropriate settings and permissions.\n",
        "\n",
        "#@markdown Note the name of the newly created repository.\n",
        "\n",
        "#@markdown You'll need to ask ChatGPT for the updated widget, which you just hit the button and everything is mostly straight forward. You just add your drive/storage information, User information and follow the boxes.\n",
        "\n",
        "import glob\n",
        "import os\n",
        "from huggingface_hub import HfApi\n",
        "from ipywidgets import *\n",
        "from IPython.display import display\n",
        "from pathlib import Path\n",
        "\n",
        "# Define widgets for user inputs\n",
        "drive_path = Text(value='/content/drive/My Drive', placeholder='Path in Google Drive', description='Drive Path:')\n",
        "hfuser = Text(placeholder='YourUSERHERE', description='HF User:')\n",
        "hfrepo = Text(placeholder='ModelNameHere', description='HF Repo:')\n",
        "update_btn = Button(description='Update Files')\n",
        "upload_btn = Button(description='Upload')\n",
        "out = Output()\n",
        "\n",
        "# Initialize Hugging Face API\n",
        "api = HfApi()\n",
        "\n",
        "# Define function to update file list\n",
        "def update_files(_):\n",
        "    global ckpt_picker\n",
        "    path = Path(drive_path.value)\n",
        "    if path.exists() and path.is_dir():\n",
        "        all_ckpts = [str(f) for f in path.glob(\"*.safetensors\")]\n",
        "        ckpt_picker.options = all_ckpts\n",
        "    else:\n",
        "        with out:\n",
        "            print(\"Invalid directory. Please check the path and try again.\")\n",
        "\n",
        "# Define upload function\n",
        "def upload_ckpts(_):\n",
        "    repo_id = f\"{hfuser.value or hfuser.placeholder}/{hfrepo.value or hfrepo.placeholder}\"\n",
        "    with out:\n",
        "        if ckpt_picker is None or len(ckpt_picker.value) < 1:\n",
        "            print(\"Nothing selected for upload, make sure to select one of the checkpoint files in the list, or verify there are checkpoint files in the specified directory.\")\n",
        "        for ckpt in ckpt_picker.value:\n",
        "            print(f\"Uploading to HF: huggingface.co/{repo_id}/{os.path.basename(ckpt)}\")\n",
        "            response = api.upload_file(\n",
        "                path_or_fileobj=ckpt,\n",
        "                path_in_repo=os.path.basename(ckpt),\n",
        "                repo_id=repo_id,\n",
        "                repo_type=None,\n",
        "                create_pr=1,\n",
        "                commit_message=\"Upload with Earth & Dusk Huggingface 🤗 Backup\"\n",
        "            )\n",
        "            display(response)\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.on_click(update_files)\n",
        "upload_btn.on_click(upload_ckpts)\n",
        "\n",
        "# Create and display the widget layout\n",
        "ckpt_picker = SelectMultiple(options=[], layout=Layout(width=\"600px\"))\n",
        "box = VBox([\n",
        "    drive_path,\n",
        "    HBox([hfuser, hfrepo]),\n",
        "    update_btn,\n",
        "    ckpt_picker,\n",
        "    upload_btn,\n",
        "    out\n",
        "])\n",
        "\n",
        "display(box)\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "KsVa2oLvRwL0"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}