{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
}
},
"cells": [
{
"cell_type": "markdown",
"source": [
"### This is a script that converts the format of the model (.ckpt/.safetensors)\n",
"#### It also converts the .yaml file included with the SD2.x series\n",
"#### It can also be saved as fp16 as an option"
],
"metadata": {
"id": "fAIY_GORNEYa"
}
},
{
"cell_type": "markdown",
"source": [
"Run both of the following two codes"
],
"metadata": {
"id": "OnuCk_wNLM_D"
}
},
{
"cell_type": "code",
"source": [
"from google.colab import drive \n",
"drive.mount(\"/content/drive\")"
],
"metadata": {
"id": "liEiK8Iioscq"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"!pip install torch safetensors\n",
"!pip install pytorch-lightning\n",
"!pip install wget"
],
"metadata": {
"id": "pXr7oNJzwwgU"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"Replace the following links, etc. with the desired ones and then run the following code"
],
"metadata": {
"id": "7Ils-K70k15Y"
}
},
{
"cell_type": "code",
"source": [
"#@title Download Models\n",
"#@markdown Please specify the model name or download link for Google Drive, separated by commas\n",
"#@markdown - If it is the model name on Google Drive, specify it as a relative path to My Drive\n",
"#@markdown - If it is a download link, copy the link address by right-clicking and paste it in place of the link below\n",
"import shutil\n",
"import urllib.parse\n",
"import urllib.request\n",
"import wget\n",
"\n",
"models = \"Specify_the_model_in_this_way_if_the_model_is_on_My_Drive.safetensors, https://huggingface.co./hakurei/waifu-diffusion-v1-4/resolve/main/wd-1-4-anime_e1.ckpt, https://huggingface.co./hakurei/waifu-diffusion-v1-4/resolve/main/wd-1-4-anime_e1.yaml\" #@param {type:\"string\"}\n",
"models = [m.strip() for m in models.split(\",\")]\n",
"for model in models:\n",
" if 0 < len(urllib.parse.urlparse(model).scheme): # if model is url\n",
" wget.download(model)\n",
" elif model.endswith((\".ckpt\", \".safetensors\", \".yaml\")):\n",
" shutil.copy(\"/content/drive/MyDrive/\" + model, \"/content/\" + model) # get the model from mydrive\n",
" else:\n",
" print(f\"\\\"{model}\\\" is not a URL and is also not a file with a proper extension\")"
],
"metadata": {
"cellView": "form",
"id": "4vd3A09AxJE0"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"#@title Convert the Models\n",
"#@markdown Specify the models to be converted, separated by commas
\n",
"#@markdown If nothing is inputted, all loaded models will be converted\n",
"import os\n",
"import glob\n",
"import torch\n",
"import safetensors.torch\n",
"from functools import partial\n",
"\n",
"models = \"wd-1-4-anime_e1.ckpt, wd-1-4-anime_e1.yaml\" #@param {type:\"string\"}\n",
"as_fp16 = True #@param {type:\"boolean\"}\n",
"save_directly_to_Google_Drive = True #@param {type:\"boolean\"}\n",
"save_type = \".safetensors\" #@param [\".safetensors\", \".ckpt\"]\n",
"\n",
"def convert_yaml(file_name):\n",
" with open(file_name) as f:\n",
" yaml = f.read()\n",
" if save_directly_to_Google_Drive:\n",
" os.chdir(\"/content/drive/MyDrive\")\n",
" is_safe = save_type == \".safetensors\"\n",
" yaml = yaml.replace(f\"use_checkpoint: {is_safe}\", f\"use_checkpoint: {not is_safe}\")\n",
" if as_fp16:\n",
" yaml = yaml.replace(\"use_fp16: False\", \"use_fp16: True\")\n",
" file_name = os.path.splitext(file_name)[0] + \"-fp16.yaml\"\n",
" with open(file_name, mode=\"w\") as f:\n",
" f.write(yaml)\n",
" os.chdir(\"/content\")\n",
"\n",
"if models == \"\":\n",
" models = [os.path.basename(m) for m in glob.glob(r\"/content/*.ckpt\") + glob.glob(r\"/content/*.safetensors\") + glob.glob(r\"/content/*.yaml\")]\n",
"else:\n",
" models = [m.strip() for m in models.split(\",\")]\n",
"\n",
"for model in models:\n",
" model_name, model_ext = os.path.splitext(model)\n",
" if model_ext == \".yaml\":\n",
" convert_yaml(model)\n",
" elif (model_ext != \".safetensors\") & (model_ext != \".ckpt\"):\n",
" print(\"The supported formats are only .ckpt, .safetensors, and .yaml\\n\" + f\"\\\"{model}\\\" is not a supported format\")\n",
" else:\n",
" load_model = partial(safetensors.torch.load_file, device=\"cpu\") if model_ext == \".safetensors\" else partial(torch.load, map_location=torch.device(\"cpu\"))\n",
" save_model = safetensors.torch.save_file if save_type == \".safetensors\" else torch.save\n",
" # convert model\n",
" with torch.no_grad():\n",
" weights = load_model(model)\n",
" if \"state_dict\" in weights:\n",
" weights = weights[\"state_dict\"]\n",
" if as_fp16:\n",
" model_name = model_name + \"-fp16\"\n",
" for key in weights.keys():\n",
" weights[key] = weights[key].half()\n",
" if save_directly_to_Google_Drive:\n",
" os.chdir(\"/content/drive/MyDrive\")\n",
" save_model(weights, model_name + save_type)\n",
" os.chdir(\"/content\")\n",
" del weights\n",
"\n",
"!reset"
],
"metadata": {
"id": "9OmSG98HxJg2",
"cellView": "form"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"If you are converting SD2.x series models, etc., be sure to download/convert the accompanying configuration file (a .yaml file with the same name as the model) at the same time.\n",
"\n",
"It can be converted in the same way as the model."
],
"metadata": {
"id": "SWTFKmGFLec6"
}
},
{
"cell_type": "markdown",
"source": [
"If you run out of memory and crash, you can use a smaller model or a paid high memory runtime.\n",
"\n",
"With the free ~12GB runtime, you can convert models up to ~10GB."
],
"metadata": {
"id": "0SUK6Alv2ItS"
}
},
{
"cell_type": "markdown",
"source": [
"Choose your favorite model from https://huggingface.co./models?other=stable-diffusion or other model link collections"
],
"metadata": {
"id": "yaLq5Nqe6an6"
}
}
]
}