Update as_safetensors+fp16.ipynb

as_safetensors+fp16.ipynb

@@ -130,12 +130,12 @@
       "source": [
         "#@title <font size=\"-0\">モデルを変換</font>\n",
         "#@markdown 変換するモデルをカンマ区切りで任意個指定<br>\n",
-        "#@markdown 何も入力されていない場合は、読み込まれている全てのモデルが変換される\n",
+        "#@markdown 何も入力されていない場合は、読み込まれている全てのモデルが変換される<br>\n",
+        "#@markdown `a.ckpt, -, b.safetensors`のような形式でモデルの引き算ができます\n",
         "import os\n",
         "import glob\n",
         "import torch\n",
         "import safetensors.torch\n",
-        "from functools import partial\n",
         "\n",
         "from sys import modules\n",
         "if \"huggingface_hub\" in modules:\n",
@@ -147,7 +147,7 @@
         "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",
+        "merge_vae = \"\" #@param [\"\", \"vae-ft-mse-840000-ema-pruned.ckpt\", \"kl-f8-anime.ckpt\", \"kl-f8-anime2.ckpt\", \"anything-v4.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",
@@ -158,8 +158,8 @@
         "    \"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",
+        "    \"anything-v4.0.vae.pt\": \"https://huggingface.co/andite/anything-v4.0/resolve/main/anything-v4.0.vae.pt\"}\n",
+        "if (merge_vae in vae_preset) and (not os.path.exists(merge_vae)):\n",
         "  wget.download(vae_preset[merge_vae])\n",
         "\n",
         "def upload_to_hugging_face(file_name):\n",
@@ -191,25 +191,74 @@
         "    return input_string[len(prefix):]\n",
         "  return input_string\n",
         "\n",
+        "load_model = lambda m: safetensors.torch.load_file(m, device=\"cpu\") if os.path.splitext(m)[1] == \".safetensors\" else torch.load(m, map_location=torch.device(\"cpu\"))\n",
+        "save_model = safetensors.torch.save_file if save_type == \".safetensors\" else torch.save\n",
+        "\n",
+        "# --- def merge ---#\n",
+        "@torch.no_grad()\n",
+        "def merged(model_a, model_b, fab, fa, fb):\n",
+        "  weights_a = load_model(model_a) if isinstance(model_a, str) else model_a\n",
+        "  weights_b = load_model(model_b) if isinstance(model_b, str) else model_b\n",
+        "  if \"state_dict\" in weights_a:\n",
+        "    weights_a = weights_a[\"state_dict\"]\n",
+        "  if \"state_dict\" in weights_b:\n",
+        "    weights_b = weights_b[\"state_dict\"]\n",
+        "  for key in list(weights_a.keys() or weights_b.keys()):\n",
+        "    if isinstance(weights_a[key], dict):\n",
+        "      del weights_a[key]\n",
+        "    if isinstance(weights_b[key], dict):\n",
+        "      del weights_b[key]\n",
+        "    if key.startswith(\"model.\") or key.startswith(\"model_ema.\"):\n",
+        "      if (key in weights_a) and (key in weights_b):\n",
+        "        weights_a[key] = fab(weights_a[key], weights_b[key])\n",
+        "        del weights_b[key]\n",
+        "      elif key in weights_a:\n",
+        "        weights_a[key] = fa(weights_a[key])\n",
+        "      elif key in weights_b:\n",
+        "        weights_a[key] = fb(weights_b[key])\n",
+        "        del weights_b[key]\n",
+        "  del weights_b\n",
+        "  return weights_a\n",
+        "\n",
+        "def add(model_a, model_b):\n",
+        "  return merged(model_a, model_b, lambda a, b: a + b, lambda a: a, lambda b: b)\n",
+        "\n",
+        "def difference(model_a, model_b):\n",
+        "  return merged(model_a, model_b, lambda a, b: a - b, lambda a: a, lambda b: -b)\n",
+        "\n",
+        "def add_difference(model_a, model_b, model_c):\n",
+        "  return add(model_a, difference(model_b, model_c))\n",
+        "# --- end merge ---#\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",
+        "  models = [os.path.basename(m) for m in glob.glob(r\"/content/*.ckpt\") + glob.glob(r\"/content/*.safetensors\") + glob.glob(r\"/content/*.yaml\") if not os.path.basename(m) in vae_preset]\n",
         "else:\n",
         "  models = [m.strip() for m in models.split(\",\")]\n",
         "\n",
-        "for model in models:\n",
+        "for i, model in enumerate(models):\n",
         "  model_name, model_ext = os.path.splitext(model)\n",
+        "  # a.ckpt, - ,b.ckpt # - or b.ckpt\n",
+        "  if (models[i] == \"-\") or (models[i - 1] == \"-\"):\n",
+        "    continue\n",
         "  if model_ext == \".yaml\":\n",
         "    convert_yaml(model)\n",
-        "  elif (model_ext != \".safetensors\") & (model_ext != \".ckpt\"):\n",
+        "  elif (model_ext != \".safetensors\") and (model_ext != \".ckpt\"):\n",
         "    print(\"対応形式は.ckpt及び.safetensors並びに.yamlのみです\\n\" + f\"\\\"{model}\\\"は対応形式ではありません\")\n",
         "  else:\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",
+        "      # a.ckpt, - ,b.ckpt # a.ckpt\n",
+        "      if (i < len(models) - 1) and (models[i + 1] == \"-\"):\n",
+        "        weights = difference(model, models[i + 2])\n",
+        "        model_name = f\"{model_name}-{os.path.splitext(models[i + 2])[0]}\"\n",
+        "      # otherwise\n",
+        "      else:\n",
+        "        weights = load_model(model)\n",
         "      if \"state_dict\" in weights:\n",
         "        weights = weights[\"state_dict\"]\n",
+        "      for key in list(weights.keys()):\n",
+        "        if isinstance(weights[key], dict):\n",
+        "          del weights[key] # to fix the broken model\n",
         "      if pruning:\n",
         "        model_name += \"-pruned\"\n",
         "        for key in list(weights.keys()):\n",
@@ -275,15 +324,6 @@
         "id": "0SUK6Alv2ItS"
       }
     },
-    {
-      "cell_type": "markdown",
-      "source": [
-        "Hugging Faceに5GB以上のファイルを投稿する場合はメモリ消費量が約2倍になります"
-      ],
-      "metadata": {
-        "id": "8KU7VgNnE0Fy"
-      }
-    },
     {
       "cell_type": "markdown",
       "source": [