DALLE2_Prior_+_MSE_VQGAN.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "DALLE2 Prior + MSE VQGAN.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "CppIQlPhhwhs"
      ]
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "d81a9f15c1124283ac326089e4683b90": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_26336e9f9d3149fb9947a05081561e72",
              "IPY_MODEL_b723c6e0ec404340a91ae3f6157a37e3",
              "IPY_MODEL_bf176e1a2cc0454ab1ee7a15a6929599"
            ],
            "layout": "IPY_MODEL_ceeffbaccc9b416fbc22552caf0338f8"
          }
        },
        "26336e9f9d3149fb9947a05081561e72": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_d77d9152d96d44e696efb7d847268239",
            "placeholder": "​",
            "style": "IPY_MODEL_079f5effc7654ab598216b4b0037525d",
            "value": "100%"
          }
        },
        "b723c6e0ec404340a91ae3f6157a37e3": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_e55fe6c508ba4e549f6e8adf54ddd648",
            "max": 553433881,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_4d2ba68f36234a49a14856b80f973bba",
            "value": 553433881
          }
        },
        "bf176e1a2cc0454ab1ee7a15a6929599": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_2098f93ec5c640698146797790594452",
            "placeholder": "​",
            "style": "IPY_MODEL_480f52532b5a4d81831e5a6d2e7829b8",
            "value": " 528M/528M [00:04&lt;00:00, 133MB/s]"
          }
        },
        "ceeffbaccc9b416fbc22552caf0338f8": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "d77d9152d96d44e696efb7d847268239": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "079f5effc7654ab598216b4b0037525d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "e55fe6c508ba4e549f6e8adf54ddd648": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "4d2ba68f36234a49a14856b80f973bba": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "2098f93ec5c640698146797790594452": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "480f52532b5a4d81831e5a6d2e7829b8": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "816f515e72bf491bb5c81253b7bfc644": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_6c34099bb5c44697b2fe0be72bcd9abe",
              "IPY_MODEL_bd36991157e8435ab3d6684f7980061e",
              "IPY_MODEL_9db9cf05439c452ebef3e9210f4b3a78"
            ],
            "layout": "IPY_MODEL_23e91f86e0d14a9ea2114ff398022bd9"
          }
        },
        "6c34099bb5c44697b2fe0be72bcd9abe": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_19b2bd53964840919e3d34e7c88f6b74",
            "placeholder": "​",
            "style": "IPY_MODEL_9c457f16b53745c1bd3967316603abf4",
            "value": ""
          }
        },
        "bd36991157e8435ab3d6684f7980061e": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "info",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_aa343266d1ca436fbe80aa94bd8b0ad9",
            "max": 1,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_52bc9f9a3c404e2790d827ad32e61b21",
            "value": 1
          }
        },
        "9db9cf05439c452ebef3e9210f4b3a78": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_7247db0f8bea4ee3b85414e44e738d8f",
            "placeholder": "​",
            "style": "IPY_MODEL_385ec88fa89445c9b1700c77bff7d156",
            "value": " 57/? [03:26&lt;00:00,  3.59s/it]"
          }
        },
        "23e91f86e0d14a9ea2114ff398022bd9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "19b2bd53964840919e3d34e7c88f6b74": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "9c457f16b53745c1bd3967316603abf4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "aa343266d1ca436fbe80aa94bd8b0ad9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": "20px"
          }
        },
        "52bc9f9a3c404e2790d827ad32e61b21": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "7247db0f8bea4ee3b85414e44e738d8f": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "385ec88fa89445c9b1700c77bff7d156": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Generates images from text prompts with VQGAN and CLIP (z+quantize method).\n",
        "\n",
        "By jbustter https://twitter.com/jbusted1 .\n",
        "Based on a notebook by Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings)\n",
        "\n",
        "# 🚧 LAION - Prior Testing 🚧"
      ],
      "metadata": {
        "id": "CppIQlPhhwhs"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!nvidia-smi"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mon May 16 23:59:43 2022       \n",
            "+-----------------------------------------------------------------------------+\n",
            "| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |\n",
            "|-------------------------------+----------------------+----------------------+\n",
            "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
            "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
            "|                               |                      |               MIG M. |\n",
            "|===============================+======================+======================|\n",
            "|   0  Tesla K80           Off  | 00000000:00:04.0 Off |                    0 |\n",
            "| N/A   51C    P8    29W / 149W |      0MiB / 11441MiB |      0%      Default |\n",
            "|                               |                      |                  N/A |\n",
            "+-------------------------------+----------------------+----------------------+\n",
            "                                                                               \n",
            "+-----------------------------------------------------------------------------+\n",
            "| Processes:                                                                  |\n",
            "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\n",
            "|        ID   ID                                                   Usage      |\n",
            "|=============================================================================|\n",
            "|  No running processes found                                                 |\n",
            "+-----------------------------------------------------------------------------+\n"
          ]
        }
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "TkUfzT60ZZ9q",
        "outputId": "9bc2982c-adc6-4c62-b94e-ce230352a061"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "!git clone https://github.com/openai/CLIP\n",
        "!git clone https://github.com/CompVis/taming-transformers\n",
        "!pip install ftfy regex tqdm omegaconf pytorch-lightning\n",
        "!pip install kornia\n",
        "!pip install einops\n",
        "!pip install scikit-learn-extra\n",
        "\n",
        "# dalle2 prior from NasirKhalid24\n",
        "!git clone https://github.com/lucidrains/DALLE2-pytorch.git\n",
        "!cd /content/DALLE2-pytorch/; git checkout 2eac7996faef08803230c7f3b2f9199e07c33b39; pip install ."
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Cloning into 'CLIP'...\n",
            "remote: Enumerating objects: 222, done.\u001b[K\n",
            "remote: Total 222 (delta 0), reused 0 (delta 0), pack-reused 222\u001b[K\n",
            "Receiving objects: 100% (222/222), 8.91 MiB | 13.15 MiB/s, done.\n",
            "Resolving deltas: 100% (113/113), done.\n",
            "Cloning into 'taming-transformers'...\n",
            "remote: Enumerating objects: 1335, done.\u001b[K\n",
            "remote: Total 1335 (delta 0), reused 0 (delta 0), pack-reused 1335\u001b[K\n",
            "Receiving objects: 100% (1335/1335), 409.77 MiB | 23.22 MiB/s, done.\n",
            "Resolving deltas: 100% (277/277), done.\n",
            "Collecting ftfy\n",
            "  Downloading ftfy-6.1.1-py3-none-any.whl (53 kB)\n",
            "\u001b[K     |████████████████████████████████| 53 kB 1.2 MB/s \n",
            "\u001b[?25hRequirement already satisfied: regex in /usr/local/lib/python3.7/dist-packages (2019.12.20)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (4.64.0)\n",
            "Collecting omegaconf\n",
            "  Downloading omegaconf-2.1.2-py3-none-any.whl (74 kB)\n",
            "\u001b[K     |████████████████████████████████| 74 kB 3.3 MB/s \n",
            "\u001b[?25hCollecting pytorch-lightning\n",
            "  Downloading pytorch_lightning-1.6.3-py3-none-any.whl (584 kB)\n",
            "\u001b[K     |████████████████████████████████| 584 kB 45.5 MB/s \n",
            "\u001b[?25hRequirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.7/dist-packages (from ftfy) (0.2.5)\n",
            "Collecting PyYAML>=5.1.0\n",
            "  Downloading PyYAML-6.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (596 kB)\n",
            "\u001b[K     |████████████████████████████████| 596 kB 44.2 MB/s \n",
            "\u001b[?25hCollecting antlr4-python3-runtime==4.8\n",
            "  Downloading antlr4-python3-runtime-4.8.tar.gz (112 kB)\n",
            "\u001b[K     |████████████████████████████████| 112 kB 44.3 MB/s \n",
            "\u001b[?25hRequirement already satisfied: tensorboard>=2.2.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (2.8.0)\n",
            "Requirement already satisfied: numpy>=1.17.2 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (1.21.6)\n",
            "Collecting fsspec[http]!=2021.06.0,>=2021.05.0\n",
            "  Downloading fsspec-2022.3.0-py3-none-any.whl (136 kB)\n",
            "\u001b[K     |████████████████████████████████| 136 kB 25.6 MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging>=17.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (21.3)\n",
            "Collecting pyDeprecate<0.4.0,>=0.3.1\n",
            "  Downloading pyDeprecate-0.3.2-py3-none-any.whl (10 kB)\n",
            "Collecting torchmetrics>=0.4.1\n",
            "  Downloading torchmetrics-0.8.2-py3-none-any.whl (409 kB)\n",
            "\u001b[K     |████████████████████████████████| 409 kB 29.8 MB/s \n",
            "\u001b[?25hRequirement already satisfied: typing-extensions>=4.0.0 in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (4.2.0)\n",
            "Requirement already satisfied: torch>=1.8.* in /usr/local/lib/python3.7/dist-packages (from pytorch-lightning) (1.11.0+cu113)\n",
            "Collecting aiohttp\n",
            "  Downloading aiohttp-3.8.1-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (1.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.1 MB 34.1 MB/s \n",
            "\u001b[?25hRequirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.23.0)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging>=17.0->pytorch-lightning) (3.0.8)\n",
            "Requirement already satisfied: google-auth<3,>=1.6.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.35.0)\n",
            "Requirement already satisfied: markdown>=2.6.8 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (3.3.6)\n",
            "Requirement already satisfied: protobuf>=3.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (3.17.3)\n",
            "Requirement already satisfied: absl-py>=0.4 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.0.0)\n",
            "Requirement already satisfied: grpcio>=1.24.3 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.44.0)\n",
            "Requirement already satisfied: werkzeug>=0.11.15 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.0.1)\n",
            "Requirement already satisfied: setuptools>=41.0.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (57.4.0)\n",
            "Requirement already satisfied: google-auth-oauthlib<0.5,>=0.4.1 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.4.6)\n",
            "Requirement already satisfied: tensorboard-data-server<0.7.0,>=0.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.6.1)\n",
            "Requirement already satisfied: tensorboard-plugin-wit>=1.6.0 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (1.8.1)\n",
            "Requirement already satisfied: wheel>=0.26 in /usr/local/lib/python3.7/dist-packages (from tensorboard>=2.2.0->pytorch-lightning) (0.37.1)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from absl-py>=0.4->tensorboard>=2.2.0->pytorch-lightning) (1.15.0)\n",
            "Requirement already satisfied: rsa<5,>=3.1.4 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (4.8)\n",
            "Requirement already satisfied: cachetools<5.0,>=2.0.0 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (4.2.4)\n",
            "Requirement already satisfied: pyasn1-modules>=0.2.1 in /usr/local/lib/python3.7/dist-packages (from google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (0.2.8)\n",
            "Requirement already satisfied: requests-oauthlib>=0.7.0 in /usr/local/lib/python3.7/dist-packages (from google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.2.0->pytorch-lightning) (1.3.1)\n",
            "Requirement already satisfied: importlib-metadata>=4.4 in /usr/local/lib/python3.7/dist-packages (from markdown>=2.6.8->tensorboard>=2.2.0->pytorch-lightning) (4.11.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=4.4->markdown>=2.6.8->tensorboard>=2.2.0->pytorch-lightning) (3.8.0)\n",
            "Requirement already satisfied: pyasn1<0.5.0,>=0.4.6 in /usr/local/lib/python3.7/dist-packages (from pyasn1-modules>=0.2.1->google-auth<3,>=1.6.3->tensorboard>=2.2.0->pytorch-lightning) (0.4.8)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (1.24.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2021.10.8)\n",
            "Requirement already satisfied: oauthlib>=3.0.0 in /usr/local/lib/python3.7/dist-packages (from requests-oauthlib>=0.7.0->google-auth-oauthlib<0.5,>=0.4.1->tensorboard>=2.2.0->pytorch-lightning) (3.2.0)\n",
            "Collecting yarl<2.0,>=1.0\n",
            "  Downloading yarl-1.7.2-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (271 kB)\n",
            "\u001b[K     |████████████████████████████████| 271 kB 47.4 MB/s \n",
            "\u001b[?25hCollecting asynctest==0.13.0\n",
            "  Downloading asynctest-0.13.0-py3-none-any.whl (26 kB)\n",
            "Collecting aiosignal>=1.1.2\n",
            "  Downloading aiosignal-1.2.0-py3-none-any.whl (8.2 kB)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (21.4.0)\n",
            "Collecting frozenlist>=1.1.1\n",
            "  Downloading frozenlist-1.3.0-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl (144 kB)\n",
            "\u001b[K     |████████████████████████████████| 144 kB 48.4 MB/s \n",
            "\u001b[?25hCollecting multidict<7.0,>=4.5\n",
            "  Downloading multidict-6.0.2-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (94 kB)\n",
            "\u001b[K     |████████████████████████████████| 94 kB 3.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: charset-normalizer<3.0,>=2.0 in /usr/local/lib/python3.7/dist-packages (from aiohttp->fsspec[http]!=2021.06.0,>=2021.05.0->pytorch-lightning) (2.0.12)\n",
            "Collecting async-timeout<5.0,>=4.0.0a3\n",
            "  Downloading async_timeout-4.0.2-py3-none-any.whl (5.8 kB)\n",
            "Building wheels for collected packages: antlr4-python3-runtime\n",
            "  Building wheel for antlr4-python3-runtime (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for antlr4-python3-runtime: filename=antlr4_python3_runtime-4.8-py3-none-any.whl size=141230 sha256=6d568a1c754d4e8c681934ae174ac618a0ee7b9504a6ea4692269cfd67efcdb1\n",
            "  Stored in directory: /root/.cache/pip/wheels/ca/33/b7/336836125fc9bb4ceaa4376d8abca10ca8bc84ddc824baea6c\n",
            "Successfully built antlr4-python3-runtime\n",
            "Installing collected packages: multidict, frozenlist, yarl, asynctest, async-timeout, aiosignal, pyDeprecate, fsspec, aiohttp, torchmetrics, PyYAML, antlr4-python3-runtime, pytorch-lightning, omegaconf, ftfy\n",
            "  Attempting uninstall: PyYAML\n",
            "    Found existing installation: PyYAML 3.13\n",
            "    Uninstalling PyYAML-3.13:\n",
            "      Successfully uninstalled PyYAML-3.13\n",
            "Successfully installed PyYAML-6.0 aiohttp-3.8.1 aiosignal-1.2.0 antlr4-python3-runtime-4.8 async-timeout-4.0.2 asynctest-0.13.0 frozenlist-1.3.0 fsspec-2022.3.0 ftfy-6.1.1 multidict-6.0.2 omegaconf-2.1.2 pyDeprecate-0.3.2 pytorch-lightning-1.6.3 torchmetrics-0.8.2 yarl-1.7.2\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "pydevd_plugins"
                ]
              }
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting kornia\n",
            "  Downloading kornia-0.6.4-py2.py3-none-any.whl (493 kB)\n",
            "\u001b[K     |████████████████████████████████| 493 kB 5.3 MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from kornia) (21.3)\n",
            "Requirement already satisfied: torch>=1.8.1 in /usr/local/lib/python3.7/dist-packages (from kornia) (1.11.0+cu113)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.8.1->kornia) (4.2.0)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->kornia) (3.0.8)\n",
            "Installing collected packages: kornia\n",
            "Successfully installed kornia-0.6.4\n",
            "Collecting einops\n",
            "  Downloading einops-0.4.1-py3-none-any.whl (28 kB)\n",
            "Installing collected packages: einops\n",
            "Successfully installed einops-0.4.1\n",
            "Cloning into 'DALLE2-pytorch'...\n",
            "remote: Enumerating objects: 1510, done.\u001b[K\n",
            "remote: Counting objects: 100% (642/642), done.\u001b[K\n",
            "remote: Compressing objects: 100% (118/118), done.\u001b[K\n",
            "remote: Total 1510 (delta 572), reused 583 (delta 523), pack-reused 868\u001b[K\n",
            "Receiving objects: 100% (1510/1510), 2.42 MiB | 13.94 MiB/s, done.\n",
            "Resolving deltas: 100% (1113/1113), done.\n",
            "Note: checking out '2eac7996faef08803230c7f3b2f9199e07c33b39'.\n",
            "\n",
            "You are in 'detached HEAD' state. You can look around, make experimental\n",
            "changes and commit them, and you can discard any commits you make in this\n",
            "state without impacting any branches by performing another checkout.\n",
            "\n",
            "If you want to create a new branch to retain commits you create, you may\n",
            "do so (now or later) by using -b with the checkout command again. Example:\n",
            "\n",
            "  git checkout -b <new-branch-name>\n",
            "\n",
            "HEAD is now at 2eac799 Additional image_embed metric (#75)\n",
            "Processing /content/DALLE2-pytorch\n",
            "\u001b[33m  DEPRECATION: A future pip version will change local packages to be built in-place without first copying to a temporary directory. We recommend you use --use-feature=in-tree-build to test your packages with this new behavior before it becomes the default.\n",
            "   pip 21.3 will remove support for this functionality. You can find discussion regarding this at https://github.com/pypa/pip/issues/7555.\u001b[0m\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (7.1.2)\n",
            "Collecting clip-anytorch\n",
            "  Downloading clip_anytorch-2.4.0-py3-none-any.whl (1.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.4 MB 5.4 MB/s \n",
            "\u001b[?25hCollecting coca-pytorch>=0.0.5\n",
            "  Downloading CoCa_pytorch-0.0.5-py3-none-any.whl (6.4 kB)\n",
            "Requirement already satisfied: einops>=0.4 in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (0.4.1)\n",
            "Collecting einops-exts>=0.0.3\n",
            "  Downloading einops_exts-0.0.3-py3-none-any.whl (3.8 kB)\n",
            "Collecting embedding-reader\n",
            "  Downloading embedding_reader-1.4.1-py3-none-any.whl (17 kB)\n",
            "Requirement already satisfied: kornia>=0.5.4 in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (0.6.4)\n",
            "Requirement already satisfied: pillow in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (7.1.2)\n",
            "Collecting resize-right>=0.0.2\n",
            "  Downloading resize_right-0.0.2-py3-none-any.whl (8.9 kB)\n",
            "Collecting rotary-embedding-torch\n",
            "  Downloading rotary_embedding_torch-0.1.5-py3-none-any.whl (4.1 kB)\n",
            "Requirement already satisfied: torch>=1.10 in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (1.11.0+cu113)\n",
            "Requirement already satisfied: torchvision in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (0.12.0+cu113)\n",
            "Requirement already satisfied: tqdm in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (4.64.0)\n",
            "Collecting vector-quantize-pytorch\n",
            "  Downloading vector_quantize_pytorch-0.6.0-py3-none-any.whl (6.5 kB)\n",
            "Collecting x-clip>=0.4.4\n",
            "  Downloading x_clip-0.5.1-py3-none-any.whl (1.4 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.4 MB 20.0 MB/s \n",
            "\u001b[?25hCollecting youtokentome\n",
            "  Downloading youtokentome-1.0.6-cp37-cp37m-manylinux2010_x86_64.whl (1.7 MB)\n",
            "\u001b[K     |████████████████████████████████| 1.7 MB 32.0 MB/s \n",
            "\u001b[?25hCollecting webdataset>=0.2.5\n",
            "  Downloading webdataset-0.2.5-py3-none-any.whl (46 kB)\n",
            "\u001b[K     |████████████████████████████████| 46 kB 4.1 MB/s \n",
            "\u001b[?25hRequirement already satisfied: fsspec>=2022.1.0 in /usr/local/lib/python3.7/dist-packages (from dalle2-pytorch==0.1.9) (2022.3.0)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from kornia>=0.5.4->dalle2-pytorch==0.1.9) (21.3)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.10->dalle2-pytorch==0.1.9) (4.2.0)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.7/dist-packages (from webdataset>=0.2.5->dalle2-pytorch==0.1.9) (1.21.6)\n",
            "Requirement already satisfied: pyyaml in /usr/local/lib/python3.7/dist-packages (from webdataset>=0.2.5->dalle2-pytorch==0.1.9) (6.0)\n",
            "Collecting braceexpand\n",
            "  Downloading braceexpand-0.1.7-py2.py3-none-any.whl (5.9 kB)\n",
            "Requirement already satisfied: regex in /usr/local/lib/python3.7/dist-packages (from x-clip>=0.4.4->dalle2-pytorch==0.1.9) (2019.12.20)\n",
            "Requirement already satisfied: ftfy in /usr/local/lib/python3.7/dist-packages (from x-clip>=0.4.4->dalle2-pytorch==0.1.9) (6.1.1)\n",
            "Requirement already satisfied: pandas<2,>=1.1.5 in /usr/local/lib/python3.7/dist-packages (from embedding-reader->dalle2-pytorch==0.1.9) (1.3.5)\n",
            "Requirement already satisfied: pyarrow<8,>=6.0.1 in /usr/local/lib/python3.7/dist-packages (from embedding-reader->dalle2-pytorch==0.1.9) (6.0.1)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.7/dist-packages (from pandas<2,>=1.1.5->embedding-reader->dalle2-pytorch==0.1.9) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas<2,>=1.1.5->embedding-reader->dalle2-pytorch==0.1.9) (2022.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.7.3->pandas<2,>=1.1.5->embedding-reader->dalle2-pytorch==0.1.9) (1.15.0)\n",
            "Requirement already satisfied: wcwidth>=0.2.5 in /usr/local/lib/python3.7/dist-packages (from ftfy->x-clip>=0.4.4->dalle2-pytorch==0.1.9) (0.2.5)\n",
            "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->kornia>=0.5.4->dalle2-pytorch==0.1.9) (3.0.8)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from torchvision->dalle2-pytorch==0.1.9) (2.23.0)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->torchvision->dalle2-pytorch==0.1.9) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->torchvision->dalle2-pytorch==0.1.9) (3.0.4)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->torchvision->dalle2-pytorch==0.1.9) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->torchvision->dalle2-pytorch==0.1.9) (2021.10.8)\n",
            "Building wheels for collected packages: dalle2-pytorch\n",
            "  Building wheel for dalle2-pytorch (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for dalle2-pytorch: filename=dalle2_pytorch-0.1.9-py3-none-any.whl size=1392106 sha256=07e3349f686d7ff8e1537784bf5a3d3266e34b929c397def55520b2d986cd64f\n",
            "  Stored in directory: /root/.cache/pip/wheels/ac/99/0b/89d881551d950e15c2b350e1dd3a10205d08eb92669de73f28\n",
            "Successfully built dalle2-pytorch\n",
            "Installing collected packages: braceexpand, youtokentome, x-clip, webdataset, vector-quantize-pytorch, rotary-embedding-torch, resize-right, embedding-reader, einops-exts, coca-pytorch, clip-anytorch, dalle2-pytorch\n",
            "Successfully installed braceexpand-0.1.7 clip-anytorch-2.4.0 coca-pytorch-0.0.5 dalle2-pytorch-0.1.9 einops-exts-0.0.3 embedding-reader-1.4.1 resize-right-0.0.2 rotary-embedding-torch-0.1.5 vector-quantize-pytorch-0.6.0 webdataset-0.2.5 x-clip-0.5.1 youtokentome-1.0.6\n"
          ]
        }
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "wSfISAhyPmyp",
        "outputId": "e9f22dc5-4177-4d4c-f1d6-7bec02192f09"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "\n",
        "!curl -L 'https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fconfigs%2Fmodel.yaml&dl=1' > vqgan_imagenet_f16_16384.yaml\n",
        "!curl -L 'https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/files/?p=%2Fckpts%2Flast.ckpt&dl=1' > vqgan_imagenet_f16_16384.ckpt\n",
        "!curl -L  https://huggingface.co/spaces/NasirKhalid24/Dalle2-Diffusion-Prior/resolve/main/larger-model.pth > model.pth # ~30k steps trained - larger model"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\n",
            "100   692  100   692    0     0   1028      0 --:--:-- --:--:-- --:--:--  1028\n",
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\n",
            "100  934M  100  934M    0     0  14.8M      0  0:01:02  0:01:02 --:--:-- 14.9M\n",
            "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
            "                                 Dload  Upload   Total   Spent    Left  Speed\n",
            "100   273  100   273    0     0   1695      0 --:--:-- --:--:-- --:--:--  1695\n",
            "100  576M  100  576M    0     0   186M      0  0:00:03  0:00:03 --:--:--  199M\n"
          ]
        }
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "FhhdWrSxQhwg",
        "outputId": "036f6a0c-cf6b-4173-bb81-54a856900bcd"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "import argparse\n",
        "import math\n",
        "from pathlib import Path\n",
        "import sys\n",
        "\n",
        "sys.path.append('./taming-transformers')\n",
        "\n",
        "from sklearn_extra.cluster import KMedoids\n",
        "from IPython import display\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from taming.models import cond_transformer, vqgan\n",
        "import torch\n",
        "from torch import nn, optim\n",
        "from torch.nn import functional as F\n",
        "from torchvision import transforms\n",
        "from torchvision.transforms import functional as TF\n",
        "from tqdm.notebook import tqdm\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from dalle2_pytorch import DiffusionPrior, DiffusionPriorNetwork\n",
        "\n",
        "from CLIP import clip\n",
        "\n",
        "import kornia.augmentation as K\n",
        "\n",
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "print('Using device:', device)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Using device: cuda:0\n"
          ]
        }
      ],
      "metadata": {
        "id": "EXMSuW2EQWsd",
        "outputId": "f123753e-fd48-4cd6-e337-47ccaf27d490",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "def noise_gen(shape):\n",
        "    n, c, h, w = shape\n",
        "    noise = torch.zeros([n, c, 1, 1])\n",
        "    for i in reversed(range(5)):\n",
        "        h_cur, w_cur = h // 2**i, w // 2**i\n",
        "        noise = F.interpolate(noise, (h_cur, w_cur), mode='bicubic', align_corners=False)\n",
        "        noise += torch.randn([n, c, h_cur, w_cur]) / 5\n",
        "    return noise\n",
        "\n",
        "\n",
        "def sinc(x):\n",
        "    return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))\n",
        "\n",
        "\n",
        "def lanczos(x, a):\n",
        "    cond = torch.logical_and(-a < x, x < a)\n",
        "    out = torch.where(cond, sinc(x) * sinc(x/a), x.new_zeros([]))\n",
        "    return out / out.sum()\n",
        "\n",
        "\n",
        "def ramp(ratio, width):\n",
        "    n = math.ceil(width / ratio + 1)\n",
        "    out = torch.empty([n])\n",
        "    cur = 0\n",
        "    for i in range(out.shape[0]):\n",
        "        out[i] = cur\n",
        "        cur += ratio\n",
        "    return torch.cat([-out[1:].flip([0]), out])[1:-1]\n",
        "\n",
        "\n",
        "def resample(input, size, align_corners=True):\n",
        "    n, c, h, w = input.shape\n",
        "    dh, dw = size\n",
        "\n",
        "    input = input.view([n * c, 1, h, w])\n",
        "\n",
        "    if dh < h:\n",
        "        kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)\n",
        "        pad_h = (kernel_h.shape[0] - 1) // 2\n",
        "        input = F.pad(input, (0, 0, pad_h, pad_h), 'reflect')\n",
        "        input = F.conv2d(input, kernel_h[None, None, :, None])\n",
        "\n",
        "    if dw < w:\n",
        "        kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)\n",
        "        pad_w = (kernel_w.shape[0] - 1) // 2\n",
        "        input = F.pad(input, (pad_w, pad_w, 0, 0), 'reflect')\n",
        "        input = F.conv2d(input, kernel_w[None, None, None, :])\n",
        "\n",
        "    input = input.view([n, c, h, w])\n",
        "    return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)\n",
        "    \n",
        "\n",
        "# def replace_grad(fake, real):\n",
        "#     return fake.detach() - real.detach() + real\n",
        "\n",
        "\n",
        "class ReplaceGrad(torch.autograd.Function):\n",
        "    @staticmethod\n",
        "    def forward(ctx, x_forward, x_backward):\n",
        "        ctx.shape = x_backward.shape\n",
        "        return x_forward\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_in):\n",
        "        return None, grad_in.sum_to_size(ctx.shape)\n",
        "\n",
        "\n",
        "class ClampWithGrad(torch.autograd.Function):\n",
        "    @staticmethod\n",
        "    def forward(ctx, input, min, max):\n",
        "        ctx.min = min\n",
        "        ctx.max = max\n",
        "        ctx.save_for_backward(input)\n",
        "        return input.clamp(min, max)\n",
        "\n",
        "    @staticmethod\n",
        "    def backward(ctx, grad_in):\n",
        "        input, = ctx.saved_tensors\n",
        "        return grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0), None, None\n",
        "\n",
        "replace_grad = ReplaceGrad.apply\n",
        "\n",
        "clamp_with_grad = ClampWithGrad.apply\n",
        "# clamp_with_grad = torch.clamp\n",
        "\n",
        "def vector_quantize(x, codebook):\n",
        "    d = x.pow(2).sum(dim=-1, keepdim=True) + codebook.pow(2).sum(dim=1) - 2 * x @ codebook.T\n",
        "    indices = d.argmin(-1)\n",
        "    x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook\n",
        "    return replace_grad(x_q, x)\n",
        "\n",
        "\n",
        "class Prompt(nn.Module):\n",
        "    def __init__(self, embed, weight=1., stop=float('-inf')):\n",
        "        super().__init__()\n",
        "        self.register_buffer('embed', embed)\n",
        "        self.register_buffer('weight', torch.as_tensor(weight))\n",
        "        self.register_buffer('stop', torch.as_tensor(stop))\n",
        "\n",
        "    def forward(self, input):\n",
        "        \n",
        "        input_normed = F.normalize(input.unsqueeze(1), dim=2)\n",
        "        embed_normed = F.normalize((self.embed).unsqueeze(0), dim=2)\n",
        "\n",
        "        dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)\n",
        "        dists = dists * self.weight.sign()\n",
        "        return self.weight.abs() * replace_grad(dists, torch.maximum(dists, self.stop)).mean()\n",
        "\n",
        "def select_best(target_embedding, embed_1, embed_2):\n",
        "  similarity1 = F.cosine_similarity(target_embedding, embed_1)\n",
        "  similarity2 = F.cosine_similarity(target_embedding, embed_2)\n",
        "\n",
        "  if similarity1 > similarity2:\n",
        "    return embed_1\n",
        "    \n",
        "  return embed_2\n",
        "\n",
        "def select_medoid(cond_embed, n_samples):\n",
        "    cond_batch = torch.cat([cond_embed for _ in range(n_samples) ])\n",
        "    prior_embeds = prior.p_sample_loop((n_samples,512), text_cond=dict(text_embed=cond_batch))\n",
        "    priors_normed = F.normalize(prior_embeds)\n",
        "    other = priors_normed\n",
        "    sims = priors_normed @ other.T\n",
        "    i = sims.mean(dim=0).argmax().item()\n",
        "    best = prior_embeds[i,:].unsqueeze(0)\n",
        "    return best\n",
        "\n",
        "def parse_prompt(prompt):\n",
        "    vals = prompt.rsplit(':', 2)\n",
        "    vals = vals + ['', '1', '-inf'][len(vals):]\n",
        "    return vals[0], float(vals[1]), float(vals[2])\n",
        "\n",
        "def one_sided_clip_loss(input, target, labels=None, logit_scale=100):\n",
        "    input_normed = F.normalize(input, dim=-1)\n",
        "    target_normed = F.normalize(target, dim=-1)\n",
        "    logits = input_normed @ target_normed.T * logit_scale\n",
        "    if labels is None:\n",
        "        labels = torch.arange(len(input), device=logits.device)\n",
        "    return F.cross_entropy(logits, labels)\n",
        "\n",
        "class MakeCutouts(nn.Module):\n",
        "    def __init__(self, cut_size, cutn, cut_pow=1.):\n",
        "        super().__init__()\n",
        "        self.cut_size = cut_size\n",
        "        self.cutn = cutn\n",
        "        self.cut_pow = cut_pow\n",
        "\n",
        "        self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))\n",
        "        self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))\n",
        "\n",
        "    def set_cut_pow(self, cut_pow):\n",
        "      self.cut_pow = cut_pow\n",
        "\n",
        "    def forward(self, input):\n",
        "        sideY, sideX = input.shape[2:4]\n",
        "        max_size = min(sideX, sideY)\n",
        "        min_size = min(sideX, sideY, self.cut_size)\n",
        "        cutouts = []\n",
        "        cutouts_full = []\n",
        "        \n",
        "        min_size_width = min(sideX, sideY)\n",
        "        lower_bound = float(self.cut_size/min_size_width)\n",
        "        \n",
        "        for ii in range(self.cutn):\n",
        "            \n",
        "            \n",
        "          size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)\n",
        "\n",
        "          offsetx = torch.randint(0, sideX - size + 1, ())\n",
        "          offsety = torch.randint(0, sideY - size + 1, ())\n",
        "          cutout = input[:, :, offsety:offsety + size, offsetx:offsetx + size]\n",
        "          cutouts.append(resample(cutout, (self.cut_size, self.cut_size)))\n",
        "\n",
        "        \n",
        "        cutouts = torch.cat(cutouts, dim=0)\n",
        "\n",
        "        if args.use_augs:\n",
        "          cutouts = augs(cutouts)\n",
        "\n",
        "        if args.noise_fac:\n",
        "          facs = cutouts.new_empty([cutouts.shape[0], 1, 1, 1]).uniform_(0, args.noise_fac)\n",
        "          cutouts = cutouts + facs * torch.randn_like(cutouts)\n",
        "        \n",
        "\n",
        "        return clamp_with_grad(cutouts, 0, 1)\n",
        "\n",
        "\n",
        "def load_vqgan_model(config_path, checkpoint_path):\n",
        "    config = OmegaConf.load(config_path)\n",
        "    if config.model.target == 'taming.models.vqgan.VQModel':\n",
        "        model = vqgan.VQModel(**config.model.params)\n",
        "        model.eval().requires_grad_(False)\n",
        "        model.init_from_ckpt(checkpoint_path)\n",
        "    elif config.model.target == 'taming.models.cond_transformer.Net2NetTransformer':\n",
        "        parent_model = cond_transformer.Net2NetTransformer(**config.model.params)\n",
        "        parent_model.eval().requires_grad_(False)\n",
        "        parent_model.init_from_ckpt(checkpoint_path)\n",
        "        model = parent_model.first_stage_model\n",
        "    elif config.model.target == 'taming.models.vqgan.GumbelVQ':\n",
        "        model = vqgan.GumbelVQ(**config.model.params)\n",
        "        model.eval().requires_grad_(False)\n",
        "        model.init_from_ckpt(checkpoint_path)\n",
        "    else:\n",
        "        raise ValueError(f'unknown model type: {config.model.target}')\n",
        "    del model.loss\n",
        "    return model\n",
        "\n",
        "def resize_image(image, out_size):\n",
        "    ratio = image.size[0] / image.size[1]\n",
        "    area = min(image.size[0] * image.size[1], out_size[0] * out_size[1])\n",
        "    size = round((area * ratio)**0.5), round((area / ratio)**0.5)\n",
        "    return image.resize(size, Image.LANCZOS)\n",
        "\n",
        "class TVLoss(nn.Module):\n",
        "    def forward(self, input):\n",
        "        input = F.pad(input, (0, 1, 0, 1), 'replicate')\n",
        "        x_diff = input[..., :-1, 1:] - input[..., :-1, :-1]\n",
        "        y_diff = input[..., 1:, :-1] - input[..., :-1, :-1]\n",
        "        diff = x_diff**2 + y_diff**2 + 1e-8\n",
        "        return diff.mean(dim=1).sqrt().mean()\n",
        "\n",
        "class GaussianBlur2d(nn.Module):\n",
        "    def __init__(self, sigma, window=0, mode='reflect', value=0):\n",
        "        super().__init__()\n",
        "        self.mode = mode\n",
        "        self.value = value\n",
        "        if not window:\n",
        "            window = max(math.ceil((sigma * 6 + 1) / 2) * 2 - 1, 3)\n",
        "        if sigma:\n",
        "            kernel = torch.exp(-(torch.arange(window) - window // 2)**2 / 2 / sigma**2)\n",
        "            kernel /= kernel.sum()\n",
        "        else:\n",
        "            kernel = torch.ones([1])\n",
        "        self.register_buffer('kernel', kernel)\n",
        "\n",
        "    def forward(self, input):\n",
        "        n, c, h, w = input.shape\n",
        "        input = input.view([n * c, 1, h, w])\n",
        "        start_pad = (self.kernel.shape[0] - 1) // 2\n",
        "        end_pad = self.kernel.shape[0] // 2\n",
        "        input = F.pad(input, (start_pad, end_pad, start_pad, end_pad), self.mode, self.value)\n",
        "        input = F.conv2d(input, self.kernel[None, None, None, :])\n",
        "        input = F.conv2d(input, self.kernel[None, None, :, None])\n",
        "        return input.view([n, c, h, w])\n",
        "\n",
        "class EMATensor(nn.Module):\n",
        "    \"\"\"implmeneted by Katherine Crowson\"\"\"\n",
        "    def __init__(self, tensor, decay):\n",
        "        super().__init__()\n",
        "        self.tensor = nn.Parameter(tensor)\n",
        "        self.register_buffer('biased', torch.zeros_like(tensor))\n",
        "        self.register_buffer('average', torch.zeros_like(tensor))\n",
        "        self.decay = decay\n",
        "        self.register_buffer('accum', torch.tensor(1.))\n",
        "        self.update()\n",
        "    \n",
        "    @torch.no_grad()\n",
        "    def update(self):\n",
        "        if not self.training:\n",
        "            raise RuntimeError('update() should only be called during training')\n",
        "\n",
        "        self.accum *= self.decay\n",
        "        self.biased.mul_(self.decay)\n",
        "        self.biased.add_((1 - self.decay) * self.tensor)\n",
        "        self.average.copy_(self.biased)\n",
        "        self.average.div_(1 - self.accum)\n",
        "\n",
        "    def forward(self):\n",
        "        if self.training:\n",
        "            return self.tensor\n",
        "        return self.average\n",
        "\n",
        "\n",
        "def load_prior(model_path=\"model.pth\"):\n",
        "  state_dict = torch.load(model_path)['model']\n",
        "\n",
        "  DIM = 512\n",
        "\n",
        "  prior_network = DiffusionPriorNetwork( \n",
        "      dim=512,\n",
        "      depth=12, \n",
        "      dim_head=64, \n",
        "      heads=12,\n",
        "      normformer=False\n",
        "  ).to(device)\n",
        "\n",
        "  diffusion_prior = DiffusionPrior( \n",
        "      net=prior_network,\n",
        "      clip=None,\n",
        "      image_embed_dim=512, \n",
        "      timesteps=1000,\n",
        "      cond_drop_prob=0.1, \n",
        "      loss_type=\"l2\", \n",
        "      condition_on_text_encodings=False\n",
        "  ).to(device)\n",
        "\n",
        "  diffusion_prior.load_state_dict(state_dict, strict=True)\n",
        "\n",
        "  return diffusion_prior\n",
        "\n",
        "%mkdir /content/vids"
      ],
      "outputs": [],
      "metadata": {
        "id": "JvnTBhPGT1gn"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# ARGS"
      ],
      "metadata": {
        "id": "WN4OtaLbHBN6"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "from typing import OrderedDict\n",
        "#@title Generation Params\n",
        "prompt = \"An oil painting of ancient ruins lost to time, by james gurney in the style of rembrandt\" #@param {type: \"string\"}\n",
        "init_weight = .2 #@param {type:\"number\"}\n",
        "mse_decay_rate = 50 #@param {type:\"number\"}\n",
        "mse_epoches =  5#@param {type:\"number\"}\n",
        "\n",
        "image_size = 256,256 #@param [\"256,256\", \"512,512\", \"256,512\", \"512,256\"] {type:\"raw\", allow-input: true}\n",
        "cutn = 64 #@param {type: \"integer\"}\n",
        "max_iterations =  200#@param {type: \"integer\"}\n",
        "#@markdown ----\n",
        "# use_prior = False #@param {type: \"boolean\"}\n",
        "seed = 1 #@param [\"256\", \"512\"] {type:\"raw\", allow-input: true}\n",
        "prior_influence = 0.2 #@param {type:\"slider\", min:0, max:1, step:0.01}\n",
        "\n",
        "#@markdown ----\n",
        "# @markdown \n",
        "\n",
        "select_best_method = 'sim_to_target' # @param ['prior_medoid', 'sim_to_target']\n",
        "n_medoid_samples =  25# @param {type: 'integer'}\n",
        "\n",
        "#@markdown prior_influence: controlls how much to weight the prior embedding (left=prior has no influnce | right=prior has 100% influence) \n",
        "#@markdown \n",
        "#@markdown ----\n",
        "#@markdown NOTE: I've found the prior works better with less iterations, more epochs, and a lower init_weight (~0.2), and a lower decay rate\n",
        "#@markdown but you should defintely experiment! I'm not sure what optimal parameters are.\n",
        "#@markdown\n",
        "#@markdown ----\n",
        "\n",
        "from pprint import pprint\n",
        "from collections import OrderedDict\n",
        "arg_summary_dict = OrderedDict({\n",
        "    \"prompt\": prompt,\n",
        "    \"cutn\":cutn,\n",
        "    \"prior_influence\": prior_influence,\n",
        "    \"init_weight\": init_weight,\n",
        "    \"mse_decay_rate\": mse_decay_rate,\n",
        "    \"mse_epochs\":mse_epoches,\n",
        "    \"max_iterations\": max_iterations,\n",
        "    \"seed\": seed,\n",
        "    \"select_best_method\": select_best_method, \n",
        "})\n",
        "\n",
        "priors = [\"without\", \"with\", \"both\"]\n",
        "\n",
        "print(\"-----please copy below if you post in #dalle2-results-----\")\n",
        "print(\"```\")\n",
        "for k,v in arg_summary_dict.items():\n",
        "  print(f\"{k} : {v}\")\n",
        "print(\"```\")\n",
        "print(\"-----please copy the above if you post in #dalle2-results-----\")\n",
        "\n",
        "args = argparse.Namespace(\n",
        "    \n",
        "    prompts=[prompt],\n",
        "    size=[*image_size], \n",
        "    init_image= None,\n",
        "    init_weight= init_weight,\n",
        "\n",
        "    # clip model settings\n",
        "    clip_model='ViT-B/32',\n",
        "    vqgan_config='vqgan_imagenet_f16_16384.yaml',         \n",
        "    vqgan_checkpoint='vqgan_imagenet_f16_16384.ckpt',\n",
        "    step_size=0.95,\n",
        "    \n",
        "    # cutouts / crops\n",
        "    cutn=cutn,\n",
        "    cut_pow=1,\n",
        "\n",
        "    # display\n",
        "    display_freq=25,\n",
        "    seed=seed,\n",
        "    use_augs = True,\n",
        "    noise_fac= 0.1,\n",
        "\n",
        "    record_generation=True,\n",
        "\n",
        "    # noise and other constraints\n",
        "    use_noise = None,\n",
        "    constraint_regions = False,#\n",
        "    \n",
        "    \n",
        "    # add noise to embedding\n",
        "    noise_prompt_weights = None,\n",
        "    noise_prompt_seeds = [14575],#\n",
        "\n",
        "    # mse settings\n",
        "    mse_withzeros = True,\n",
        "    mse_decay_rate = mse_decay_rate,\n",
        "    mse_epoches = mse_epoches,\n",
        "\n",
        "    # end itteration\n",
        "    max_itter = max_iterations,\n",
        "\n",
        "    # prior sampling\n",
        "    select_best_method = select_best_method,\n",
        "    n_medoid_samples = n_medoid_samples,\n",
        ")\n",
        "\n",
        "mse_decay = 0\n",
        "if args.init_weight:\n",
        "  mse_decay = args.init_weight / args.mse_epoches\n",
        "\n",
        "# <AUGMENTATIONS>\n",
        "augs = nn.Sequential(\n",
        "    \n",
        "    K.RandomHorizontalFlip(p=0.5),\n",
        "    K.RandomAffine(degrees=30, translate=0.1, p=0.8, padding_mode='border'), # padding_mode=2\n",
        "    K.RandomPerspective(0.2,p=0.4, ),\n",
        "    K.ColorJitter(hue=0.01, saturation=0.01, p=0.7),\n",
        "\n",
        "    )\n",
        "\n",
        "noise = noise_gen([1, 3, args.size[0], args.size[1]])\n",
        "image = TF.to_pil_image(noise.div(5).add(0.5).clamp(0, 1)[0])\n",
        "image.save('init3.png')"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "-----please copy below if you post in #dalle2-results-----\n",
            "```\n",
            "prompt : An oil painting of ancient ruins lost to time, by james gurney in the style of rembrandt\n",
            "cutn : 64\n",
            "prior_influence : 0.2\n",
            "init_weight : 0.2\n",
            "mse_decay_rate : 50\n",
            "mse_epochs : 5\n",
            "max_iterations : 200\n",
            "seed : 1\n",
            "```\n",
            "-----please copy the above if you post in #dalle2-results-----\n"
          ]
        }
      ],
      "metadata": {
        "id": "tLw9p5Rzacso",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6fd7996e-32f0-4bde-978e-605b76a0964c"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 🎨"
      ],
      "metadata": {
        "id": "QXgTa_JWi7Sn"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#@title Actually do the run!\n",
        "for use_prior in priors:\n",
        "  tv_loss = TVLoss() \n",
        "  mse_weight = args.init_weight\n",
        "\n",
        "  model = load_vqgan_model(args.vqgan_config, args.vqgan_checkpoint).to(device)\n",
        "  perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)\n",
        "  \n",
        "  if use_prior == \"without\":\n",
        "    print(\"using only vqgan + clip\")\n",
        "    append = \"_without_prior\"\n",
        "\n",
        "  elif use_prior == \"with\":\n",
        "    print(\"using only prior\")\n",
        "    prior = load_prior(\"model.pth\")\n",
        "    append = \"_with_prior\"\n",
        "\n",
        "  elif use_prior == \"both\":\n",
        "    print(f\"using both with weight {prior_influence}\")\n",
        "    prior = load_prior(\"model.pth\")\n",
        "    append = \"_with_both\"\n",
        "\n",
        "\n",
        "  cut_size = perceptor.visual.input_resolution\n",
        "\n",
        "  if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':\n",
        "      e_dim = 256\n",
        "      n_toks = model.quantize.n_embed\n",
        "      z_min = model.quantize.embed.weight.min(dim=0).values[None, :, None, None]\n",
        "      z_max = model.quantize.embed.weight.max(dim=0).values[None, :, None, None]\n",
        "  else:\n",
        "      e_dim = model.quantize.e_dim\n",
        "      n_toks = model.quantize.n_e\n",
        "      z_min = model.quantize.embedding.weight.min(dim=0).values[None, :, None, None]\n",
        "      z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]\n",
        "\n",
        "\n",
        "  make_cutouts = MakeCutouts(cut_size, args.cutn, cut_pow=args.cut_pow)\n",
        "\n",
        "  f = 2**(model.decoder.num_resolutions - 1)\n",
        "  toksX, toksY = args.size[0] // f, args.size[1] // f\n",
        "\n",
        "  if args.seed is not None:\n",
        "      torch.manual_seed(args.seed)\n",
        "\n",
        "  if args.init_image:\n",
        "      pil_image = Image.open(args.init_image).convert('RGB')\n",
        "      pil_image = pil_image.resize((toksX * 16, toksY * 16), Image.LANCZOS)\n",
        "      pil_image = TF.to_tensor(pil_image)\n",
        "      if args.use_noise:\n",
        "        pil_image = pil_image + args.use_noise * torch.randn_like(pil_image) \n",
        "      z, *_ = model.encode(pil_image.to(device).unsqueeze(0) * 2 - 1)\n",
        "\n",
        "  else:\n",
        "      \n",
        "      one_hot = F.one_hot(torch.randint(n_toks, [toksY * toksX], device=device), n_toks).float()\n",
        "\n",
        "      if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':\n",
        "          z = one_hot @ model.quantize.embed.weight\n",
        "      else:\n",
        "          z = one_hot @ model.quantize.embedding.weight\n",
        "      z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2)\n",
        "\n",
        "  if args.mse_withzeros and not args.init_image:\n",
        "    z_orig = torch.zeros_like(z)\n",
        "  else:\n",
        "    z_orig = z.clone()\n",
        "\n",
        "  z.requires_grad = True\n",
        "\n",
        "  opt = optim.Adam([z], lr=args.step_size, weight_decay=0.00000000)\n",
        "\n",
        "  normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
        "                                  std=[0.26862954, 0.26130258, 0.27577711])\n",
        "\n",
        "  pMs = []\n",
        "\n",
        "  if args.noise_prompt_weights and args.noise_prompt_seeds:\n",
        "    for seed, weight in zip(args.noise_prompt_seeds, args.noise_prompt_weights):\n",
        "      gen = torch.Generator().manual_seed(seed)\n",
        "      embed = torch.empty([1, perceptor.visual.output_dim]).normal_(generator=gen)\n",
        "      pMs.append(Prompt(embed, weight).to(device))\n",
        "\n",
        "  if args.select_best_method == 'sim_to_target':\n",
        "    for prompt in args.prompts:\n",
        "        txt, weight, stop = parse_prompt(prompt)\n",
        "        embed = perceptor.encode_text(clip.tokenize(txt).to(device)).float()\n",
        "        prior_embed = embed\n",
        "        if use_prior == \"with\":\n",
        "          pred_embed_1 = prior.p_sample_loop((1,512), text_cond=dict(text_embed=prior_embed))\n",
        "          pred_embed_2 = prior.p_sample_loop((1,512), text_cond=dict(text_embed=prior_embed))\n",
        "          embed = select_best(prior_embed, pred_embed_1, pred_embed_2)\n",
        "        elif use_prior == \"both\":\n",
        "          pred_embed_1 = prior.p_sample_loop((1,512), text_cond=dict(text_embed=prior_embed))\n",
        "          pred_embed_2 = prior.p_sample_loop((1,512), text_cond=dict(text_embed=prior_embed))\n",
        "          pred_embed = select_best(prior_embed, pred_embed_1, pred_embed_2)\n",
        "\n",
        "          weighted_pred = prior_influence*pred_embed\n",
        "          weighted_text = (1-prior_influence)*embed\n",
        "          embed = weighted_pred+weighted_text\n",
        "\n",
        "        pMs.append(Prompt(embed, weight, stop).to(device))\n",
        "  elif args.select_best_method == 'prior_medoid':\n",
        "    print(\"using medoid prior estimation\")\n",
        "    for prompt in args.prompts:\n",
        "      txt, weight, stop = parse_prompt(prompt)\n",
        "      cond_embed = perceptor.encode_text(clip.tokenize(txt).to(device)).float()\n",
        "      if use_prior in ('with','both'):\n",
        "        prior_embed = select_medoid(cond_embed, args.n_medoid_samples)\n",
        "        if use_prior == 'both':\n",
        "        \n",
        "          embed = (prior_influence * prior_embed) + ((1-prior_influence) * cond_embed)\n",
        "        else:\n",
        "          embed = prior_embed\n",
        "      else:\n",
        "        embed = cond_embed\n",
        "    pMs.append(Prompt(embed, weight, stop).to(device))\n",
        "\n",
        "\n",
        "\n",
        "  def synth(z, quantize=True):\n",
        "      if args.constraint_regions:\n",
        "        z = replace_grad(z, z * z_mask)\n",
        "\n",
        "      if quantize:\n",
        "        if args.vqgan_checkpoint == 'vqgan_openimages_f16_8192.ckpt':\n",
        "          z_q = vector_quantize(z.movedim(1, 3), model.quantize.embed.weight).movedim(3, 1)\n",
        "        else:\n",
        "          z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)\n",
        "\n",
        "      else:\n",
        "        z_q = z.model\n",
        "\n",
        "      return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def checkin(i, losses):\n",
        "      losses_str = ', '.join(f'{loss.item():g}' for loss in losses)\n",
        "      tqdm.write(f'i: {i}, loss: {sum(losses).item():g}, losses: {losses_str}')\n",
        "      out = synth(z, True)\n",
        "\n",
        "      TF.to_pil_image(out[0].cpu()).save('progress' + append + '.png')   \n",
        "      display.display(display.Image('progress' + append + '.png')) \n",
        "\n",
        "\n",
        "  def ascend_txt():\n",
        "      global mse_weight\n",
        "\n",
        "      out = synth(z)\n",
        "\n",
        "      if args.record_generation:\n",
        "        with torch.no_grad():\n",
        "          global vid_index\n",
        "          out_a = synth(z, True)\n",
        "          TF.to_pil_image(out_a[0].cpu()).save(f'/content/vids/{vid_index}.png')\n",
        "          vid_index += 1\n",
        "\n",
        "      cutouts = make_cutouts(out)\n",
        "\n",
        "      iii = perceptor.encode_image(normalize(cutouts)).float()\n",
        "\n",
        "      result = []\n",
        "\n",
        "      if args.init_weight:\n",
        "          \n",
        "          global z_orig\n",
        "          \n",
        "          result.append(F.mse_loss(z, z_orig) * mse_weight / 2)\n",
        "          # result.append(F.mse_loss(z, z_orig) * ((1/torch.tensor((i)*2 + 1))*mse_weight) / 2)\n",
        "\n",
        "          with torch.no_grad():\n",
        "            if i > 0 and i%args.mse_decay_rate==0 and i <= args.mse_decay_rate*args.mse_epoches:\n",
        "\n",
        "              if mse_weight - mse_decay > 0 and mse_weight - mse_decay >= mse_decay:\n",
        "                mse_weight = mse_weight - mse_decay\n",
        "                print(f\"updated mse weight: {mse_weight}\")\n",
        "              else:\n",
        "                mse_weight = 0\n",
        "                print(f\"updated mse weight: {mse_weight}\")\n",
        "\n",
        "      for prompt in pMs:\n",
        "          result.append(prompt(iii))\n",
        "\n",
        "      return result\n",
        "\n",
        "  vid_index = 0\n",
        "  def train(i):\n",
        "      \n",
        "      opt.zero_grad()\n",
        "      lossAll = ascend_txt()\n",
        "\n",
        "      if (i+1) % args.display_freq == 0:\n",
        "          checkin(i, lossAll)\n",
        "      \n",
        "      loss = sum(lossAll)\n",
        "\n",
        "      loss.backward()\n",
        "      opt.step()\n",
        "\n",
        "  i = 0\n",
        "  try:\n",
        "      with tqdm() as pbar:\n",
        "          while True and i != args.max_itter:\n",
        "\n",
        "              train(i)\n",
        "\n",
        "              if i > 0 and i%args.mse_decay_rate==0 and i <= args.mse_decay_rate * args.mse_epoches:\n",
        "                \n",
        "                opt = optim.Adam([z], lr=args.step_size, weight_decay=0.00000000)\n",
        "\n",
        "              i += 1\n",
        "              pbar.update()\n",
        "\n",
        "  except KeyboardInterrupt:\n",
        "      pass\n",
        "\n",
        "im1 = Image.open(\"progress_without_prior.png\")\n",
        "im2 = Image.open(\"progress_with_both.png\")\n",
        "im3 = Image.open(\"progress_with_prior.png\")\n",
        "\n",
        "\n",
        "parameter_text = ''\n",
        "for k,v in arg_summary_dict.items():\n",
        "  parameter_text += f\"{k} : {v}\\n\"\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "fig, ax = plt.subplots(1,3, figsize=(12, 9))\n",
        "fig.patch.set_facecolor('white')\n",
        "\n",
        "ax[0].title.set_text('Without prior')\n",
        "ax[0].axis('off')\n",
        "ax[0].imshow(im1)\n",
        "\n",
        "ax[1].title.set_text('Combined')\n",
        "ax[1].axis('off')\n",
        "ax[1].imshow(im2)\n",
        "\n",
        "ax[2].title.set_text('With prior')\n",
        "ax[2].axis('off')\n",
        "ax[2].imshow(im3)\n",
        "\n",
        "fig.suptitle(parameter_text)\n",
        "print(\"Just copy the image (right click and copy image) and paste it into the #dalle2-results channel in the discord!\")"
      ],
      "outputs": [],
      "metadata": {
        "id": "g7EDme5RYCrt",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000,
          "referenced_widgets": [
            "d81a9f15c1124283ac326089e4683b90",
            "26336e9f9d3149fb9947a05081561e72",
            "b723c6e0ec404340a91ae3f6157a37e3",
            "bf176e1a2cc0454ab1ee7a15a6929599",
            "ceeffbaccc9b416fbc22552caf0338f8",
            "d77d9152d96d44e696efb7d847268239",
            "079f5effc7654ab598216b4b0037525d",
            "e55fe6c508ba4e549f6e8adf54ddd648",
            "4d2ba68f36234a49a14856b80f973bba",
            "2098f93ec5c640698146797790594452",
            "480f52532b5a4d81831e5a6d2e7829b8",
            "816f515e72bf491bb5c81253b7bfc644",
            "6c34099bb5c44697b2fe0be72bcd9abe",
            "bd36991157e8435ab3d6684f7980061e",
            "9db9cf05439c452ebef3e9210f4b3a78",
            "23e91f86e0d14a9ea2114ff398022bd9",
            "19b2bd53964840919e3d34e7c88f6b74",
            "9c457f16b53745c1bd3967316603abf4",
            "aa343266d1ca436fbe80aa94bd8b0ad9",
            "52bc9f9a3c404e2790d827ad32e61b21",
            "7247db0f8bea4ee3b85414e44e738d8f",
            "385ec88fa89445c9b1700c77bff7d156"
          ]
        },
        "outputId": "b93578a6-85ee-4978-ebbc-95b6d4133607",
        "cellView": "form"
      }
    }
  ]
}