tiny_autoencoder.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e1b7be2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget -q https://huggingface.co/madebyollin/taesd/resolve/main/taesd_encoder.safetensors\n",
    "# !wget -q https://huggingface.co/madebyollin/taesd/resolve/main/taesd_decoder.safetensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "09282d9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !ls -lf *.safetensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "49b51813",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b62e06f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from diffusers import TinyAutoencoder\n",
    "\n",
    "tiny = TinyAutoencoder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "398ad678",
   "metadata": {},
   "outputs": [],
   "source": [
    "import safetensors.torch \n",
    "\n",
    "encoder = safetensors.torch.load_file(\"taesd_encoder.safetensors\")\n",
    "decoder = safetensors.torch.load_file(\"taesd_decoder.safetensors\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "938f9805",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_state_dict = {}\n",
    "\n",
    "for k in encoder:\n",
    "    new_state_dict.update({f\"encoder.layers.{k}\": encoder[k]})\n",
    "for k in decoder:\n",
    "    new_state_dict.update({f\"decoder.layers.{k}\": decoder[k]})\n",
    "    \n",
    "tiny.load_state_dict(new_state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a488d8fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget -q https://huggingface.co/datasets/diffusers/docs-images/resolve/main/if/person.png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c519a0b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 3, 256, 256])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from PIL import Image\n",
    "import torchvision.transforms.functional as TF\n",
    "\n",
    "img = Image.open(\"person.png\").convert(\"RGB\").resize((256, 256))\n",
    "img = TF.to_tensor(img).unsqueeze(0).to(\"cpu\")\n",
    "img.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a7213e29",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 4, 32, 32])\n",
      "True\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "# Comes from\n",
    "# https://colab.research.google.com/gist/sayakpaul/de9bd84ddd4780256e0464b2303bbf3d/scratchpad.ipynb\n",
    "expected = np.array([0.8008, 0.7907, 0.9320, 1.1116, 0.6613, 0.6602, 0.6918, 0.5838])\n",
    "\n",
    "with torch.no_grad():\n",
    "    enc = tiny.encode(img, return_dict=False)[0]\n",
    "    print(enc.shape)\n",
    "    outputs = enc[0, :2, :2, :2].flatten().numpy()\n",
    "    print(np.allclose(outputs, expected, atol=1e-4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3eaeb7da",
   "metadata": {},
   "outputs": [],
   "source": [
    "tiny.save_pretrained(\"tiny_autoencoder\")\n",
    "tiny.save_pretrained(\"tiny_autoencoder\", safe_serialization=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4881f0b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "total 39744\r\n",
      "-rw-r--r--  1 sayakpaul  staff   501B Jul 31 08:41 config.json\r\n",
      "-rw-r--r--  1 sayakpaul  staff   9.4M Jul 31 08:41 diffusion_pytorch_model.bin\r\n",
      "-rw-r--r--  1 sayakpaul  staff   9.3M Jul 31 08:41 diffusion_pytorch_model.safetensors\r\n"
     ]
    }
   ],
   "source": [
    "!ls -lh tiny_autoencoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f38eea9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca16644b752d44158b79a3e2ccd51262",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "diffusion_pytorch_model.safetensors:   0%|          | 0.00/9.79M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6a0685ffeaa548cd9eca14c74e305738",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Upload 1 LFS files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "'https://huggingface.co/sayakpaul/taesd-diffusers/tree/main/'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from huggingface_hub import HfApi, create_repo\n",
    "\n",
    "repo_id = create_repo(\"sayakpaul/taesd-diffusers\", exist_ok=True).repo_id\n",
    "api = HfApi()\n",
    "api.upload_folder(repo_id=repo_id, folder_path=\"tiny_autoencoder\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}