Skip to content

Instantly share code, notes, and snippets.

@sayakpaul
Created July 31, 2023 03:24
Show Gist options
  • Save sayakpaul/337b0988f08bd2cf2b248206f760e28f to your computer and use it in GitHub Desktop.
Save sayakpaul/337b0988f08bd2cf2b248206f760e28f to your computer and use it in GitHub Desktop.
Display the source blob
Display the rendered blob
Raw
{
"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
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment