Update as_safetensors+fp16.ipynb

as_safetensors+fp16.ipynb

@@ -107,6 +107,7 @@
  "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",
@@ -139,7 +140,7 @@
  " elif (model_ext != \".safetensors\") & (model_ext != \".ckpt\"):\n",
  " print(\"対応形式は.ckpt及び.safetensors並びに.yamlのみです\\n\" + f\"\\\"{model}\\\"は対応形式ではありません\")\n",
  " else:\n",
- " load_model = safetensors.torch.load_file if model_ext == \".safetensors\" else torch.load\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",