Update as_safetensors+fp16_en.ipynb

as_safetensors+fp16_en.ipynb

@@ -14,6 +14,17 @@
  }
  },
  "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": [
@@ -39,6 +50,7 @@
  "cell_type": "code",
  "source": [
  "!pip install torch safetensors\n",
+ "!pip install pytorch-lightning\n",
  "!pip install wget"
  ],
  "metadata": {
@@ -63,27 +75,20 @@
  "#@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",
- "\n",
  "import shutil\n",
  "import urllib.parse\n",
  "import urllib.request\n",
  "import wget\n",
  "\n",
- "models = \"Please use your own model in place of this example, example.safetensors, https://huggingface.co/stabilityai/stable-diffusion-2-1/resolve/main/v2-1_768-ema-pruned.ckpt\" #@param {type:\"string\"}\n",
- "models = [m.strip() for m in models.split(\",\") if not models == \"\"]\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",
- " # once the bug on python 3.8 is fixed, replace the above code with the following code\n",
- " ## model_data = urllib.request.urlopen(model).read()\n",
- " ## with open(os.path.basename(model), mode=\"wb\") as f:\n",
- " ## f.write(model_data)\n",
- " elif model.endswith((\".ckpt\", \".safetensors\", \".pt\", \".pth\")):\n",
- " from_ = \"/content/drive/MyDrive/\" + model\n",
- " to_ = \"/content/\" + model\n",
- " shutil.copy(from_, to_)\n",
+ " elif model.endswith((\".ckpt\", \".safetensors\")):\n",
+ " shutil.copy(\"/content/drive/MyDrive/\" + model, \"/content/\" + model) # get the model from mydrive\n",
  " else:\n",
- " print(f\"\\\"{model}\\\"はURLではなく、正しい形式のファイルでもありません\")"
+ " print(f\"\\\"{model}\\\" is not a URL and is also not a file with a proper extension\")"
  ],
  "metadata": {
  "cellView": "form",
@@ -92,66 +97,65 @@
  "execution_count": null,
  "outputs": []
  },
- {
- "cell_type": "markdown",
- "source": [
- "if you use a relatively newer model such as SD2.1, run the following code"
- ],
- "metadata": {
- "id": "m1mHzOMjcDhz"
- }
- },
- {
- "cell_type": "code",
- "source": [
- "!pip install pytorch-lightning"
- ],
- "metadata": {
- "id": "TkrmByc0aYVN"
- },
- "execution_count": null,
- "outputs": []
- },
- {
- "cell_type": "markdown",
- "source": [
- "Run either of the following two codes. If you run out of memory and crash, use a smaller model or a paid high-memory runtime"
- ],
- "metadata": {
- "id": "0SUK6Alv2ItS"
- }
- },
  {
  "cell_type": "code",
  "source": [
- "#@title <font size=\"-0\">If you specify the name of the model you want to convert and convert it manually</font>\n",
+ "#@title <font size=\"-0\">Convert the Models</font>\n",
+ "#@markdown Specify the models to be converted, separated by commas<br>\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",
- "model = \"v2-1_768-ema-pruned.ckpt\" #@param {type:\"string\"}\n",
- "model_name, model_ext = os.path.splitext(model)\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",
- "with torch.no_grad():\n",
- " if model_ext == \".safetensors\":\n",
- " weights = safetensors.torch.load_file(model_name + model_ext, device=\"cpu\")\n",
- " elif model_ext == \".ckpt\":\n",
- " weights = torch.load(model_name + model_ext, map_location=torch.device('cpu'))[\"state_dict\"]\n",
- " else:\n",
- " raise Exception(\"対応形式は.ckptと.safetensorsです\\n\" + f\"\\\"{model}\\\"は対応形式ではありません\")\n",
- " if as_fp16:\n",
- " model_name = model_name + \"-fp16\"\n",
- " for key in weights.keys():\n",
- " weights[key] = weights[key].half()\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",
- " safetensors.torch.save_file(weights, model_name + \".safetensors\")\n",
- " os.chdir(\"/content\")\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",
- " safetensors.torch.save_file(weights, model_name + \".safetensors\")\n",
- " del weights\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"
  ],
@@ -163,48 +167,26 @@
  "outputs": []
  },
  {
- "cell_type": "code",
+ "cell_type": "markdown",
  "source": [
- "#@title <font size=\"-0\">If you automatically convert all pre-loaded models</font>\n",
- "import os\n",
- "import glob\n",
- "import torch\n",
- "import safetensors.torch\n",
- "\n",
- "as_fp16 = True #@param {type:\"boolean\"}\n",
- "save_directly_to_Google_Drive = True #@param {type:\"boolean\"}\n",
+ "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",
- "with torch.no_grad():\n",
- " model_paths = glob.glob(r\"/content/*.ckpt\") + glob.glob(r\"/content/*.safetensors\") + glob.glob(r\"/content/*.pt\") + glob.glob(r\"/content/*.pth\")\n",
- " for model_path in model_paths:\n",
- " model_name, model_ext = os.path.splitext(os.path.basename(model_path))\n",
- " if model_ext == \".safetensors\":\n",
- " weights = safetensors.torch.load_file(model_name + model_ext, device=\"cpu\")\n",
- " elif model_ext == \".ckpt\":\n",
- " weights = torch.load(model_name + model_ext, map_location=torch.device('cpu'))[\"state_dict\"]\n",
- " else:\n",
- " print(\"対応形式は.ckpt\tと.safetensorsです\\n\" + f\"\\\"{model}\\\"は対応形式ではありません\")\n",
- " break\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",
- " safetensors.torch.save_file(weights, model_name + \".safetensors\")\n",
- " os.chdir(\"/content\")\n",
- " else:\n",
- " safetensors.torch.save_file(weights, model_name + \".safetensors\")\n",
- " del weights\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",
- "!reset"
+ "With the free ~12GB runtime, you can convert models up to ~10GB."
  ],
  "metadata": {
- "id": "5TUvrW5VzLst",
- "cellView": "form"
- },
- "execution_count": null,
- "outputs": []
+ "id": "0SUK6Alv2ItS"
+ }
  },
  {
  "cell_type": "markdown",