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",