Update as_safetensors+fp16.ipynb

as_safetensors+fp16.ipynb

@@ -121,7 +121,7 @@
         "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",
+        "  elif model.endswith((\".ckpt\", \".safetensors\", \".yaml\", \".pt\")):\n",
         "    shutil.copy(\"/content/drive/MyDrive/\" + model, \"/content/\" + model) # get the model from mydrive\n",
         "  else:\n",
         "    print(f\"\\\"{model}\\\"はURLではなく、正しい形式のファイルでもありません\")"
@@ -150,16 +150,26 @@
         "  from huggingface_hub import HfApi, Repository\n",
         "\n",
         "models = \"wd-1-4-anime_e1.ckpt, wd-1-4-anime_e1.yaml\" #@param {type:\"string\"}\n",
+        "pruning = True #@param {type:\"boolean\"}\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",
         "clip_fix = \"fix err key\" #@param [\"off\", \"fix err key\", \"del err key\"]\n",
         "uninvited_key = \"cond_stage_model.transformer.text_model.embeddings.position_ids\"\n",
+        "save_type = \".safetensors\" #@param [\".safetensors\", \".ckpt\"]\n",
+        "merge_vae = \"\" #@param [\"\", \"vae-ft-mse-840000-ema-pruned.ckpt\", \"kl-f8-anime.ckpt\", \"kl-f8-anime2.ckpt\", \"Anything-V3.0.vae.pt\"] {allow-input: true}\n",
+        "save_directly_to_Google_Drive = False #@param {type:\"boolean\"}\n",
         "#@markdown 変換したモデルをHugging Faceに投稿する場合は「yourname/yourrepo」の形式で投稿先リポジトリを指定<br>\n",
         "#@markdown 投稿しない場合は何も入力しない<br>\n",
         "# 5GB以上のファイルを投稿する場合は、投稿先リポジトリを丸ごとダウンロードする工程が挟まるので、時間がかかる場合があります\n",
         "repo_id = \"\" #@param {type:\"string\"}\n",
         "\n",
+        "vae_preset = {\n",
+        "    \"vae-ft-mse-840000-ema-pruned.ckpt\": \"https://huggingface.co/stabilityai/sd-vae-ft-mse-original/resolve/main/vae-ft-mse-840000-ema-pruned.ckpt\",\n",
+        "    \"kl-f8-anime.ckpt\": \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime.ckpt\",\n",
+        "    \"kl-f8-anime2.ckpt\": \"https://huggingface.co/hakurei/waifu-diffusion-v1-4/resolve/main/vae/kl-f8-anime2.ckpt\",\n",
+        "    \"Anything-V3.0.vae.pt\": \"https://huggingface.co/Linaqruf/anything-v3.0/resolve/main/Anything-V3.0.vae.pt\"}\n",
+        "if (merge_vae in vae_preset) & (not os.path.exists(merge_vae)):\n",
+        "  wget.download(vae_preset[merge_vae])\n",
+        "\n",
         "def upload_to_hugging_face(file_name):\n",
         "  api = HfApi()\n",
         "  api.upload_file(path_or_fileobj=file_name,\n",
@@ -183,6 +193,12 @@
         "    upload_to_hugging_face(file_name)\n",
         "  os.chdir(\"/content\")\n",
         "\n",
+        "#use `str.removeprefix(p)` in python 3.9+\n",
+        "def remove_prefix(input_string, prefix):\n",
+        "  if prefix and input_string.startswith(prefix):\n",
+        "    return input_string[len(prefix):]\n",
+        "  return input_string\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",
@@ -195,15 +211,20 @@
         "  elif (model_ext != \".safetensors\") & (model_ext != \".ckpt\"):\n",
         "    print(\"対応形式は.ckpt及び.safetensors並びに.yamlのみです\\n\" + f\"\\\"{model}\\\"は対応形式ではありません\")\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",
+        "    load_model = lambda filename: partial(safetensors.torch.load_file, device=\"cpu\")(filename) if os.path.splitext(filename)[1] == \".safetensors\" else partial(torch.load, map_location=torch.device(\"cpu\"))(filename)\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 pruning:\n",
+        "        model_name += \"-pruned\"\n",
+        "        for key in list(weights.keys()):\n",
+        "          if key.startswith(\"model_ema.\"):\n",
+        "            del weights[key]\n",
         "      if as_fp16:\n",
-        "        model_name = model_name + \"-fp16\"\n",
+        "        model_name += \"-fp16\"\n",
         "        for key in weights.keys():\n",
         "          weights[key] = weights[key].half()\n",
         "      if uninvited_key in weights:\n",
@@ -211,6 +232,14 @@
         "          del weights[uninvited_key]\n",
         "        if clip_fix == \"fix err key\":\n",
         "          weights[uninvited_key] = torch.tensor([list(range(77))],dtype=torch.int64)\n",
+        "      if merge_vae != \"\":\n",
+        "        vae_weights = load_model(merge_vae)\n",
+        "        if \"state_dict\" in vae_weights:\n",
+        "          vae_weights = vae_weights[\"state_dict\"]\n",
+        "        for key in weights.keys():\n",
+        "          if key.startswith(\"first_stage_model.\"):\n",
+        "            weights[key] = vae_weights[remove_prefix(key, \"first_stage_model.\")]\n",
+        "        del vae_weights\n",
         "      if save_directly_to_Google_Drive:\n",
         "        os.chdir(\"/content/drive/MyDrive\")\n",
         "      save_model(weights, saved_model := model_name + save_type)\n",
@@ -226,8 +255,8 @@
         "!reset"
       ],
       "metadata": {
-        "id": "9OmSG98HxJg2",
-        "cellView": "form"
+        "cellView": "form",
+        "id": "QSzZqGygdXM9"
       },
       "execution_count": null,
       "outputs": []
@@ -254,6 +283,15 @@
         "id": "0SUK6Alv2ItS"
       }
     },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "Hugging Faceに5GB以上のファイルを投稿する場合はメモリ消費量が約2倍になります"
+      ],
+      "metadata": {
+        "id": "8KU7VgNnE0Fy"
+      }
+    },
     {
       "cell_type": "markdown",
       "source": [