Google Colab notebook for ConsistencyDecoder

ConsistencyDecoder_colab_notebook.ipynb

@@ -0,0 +1,112 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "gpuType": "T4"
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "rRvsavCyFB05"
+      },
+      "outputs": [],
+      "source": [
+        "# @title # Installation\n",
+        "!pip -q install git+https://github.com/openai/consistencydecoder.git\n",
+        "\n",
+        "!pip -q install numpy open-clip-torch scikit-image>=0.19 timm tomesd torch torchdiffeq torchsde diffusers transformers==4.30.2 accelerate==0.21.0"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# @title # Uploading a picture\n",
+        "\n",
+        "import locale\n",
+        "from google.colab import files\n",
+        "import os\n",
+        "import matplotlib.pyplot as plt\n",
+        "from PIL import Image\n",
+        "\n",
+        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
+        "img_folder = '/content/input' #@param {type:\"string\"}\n",
+        "!rm -rf {img_folder}\n",
+        "!mkdir -p {img_folder}\n",
+        "%cd {img_folder}\n",
+        "IMAGE = ''\n",
+        "\n",
+        "uploaded = files.upload()\n",
+        "\n",
+        "for file_name in uploaded.keys():\n",
+        "  IMAGE = os.path.join(img_folder, file_name)\n",
+        "\n",
+        "print(f\"Original image: {IMAGE}\")\n",
+        "plt.figure(figsize=(7, 5))\n",
+        "plt.axis('off')\n",
+        "plt.imshow(Image.open(IMAGE))\n",
+        "plt.show()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "XtgZFZsxFJDI"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# @title # Processing ConsistencyDecoder\n",
+        "\n",
+        "import torch, gc\n",
+        "from diffusers import StableDiffusionPipeline\n",
+        "from consistencydecoder import ConsistencyDecoder, save_image, load_image\n",
+        "import matplotlib.pyplot as plt\n",
+        "from PIL import Image\n",
+        "\n",
+        "OUTPUT = '/content/output.png'\n",
+        "\n",
+        "gc.collect()\n",
+        "torch.cuda.empty_cache()\n",
+        "pipe = StableDiffusionPipeline.from_pretrained(\n",
+        "    \"stablediffusionapi/deliberate3\", torch_dtype=torch.float16\n",
+        ")\n",
+        "pipe = pipe.to(\"cuda\")\n",
+        "pipe.vae.cuda()\n",
+        "decoder_consistency = ConsistencyDecoder(device=\"cuda:0\")\n",
+        "image = load_image(IMAGE)\n",
+        "\n",
+        "latent = pipe.vae.encode(image.half().cuda()).latent_dist.mean\n",
+        "\n",
+        "sample_consistency = decoder_consistency(latent)\n",
+        "save_image(sample_consistency, OUTPUT)\n",
+        "\n",
+        "print(f\"Processed image: {OUTPUT}\")\n",
+        "plt.figure(figsize=(7, 5))\n",
+        "plt.axis('off')\n",
+        "plt.imshow(Image.open(OUTPUT))\n",
+        "plt.show()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "ljJSPhpzGAuo"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}