Doohickey_Prompt_Engine.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "XIOfbBNOrjfE"
      ]
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "947b41c2a67d45988437ec78a17c78b6": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "VBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "VBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "VBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_45d06ae73f794835bf210731d6516ccc",
              "IPY_MODEL_7d8b8c000f62406fa7c28e5124f9fa0a",
              "IPY_MODEL_0bad96652bab41dda231489747fcee21",
              "IPY_MODEL_71b9e0a7a668425996b346f273c1be91"
            ],
            "layout": "IPY_MODEL_6f9f683a3c35437e9f998588fc3f14b3"
          }
        },
        "45d06ae73f794835bf210731d6516ccc": {
          "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_0ebf02bb1bcd4c9d9be165901b7c8157",
            "placeholder": "​",
            "style": "IPY_MODEL_09628a43b24d4d98b4dbe77190139fc7",
            "value": "<center> <img\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.svg\nalt='Hugging Face'> <br> Copy a token from <a\nhref=\"https://huggingface.co/settings/tokens\" target=\"_blank\">your Hugging Face\ntokens page</a> and paste it below. <br> Immediately click login after copying\nyour token or it might be stored in plain text in this notebook file. </center>"
          }
        },
        "7d8b8c000f62406fa7c28e5124f9fa0a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "PasswordModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "PasswordModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "PasswordView",
            "continuous_update": true,
            "description": "Token:",
            "description_tooltip": null,
            "disabled": false,
            "layout": "IPY_MODEL_b914444ac7554c2e8b979a84170cd3f7",
            "placeholder": "​",
            "style": "IPY_MODEL_37643af36a0a4e289d9bf4c0620a7f76",
            "value": ""
          }
        },
        "0bad96652bab41dda231489747fcee21": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "",
            "description": "Login",
            "disabled": false,
            "icon": "",
            "layout": "IPY_MODEL_87c1e744fbc146298e596facf6f4b30d",
            "style": "IPY_MODEL_263bd008bf8844709d1458339da7666a",
            "tooltip": ""
          }
        },
        "71b9e0a7a668425996b346f273c1be91": {
          "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_f55c2bf8d9564141935a8357bb4c4a3c",
            "placeholder": "​",
            "style": "IPY_MODEL_d3896a88512a48718dbf11c93a92296c",
            "value": "\n<b>Pro Tip:</b> If you don't already have one, you can create a dedicated\n'notebooks' token with 'write' access, that you can then easily reuse for all\nnotebooks. </center>"
          }
        },
        "6f9f683a3c35437e9f998588fc3f14b3": {
          "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": "center",
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": "flex",
            "flex": null,
            "flex_flow": "column",
            "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": "50%"
          }
        },
        "0ebf02bb1bcd4c9d9be165901b7c8157": {
          "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
          }
        },
        "09628a43b24d4d98b4dbe77190139fc7": {
          "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": ""
          }
        },
        "b914444ac7554c2e8b979a84170cd3f7": {
          "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
          }
        },
        "37643af36a0a4e289d9bf4c0620a7f76": {
          "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": ""
          }
        },
        "87c1e744fbc146298e596facf6f4b30d": {
          "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
          }
        },
        "263bd008bf8844709d1458339da7666a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ButtonStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        },
        "f55c2bf8d9564141935a8357bb4c4a3c": {
          "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
          }
        },
        "d3896a88512a48718dbf11c93a92296c": {
          "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": "code",
      "source": [
        "#@title Install libraries + Log in to 🤗\n",
        "import os\n",
        "from IPython.display import clear_output\n",
        "import time\n",
        "if not os.path.exists(\"installed.txt\"):\n",
        "    # red lines, it's fines, that's what i always say\n",
        "    !pip install transformers diffusers lpips -q\n",
        "    # !pip install git+https://github.com/openai/CLIP -q\n",
        "    !pip install open_clip_torch -q\n",
        "    !pip install wget -q\n",
        "    !sudo apt-get install git-lfs\n",
        "    !cat \"test\" > installed.txt\n",
        "    !mkdir /content/output\n",
        "    print(\"Installed libraries\")\n",
        "    time.sleep(1) # just so that the user can see a glimpse of the print to know it went succesfuly\n",
        "    clear_output(wait=False)\n",
        "else:\n",
        "    print(\"Libraries already installed.\")\n",
        "\n",
        "from huggingface_hub import notebook_login\n",
        "notebook_login()"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 288,
          "referenced_widgets": [
            "947b41c2a67d45988437ec78a17c78b6",
            "45d06ae73f794835bf210731d6516ccc",
            "7d8b8c000f62406fa7c28e5124f9fa0a",
            "0bad96652bab41dda231489747fcee21",
            "71b9e0a7a668425996b346f273c1be91",
            "6f9f683a3c35437e9f998588fc3f14b3",
            "0ebf02bb1bcd4c9d9be165901b7c8157",
            "09628a43b24d4d98b4dbe77190139fc7",
            "b914444ac7554c2e8b979a84170cd3f7",
            "37643af36a0a4e289d9bf4c0620a7f76",
            "87c1e744fbc146298e596facf6f4b30d",
            "263bd008bf8844709d1458339da7666a",
            "f55c2bf8d9564141935a8357bb4c4a3c",
            "d3896a88512a48718dbf11c93a92296c"
          ]
        },
        "cellView": "form",
        "id": "TVkSjx4tueAy",
        "outputId": "65cbe186-31ff-4006-8004-f05244f25634"
      },
      "execution_count": 1,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Libraries already installed.\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "947b41c2a67d45988437ec78a17c78b6"
            }
          },
          "metadata": {}
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup program files\n",
        "\n",
        "top 10 pages of danbooru tags, prompt building blocks from clip interrogator"
      ],
      "metadata": {
        "id": "XIOfbBNOrjfE"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%writefile doohickey-slim.py\n",
        "import torch\n",
        "torch.manual_seed(0)\n",
        "from transformers import CLIPTextModel, CLIPTokenizer\n",
        "from diffusers import AutoencoderKL, UNet2DConditionModel\n",
        "from diffusers import LMSDiscreteScheduler\n",
        "from tqdm.auto import tqdm, trange\n",
        "from torch import autocast\n",
        "import PIL.Image as PImage\n",
        "import numpy\n",
        "from torchvision import transforms\n",
        "import torchvision.transforms.functional as f\n",
        "import random\n",
        "import requests\n",
        "from io import BytesIO\n",
        "import sys\n",
        "\n",
        "model_name = sys.argv[4]\n",
        "model_type = \"diffusers\"\n",
        "\n",
        "# Set device\n",
        "torch_device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "offload_device = \"cpu\"\n",
        "\n",
        "vae = AutoencoderKL.from_pretrained(model_name, subfolder=\"vae\", use_auth_token=True).cpu().to(torch.bfloat16)\n",
        "tokenizer = CLIPTokenizer.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
        "text_encoder = CLIPTextModel.from_pretrained(model_name, subfolder=\"text_encoder\", use_auth_token=True).cpu().to(torch.bfloat16)\n",
        "unet = UNet2DConditionModel.from_pretrained(model_name, subfolder=\"unet\", use_auth_token=True).cpu().to(torch.bfloat16)\n",
        "\n",
        "def requires_grad(model, val=False):\n",
        "    for param in model.parameters():\n",
        "        param.requires_grad = val\n",
        "\n",
        "requires_grad(vae)\n",
        "requires_grad(text_encoder)\n",
        "requires_grad(unet)\n",
        "\n",
        "#@title Set up generation loop\n",
        "\n",
        "to_tensor_tfm = transforms.ToTensor()\n",
        "\n",
        "# mismatch of tons of image encoding / decoding / loading functions i cant be asked to clean up right now\n",
        "\n",
        "def pil_to_latent(input_im):\n",
        "  # Single image -> single latent in a batch (so size 1, 4, 64, 64)\n",
        "  with torch.no_grad():\n",
        "      with autocast(\"cuda\"):\n",
        "        latent = vae.encode(to_tensor_tfm(input_im.convert(\"RGB\")).unsqueeze(0).to(torch_device)*2-1).latent_dist # Note scaling\n",
        "#   print(latent)\n",
        "  return 0.18215 * latent.mode() # or .mean or .sample\n",
        "\n",
        "def latents_to_pil(latents):\n",
        "  # bath of latents -> list of images\n",
        "  latents = (1 / 0.18215) * latents\n",
        "  with torch.no_grad():\n",
        "    image = vae.decode(latents)\n",
        "  image = (image / 2 + 0.5).clamp(0, 1)\n",
        "  image = image.detach().cpu().permute(0, 2, 3, 1).numpy()\n",
        "  images = (image * 255).round().astype(\"uint8\")\n",
        "  pil_images = [Image.fromarray(image) for image in images]\n",
        "  return pil_images\n",
        "\n",
        "def get_latent_from_url(url, size=(512,512)):\n",
        "    response = requests.get(url)\n",
        "    img = PImage.open(BytesIO(response.content))\n",
        "    img = img.resize(size).convert(\"RGB\")\n",
        "    latent = pil_to_latent(img)\n",
        "    return latent\n",
        "\n",
        "def scale_and_decode(latents):\n",
        "    with autocast(\"cuda\"):\n",
        "        # scale and decode the image latents with vae\n",
        "        latents = 1 / 0.18215 * latents\n",
        "        with torch.no_grad():\n",
        "            image = vae.decode(latents).sample.squeeze(0)\n",
        "        image = f.to_pil_image((image / 2 + 0.5).clamp(0, 1))\n",
        "        return image\n",
        "\n",
        "def fetch(url_or_path):\n",
        "    import io\n",
        "    if str(url_or_path).startswith('http://') or str(url_or_path).startswith('https://'):\n",
        "        r = requests.get(url_or_path)\n",
        "        r.raise_for_status()\n",
        "        fd = io.BytesIO()\n",
        "        fd.write(r.content)\n",
        "        fd.seek(0)\n",
        "        return PImage.open(fd).convert('RGB')\n",
        "    return PImage.open(open(url_or_path, 'rb')).convert('RGB')\n",
        "\n",
        "\"\"\"\n",
        "grabs all text up to the first occurrence of ':' \n",
        "uses the grabbed text as a sub-prompt, and takes the value following ':' as weight\n",
        "if ':' has no value defined, defaults to 1.0\n",
        "repeats until no text remaining\n",
        "\"\"\"\n",
        "def split_weighted_subprompts(text, split=\":\"):\n",
        "    remaining = len(text)\n",
        "    prompts = []\n",
        "    weights = []\n",
        "    while remaining > 0:\n",
        "        if split in text:\n",
        "            idx = text.index(split) # first occurrence from start\n",
        "            # grab up to index as sub-prompt\n",
        "            prompt = text[:idx]\n",
        "            remaining -= idx\n",
        "            # remove from main text\n",
        "            text = text[idx+1:]\n",
        "            # find value for weight \n",
        "            if \" \" in text:\n",
        "                idx = text.index(\" \") # first occurence\n",
        "            else: # no space, read to end\n",
        "                idx = len(text)\n",
        "            if idx != 0:\n",
        "                try:\n",
        "                    weight = float(text[:idx])\n",
        "                except: # couldn't treat as float\n",
        "                    print(f\"Warning: '{text[:idx]}' is not a value, are you missing a space?\")\n",
        "                    weight = 1.0\n",
        "            else: # no value found\n",
        "                weight = 1.0\n",
        "            # remove from main text\n",
        "            remaining -= idx\n",
        "            text = text[idx+1:]\n",
        "            # append the sub-prompt and its weight\n",
        "            prompts.append(prompt)\n",
        "            weights.append(weight)\n",
        "        else: # no : found\n",
        "            if len(text) > 0: # there is still text though\n",
        "                # take remainder as weight 1\n",
        "                prompts.append(text)\n",
        "                weights.append(1.0)\n",
        "            remaining = 0\n",
        "    print(prompts, weights)\n",
        "    return prompts, weights \n",
        "\n",
        "\n",
        "# from some stackoverflow comment\n",
        "import numpy as np\n",
        "def lerp(a, b, x):\n",
        "    \"linear interpolation\"\n",
        "    return a + x * (b - a)\n",
        "def fade(t):\n",
        "    \"6t^5 - 15t^4 + 10t^3\"\n",
        "    return 6 * t**5 - 15 * t**4 + 10 * t**3\n",
        "def gradient(h, x, y):\n",
        "    \"grad converts h to the right gradient vector and return the dot product with (x,y)\"\n",
        "    vectors = np.array([[0, 1], [0, -1], [1, 0], [-1, 0]])\n",
        "    g = vectors[h % 4]\n",
        "    return g[:, :, 0] * x + g[:, :, 1] * y\n",
        "def perlin(x, y, seed=0):\n",
        "    # permutation table\n",
        "    np.random.seed(seed)\n",
        "    p = np.arange(256, dtype=int)\n",
        "    np.random.shuffle(p)\n",
        "    p = np.stack([p, p]).flatten()\n",
        "    # coordinates of the top-left\n",
        "    xi, yi = x.astype(int), y.astype(int)\n",
        "    # internal coordinates\n",
        "    xf, yf = x - xi, y - yi\n",
        "    # fade factors\n",
        "    u, v = fade(xf), fade(yf)\n",
        "    # noise components\n",
        "    n00 = gradient(p[p[xi] + yi], xf, yf)\n",
        "    n01 = gradient(p[p[xi] + yi + 1], xf, yf - 1)\n",
        "    n11 = gradient(p[p[xi + 1] + yi + 1], xf - 1, yf - 1)\n",
        "    n10 = gradient(p[p[xi + 1] + yi], xf - 1, yf)\n",
        "    # combine noises\n",
        "    x1 = lerp(n00, n10, u)\n",
        "    x2 = lerp(n01, n11, u)  # FIX1: I was using n10 instead of n01\n",
        "    return lerp(x1, x2, v)  # FIX2: I also had to reverse x1 and x2 here\n",
        "\n",
        "def sample(args):\n",
        "    global in_channels\n",
        "    global text_encoder # uugghhhghhghgh\n",
        "    global vae # UUGHGHHGHGH\n",
        "    global unet # .hggfkgjks;ldjf\n",
        "    # prompt = args.prompt\n",
        "    prompts, weights = split_weighted_subprompts(args.prompt)\n",
        "    h,w = args.size\n",
        "    steps = args.steps\n",
        "    scale = args.scale\n",
        "    classifier_guidance = args.classifier_guidance\n",
        "    use_init = len(args.init_img)>1\n",
        "    if args.seed!=-1:\n",
        "        seed = args.seed\n",
        "        generator = torch.manual_seed(seed)\n",
        "    else:\n",
        "        seed = random.randint(0,10_000)\n",
        "        generator = torch.manual_seed(seed)\n",
        "    print(f\"Generating with seed {seed}...\")\n",
        "    \n",
        "    # tokenize / encode text\n",
        "    tokens = [tokenizer(prompt, padding=\"max_length\", max_length=tokenizer.model_max_length, truncation=True, return_tensors=\"pt\") for prompt in prompts]\n",
        "    with torch.no_grad():\n",
        "        # move CLIP to cuda\n",
        "        text_encoder = text_encoder.to(torch_device)\n",
        "        text_embeddings = [text_encoder(tok.input_ids.to(torch_device))[0].unsqueeze(0) for tok in tokens]\n",
        "        text_embeddings = [text_embeddings[i]*weights[i] for i in range(len(text_embeddings))]\n",
        "        text_embeddings = torch.cat(text_embeddings, 0).sum(0)\n",
        "        max_length = 77\n",
        "        uncond_input = tokenizer(\n",
        "            [\"\"], padding=\"max_length\", max_length=max_length, return_tensors=\"pt\"\n",
        "        )\n",
        "        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0] \n",
        "        text_embeddings = torch.cat([uncond_embeddings, text_embeddings])\n",
        "        # move it back to CPU so there's more vram for generating\n",
        "        text_encoder = text_encoder.to(offload_device)\n",
        "    images = []\n",
        "\n",
        "    if args.lpips_guidance:\n",
        "        import lpips\n",
        "        lpips_model = lpips.LPIPS(net='vgg').to(torch_device)\n",
        "        init = to_tensor_tfm(fetch(args.init_img).resize(args.size)).to(torch_device)\n",
        "\n",
        "    for batch_n in trange(args.batches):\n",
        "        with autocast(\"cuda\"):\n",
        "            # unet = unet.to(torch_device)\n",
        "            scheduler.set_timesteps(steps)\n",
        "            if not use_init or args.start_step==0:\n",
        "                latents = torch.randn(\n",
        "                    (1, in_channels, h//8, w//8),\n",
        "                    generator=generator\n",
        "                )\n",
        "                latents = latents.to(torch_device)\n",
        "                latents = latents * scheduler.sigmas[0]\n",
        "                start_step = args.start_step\n",
        "            else:\n",
        "                # Start step\n",
        "                start_step = args.start_step -1\n",
        "                start_sigma = scheduler.sigmas[start_step]\n",
        "                start_timestep = int(scheduler.timesteps[start_step])\n",
        "\n",
        "                # Prep latents\n",
        "                vae = vae.to(torch_device)\n",
        "                encoded = get_latent_from_url(args.init_img)\n",
        "                if not classifier_guidance:\n",
        "                    vae = vae.to(offload_device)\n",
        "\n",
        "                # ???????????????????????????????????????\n",
        "                encoded = f.resize(encoded, (h//8,w//8))\n",
        "\n",
        "                noise = torch.randn_like(encoded)\n",
        "                sigmas = scheduler.match_shape(scheduler.sigmas[start_step], noise)\n",
        "                noisy_samples = encoded + noise * sigmas\n",
        "\n",
        "                latents = noisy_samples.to(torch_device).to(torch.bfloat16)\n",
        "\n",
        "                \n",
        "            \n",
        "            if args.perlin_multi != 0 and args.start_step==0:\n",
        "                linx = np.linspace(0, 5, h // 8, endpoint=False)\n",
        "                liny = np.linspace(0, 5, w // 8, endpoint=False)\n",
        "                x, y = np.meshgrid(liny, linx)\n",
        "                p = [np.expand_dims(perlin(x, y, seed=i), 0) for i in range(4)] # reproducable seed\n",
        "                p = np.concatenate(p, 0)\n",
        "                p = torch.tensor(p).unsqueeze(0).cuda()\n",
        "                # latents = latents + (p * args.perlin_multi).to(torch_device).to(torch.bfloat16)\n",
        "                latents = latents*(1-(args.perlin_multi*0.1)) + (p*args.perlin_multi).to(torch_device).to(torch.bfloat16)\n",
        "\n",
        "                \n",
        "            for i, t in enumerate(scheduler.timesteps):\n",
        "                if i > args.start_step:\n",
        "                    latent_model_input = torch.cat([latents]*2)\n",
        "                    sigma = scheduler.sigmas[i]\n",
        "                    latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)\n",
        "\n",
        "                    with torch.no_grad():\n",
        "                        # noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)[\"sample\"]\n",
        "                        # noise_pred = unet(latent_model_input, torch.tensor(t, dtype=torch.float32).cuda().to(torch.bfloat16), text_embeddings)#[\"sample\"]\n",
        "                        if classifier_guidance: unet.cuda()\n",
        "                        if traced and model_type!=\"compvis\":# and unet_path!=None:\n",
        "                            noise_pred = unet(latent_model_input, torch.tensor(t, dtype=torch.float32).cuda(), text_embeddings)#[\"sample\"]\n",
        "                        else:\n",
        "                            noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings)[\"sample\"]\n",
        "                        if classifier_guidance: unet.cpu()\n",
        "                    # cfg\n",
        "                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
        "                    noise_pred = noise_pred_uncond + scale * (noise_pred_text - noise_pred_uncond)\n",
        "\n",
        "                    # cg\n",
        "                    if classifier_guidance:\n",
        "                        # vae = vae.to(torch_device)\n",
        "                        if vae.device != latents.device:\n",
        "                            vae = vae.to(latents.device)\n",
        "                        latents = latents.detach().requires_grad_()\n",
        "                        latents_x0 = latents - sigma * noise_pred\n",
        "                        denoised_images = vae.decode((1 / 0.18215) * latents_x0).sample / 2 + 0.5\n",
        "                        if args.clip_scale != 0:\n",
        "                            loss = args.loss_fn(denoised_images, \"clip\") * args.clip_scale\n",
        "                        if args.tv_scale != 0:\n",
        "                            loss = args.loss_fn(denoised_images, \"tv\") * args.tv_scale\n",
        "                        if args.lpips_scale != 0:\n",
        "                            loss = 0\n",
        "                            # dude oh my god\n",
        "                            denoised_images = f.resize(denoised_images, (512,512))\n",
        "                            init = f.resize(init, (512,512))\n",
        "                            init_losses = lpips_model(denoised_images, init)\n",
        "                            loss = loss + init_losses.sum() * args.lpips_scale\n",
        "                        cond_grad = -torch.autograd.grad(loss, latents)[0]\n",
        "                        latents = latents.detach() + cond_grad * sigma**2\n",
        "                        # vae = vae.to(offload_device)\n",
        "\n",
        "                    latents = scheduler.step(noise_pred, i, latents)[\"prev_sample\"]\n",
        "\n",
        "        # yaaaaay juggling but guess what it DOESNT WORK!!!!\n",
        "        vae = vae.to(torch_device).to(torch.bfloat16)\n",
        "        unet = unet.to(offload_device)\n",
        "        text_encoder = text_encoder.to(offload_device)\n",
        "\n",
        "        output_image = scale_and_decode(latents.detach().requires_grad_(False).to(torch.bfloat16))\n",
        "\n",
        "        vae = vae.to(offload_device)\n",
        "        unet = unet.to(torch_device)\n",
        "        text_encoder = text_encoder.to(torch_device)\n",
        "        images.append(output_image)\n",
        "\n",
        "        import gc\n",
        "        gc.collect()\n",
        "        torch.cuda.empty_cache()\n",
        "\n",
        "        images[-1].save(f\"progress.png\")\n",
        "    return images\n",
        "\n",
        "torch_device = \"cuda\"\n",
        "offload_device = \"cpu\"\n",
        "\n",
        "do_that = True\n",
        "in_channels = 4 # for later, since the traced version doesn't have this attribute\n",
        "#unet.set_attention_slice(8)\n",
        "unet.cuda()\n",
        "scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", num_train_timesteps=1000)\n",
        "traced = False\n",
        "\n",
        "\n",
        "# idk how people normally do this and i cba to look\n",
        "# prompt = \"By Artgerm and Greg Rutkowski and Alphonse Mucha and hiromumaru, artstation hq, (full body) shot of a ((chinese cultivator)), skimpy ancient Chinese clothing, (((intricate human hands and fingers))), (((pretty face))), (playful eyes) ((nsfw)):1\" #@param {\"type\":\"string\"}\n",
        "import sys \n",
        "prompt = sys.argv[1]\n",
        "\n",
        "bracket_base = 0.0\n",
        "bracket_multiplier = 1.\n",
        "init_img = \"\" \n",
        "size = [sys.argv[2], sys.argv[3]]\n",
        "size = [int(i) for i in size]\n",
        "steps = int(sys.argv[5]) \n",
        "start_step = 0 \n",
        "perlin_multi = 0.72 \n",
        "scale = float(sys.argv[6]) \n",
        "seed = -1 \n",
        "batches = 1\n",
        "\n",
        "# a few \"styles\" from prompts i      stole      from lexica that I know work well, for easy prompt building if you don't have an idea of what to do to improve your prompt\n",
        "prompt_suffix_map = {\n",
        "    \"{artstation}\": \"by ross tran, greg rutkowski, trending on artstation, photograph, hyperreal, octane render, oil on canvas\",\n",
        "    \"{overwatch}\": \"from overwatch, character portrait, close up, concept art, intricate details, highly detailed photorealistic in the style of marco plouffe, keos masons, joel torres, seseon yoon, artgerm and warren louw\",\n",
        "    \"{ghibli}\": \"still from studio ghibli movie; very detailed, focused, colorful, antoine pierre mongin, trending on artstation, 8 k\",\n",
        "    \"{intricate}\": \"4 k resolution, trending on artstation, very very detailed, masterpiece, stunning, intricate\"\n",
        "}\n",
        "def add_suffixes(prompt):\n",
        "    for i in prompt_suffix_map.keys():\n",
        "        prompt = prompt.replace(i,prompt_suffix_map[i])\n",
        "    return prompt\n",
        "prompt = add_suffixes(prompt)\n",
        "\n",
        "\n",
        "def count(string, start=\"(\", end=\")\", negative=True):\n",
        "    temp_string = \"\"\n",
        "    temp_multiplier = bracket_base\n",
        "    mode = \"neutral\"\n",
        "    extension = \"\"\n",
        "    for char in string:\n",
        "        if char == start and mode == \"neutral\":\n",
        "            mode = \"writing\"\n",
        "            temp_multiplier = bracket_base if not negative else -bracket_base\n",
        "        if char == start and mode == \"writing\":\n",
        "            temp_multiplier *= bracket_multiplier\n",
        "        if char == end and mode == \"writing\":\n",
        "            extension += f\" {temp_string}:{str(temp_multiplier)}\"\n",
        "            mode = \"neutral\"\n",
        "            temp_multiplier = bracket_base if not negative else -bracket_base\n",
        "            temp_string = \"\"\n",
        "        if char not in [start, end] and mode == \"writing\":\n",
        "            temp_string += char\n",
        "    for char in [start, end]:\n",
        "        string = string.replace(char, \"\")\n",
        "    return string, extension\n",
        "\n",
        "def add_brackets(prompt):\n",
        "    if \":\" not in prompt[-5:]:\n",
        "        prompt += \":1\"\n",
        "    clean, ext_p = count(prompt, start=\"(\", end=\")\", negative=False)\n",
        "    clean, ext_n = count(clean, start=\"[\", end=\"]\", negative=True)\n",
        "    return prompt + ext_p + ext_n #  make it work more like automatics so the prompts are more cross-compatible\n",
        "\n",
        "prompt = add_brackets(prompt)\n",
        "#prompt = prompt + \" out of frame, bad anatomy, deformed hands, ugly, extra limbs,uneven unnatural eyes, blurry:-0.12\"\n",
        "\n",
        "# classifier_guidance = True\n",
        "# lpips_guidance = True \n",
        "lpips_scale = 0 \n",
        "clip_scale = 0.\n",
        "tv_scale = 0 \n",
        "\n",
        "classifier_guidance = (lpips_scale!=0) or (clip_scale!=0) or (tv_scale!=0)\n",
        "lpips_guidance = lpips_scale!=0\n",
        "\n",
        "\n",
        "class BlankClass():\n",
        "    def __init__(self):\n",
        "        bruh = 'BRUH'\n",
        "args = BlankClass()\n",
        "args.prompt = prompt\n",
        "args.init_img = init_img\n",
        "args.size = size \n",
        "args.steps = steps \n",
        "args.start_step = start_step \n",
        "args.scale = scale\n",
        "args.perlin_multi = perlin_multi\n",
        "args.seed = seed\n",
        "args.batches = batches \n",
        "args.classifier_guidance = classifier_guidance\n",
        "args.lpips_guidance = lpips_guidance\n",
        "args.lpips_scale = lpips_scale\n",
        "# args.loss_scale = clip_scale\n",
        "args.clip_scale = clip_scale\n",
        "args.tv_scale = tv_scale\n",
        "\n",
        "if args.classifier_guidance:\n",
        "    # import clip\n",
        "    import open_clip as clip\n",
        "    from torch import nn\n",
        "    import torch.nn.functional as F\n",
        "    import io\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",
        "        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",
        "            for _ in range(self.cutn):\n",
        "                size = int(torch.rand([])**self.cut_pow * (max_size - min_size) + min_size)\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(F.adaptive_avg_pool2d(cutout, self.cut_size))\n",
        "            return torch.cat(cutouts)\n",
        "    # make_cutouts = MakeCutouts(224, 16)\n",
        "    \n",
        "    clip_text_prompt = \"out of frame, bad anatomy, deformed hands, ugly, extra limbs,uneven unnatural eyes, blurry\" \n",
        "    if clip_scale != 0:\n",
        "        clip_text_prompt = add_suffixes(clip_text_prompt)\n",
        "        clip_text_prompt = add_brackets(clip_text_prompt)\n",
        "\n",
        "    clip_image_prompt = \"\"\n",
        "\n",
        "    if clip_scale != 0:\n",
        "        # clip_model = clip.load(\"ViT-B/32\", jit=False)[0].eval().requires_grad_(False).to(torch_device)\n",
        "        clip_model_name = \"ViT-H-14\" \n",
        "        clip_model_pretrained = \"laion2b_s32b_b79k\" \n",
        "        clip_model, _, preprocess = clip.create_model_and_transforms(clip_model_name, pretrained=clip_model_pretrained)\n",
        "        clip_model = clip_model.eval().requires_grad_(False).to(torch_device)\n",
        "\n",
        "        cutn = 2 \n",
        "        make_cutouts = MakeCutouts(clip_model.visual.image_size if type(clip_model.visual.image_size)!= tuple else clip_model.visual.image_size[0], cutn)\n",
        "\n",
        "    target = None\n",
        "    if len(clip_text_prompt) > 1 and clip_scale != 0:\n",
        "        clip_text_prompt, clip_text_weights = split_weighted_subprompts(clip_text_prompt)\n",
        "        target = clip_model.encode_text(clip.tokenize(clip_text_prompt).to(torch_device)) * torch.tensor(clip_text_weights).view(len(clip_text_prompt), 1).to(torch_device)\n",
        "    if len(clip_image_prompt) > 1 and clip_scale != 0:\n",
        "        clip_image_prompt, clip_image_weights = split_weighted_subprompts(clip_image_prompt, split=\"|\")\n",
        "        # pesky spaces\n",
        "        clip_image_prompt = [p.replace(\" \", \"\") for p in clip_image_prompt]\n",
        "        images = [fetch(image) for image in clip_image_prompt]\n",
        "        images = [f.to_tensor(i).unsqueeze(0) for i in images]\n",
        "        images = [make_cutouts(i) for i in images]\n",
        "        encodings = [clip_model.encode_image(i.to(torch_device)).mean(0) for i in images]\n",
        "        \n",
        "        for i in range(len(encodings)):\n",
        "            encodings[i] = (encodings[i] * clip_image_weights[i]).unsqueeze(0)\n",
        "        # print(encodings.shape)\n",
        "        encodings = torch.cat(encodings, 0)\n",
        "        encoding = encodings.sum(0)\n",
        "\n",
        "        if target!=None:\n",
        "            target = target + encoding\n",
        "        else:\n",
        "            target = encoding\n",
        "        target = target.to(torch.bfloat16).to(torch_device)\n",
        "\n",
        "    # free a little memory, we dont use the text encoder after this so just delete it\n",
        "    if clip_scale != 0:\n",
        "        clip_model.transformer = None\n",
        "    import gc\n",
        "    gc.collect()\n",
        "    torch.cuda.empty_cache()\n",
        "    def spherical_distance(x, y):\n",
        "        x = F.normalize(x, dim=-1)\n",
        "        y = F.normalize(y, dim=-1)\n",
        "        l = (x - y).norm(dim=-1).div(2).arcsin().pow(2).mul(2).mean()\n",
        "        return l \n",
        "    def tv_loss(input):\n",
        "        input = F.pad(input, (0, 1, 0, 1), 'replicate')\n",
        "        return ((input[..., :-1, 1:] - input[..., :-1, :-1])**2 + (input[..., 1:, :-1] - input[..., :-1, :-1])**2).mean()\n",
        "    def loss_fn(x,mode):\n",
        "        # crappy way of handling it, i know\n",
        "        if mode==\"clip\":\n",
        "            with torch.autocast(\"cuda\"):\n",
        "                cutouts = make_cutouts(x)\n",
        "                encoding = clip_model.encode_image(cutouts.float()).to(torch.bfloat16)\n",
        "                loss = spherical_distance(encoding, target)\n",
        "                return loss.mean()\n",
        "        if mode==\"tv\":\n",
        "            return tv_loss(x).mean()\n",
        "\n",
        "    args.loss_fn = loss_fn\n",
        "notify_me_on_every_image = True\n",
        "args.notif = notify_me_on_every_image\n",
        "dtype = torch.float16\n",
        "\n",
        "try:\n",
        "    with torch.amp.autocast(device_type=torch_device, dtype=dtype):\n",
        "        output = sample(args)\n",
        "except KeyboardInterrupt:\n",
        "    print('Interrupting generation..')\n",
        "else:\n",
        "    print('No errors caught!')\n",
        "\n",
        "print(\"Done!\")\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "id": "pJieA9bTt9tN",
        "outputId": "c1c0c0af-a2cf-4351-aa52-8b2869a6b3c2"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Overwriting doohickey-slim.py\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!wget -O danbooru_tags.txt https://gist.githubusercontent.com/aicrumb/000f801df0c30398471cdceff57724bd/raw/920d1108c8f500938d895f43f800bc268c3a5910/danbooru_tags.txt\n",
        "!wget -O danbooru_2.txt https://gist.githubusercontent.com/aicrumb/78fb95439e3b3ccd043fba628de3454f/raw/629b00f55a12bd1fcb0a963c2134c2998a68e803/danbooru_10-20.txt\n",
        "!wget -O artists.txt https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/artists.txt\n",
        "!wget -O mediums.txt https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/mediums.txt\n",
        "!wget -O flavors.txt https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/flavors.txt\n",
        "!wget -O movements.txt https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/movements.txt"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BIzDt3QErBie",
        "outputId": "f49ad7fc-c51a-478f-c680-5b698abba4b8"
      },
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "--2022-09-24 19:02:11--  https://gist.githubusercontent.com/aicrumb/000f801df0c30398471cdceff57724bd/raw/920d1108c8f500938d895f43f800bc268c3a5910/danbooru_tags.txt\n",
            "Resolving gist.githubusercontent.com (gist.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to gist.githubusercontent.com (gist.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1980 (1.9K) [text/plain]\n",
            "Saving to: ‘danbooru_tags.txt’\n",
            "\n",
            "danbooru_tags.txt   100%[===================>]   1.93K  --.-KB/s    in 0s      \n",
            "\n",
            "2022-09-24 19:02:11 (28.1 MB/s) - ‘danbooru_tags.txt’ saved [1980/1980]\n",
            "\n",
            "--2022-09-24 19:02:12--  https://gist.githubusercontent.com/aicrumb/78fb95439e3b3ccd043fba628de3454f/raw/629b00f55a12bd1fcb0a963c2134c2998a68e803/danbooru_10-20.txt\n",
            "Resolving gist.githubusercontent.com (gist.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to gist.githubusercontent.com (gist.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2168 (2.1K) [text/plain]\n",
            "Saving to: ‘danbooru_2.txt’\n",
            "\n",
            "danbooru_2.txt      100%[===================>]   2.12K  --.-KB/s    in 0s      \n",
            "\n",
            "2022-09-24 19:02:12 (33.2 MB/s) - ‘danbooru_2.txt’ saved [2168/2168]\n",
            "\n",
            "--2022-09-24 19:02:12--  https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/artists.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 81865 (80K) [text/plain]\n",
            "Saving to: ‘artists.txt’\n",
            "\n",
            "artists.txt         100%[===================>]  79.95K  --.-KB/s    in 0.01s   \n",
            "\n",
            "2022-09-24 19:02:12 (6.03 MB/s) - ‘artists.txt’ saved [81865/81865]\n",
            "\n",
            "--2022-09-24 19:02:12--  https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/mediums.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.110.133, 185.199.111.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 1606 (1.6K) [text/plain]\n",
            "Saving to: ‘mediums.txt’\n",
            "\n",
            "mediums.txt         100%[===================>]   1.57K  --.-KB/s    in 0s      \n",
            "\n",
            "2022-09-24 19:02:12 (25.1 MB/s) - ‘mediums.txt’ saved [1606/1606]\n",
            "\n",
            "--2022-09-24 19:02:12--  https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/flavors.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.108.133, 185.199.109.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 4958 (4.8K) [text/plain]\n",
            "Saving to: ‘flavors.txt’\n",
            "\n",
            "flavors.txt         100%[===================>]   4.84K  --.-KB/s    in 0s      \n",
            "\n",
            "2022-09-24 19:02:12 (58.0 MB/s) - ‘flavors.txt’ saved [4958/4958]\n",
            "\n",
            "--2022-09-24 19:02:12--  https://raw.githubusercontent.com/pharmapsychotic/clip-interrogator/main/data/movements.txt\n",
            "Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.108.133, 185.199.109.133, 185.199.110.133, ...\n",
            "Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.108.133|:443... connected.\n",
            "HTTP request sent, awaiting response... 200 OK\n",
            "Length: 2683 (2.6K) [text/plain]\n",
            "Saving to: ‘movements.txt’\n",
            "\n",
            "movements.txt       100%[===================>]   2.62K  --.-KB/s    in 0s      \n",
            "\n",
            "2022-09-24 19:02:13 (51.1 MB/s) - ‘movements.txt’ saved [2683/2683]\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### prompt engine"
      ],
      "metadata": {
        "id": "oXGiSjNgrowU"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "cellView": "form",
        "id": "1dD4B6kghPTq",
        "outputId": "e1ae2fd7-2c76-4d16-e7d4-95af491b41c0"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Current Generation 0\n",
            "====================\n",
            "['Stephen Gilbert Henri Alphonse Barnoin Kahlo Austin Briggs Rowena Meeks Abdy Cornelis Anthonisz Alfons Karpiński Julia Margaret Cameron '] [1.0]\n",
            "Generating with seed 7009...\n",
            "  0% 0/1 [00:00<?, ?it/s]"
          ]
        }
      ],
      "source": [
        "#@markdown code for genetic algorithm *mostly* generated by code-davinci-002 <br>\n",
        "#@markdown comes pre-loaded with the first 20 pages of the most popular danbooru tags (https://danbooru.donmai.us/tags?commit=Search&page=1&search%5Bhide_empty%5D=yes&search%5Border%5D=count) <br>\n",
        "#@markdown and the word lists here: https://github.com/pharmapsychotic/clip-interrogator/tree/main/data <br>\n",
        "#@markdown **the first generation will take a while, as it's downloading the model** <br><br>\n",
        "#@markdown use any scoring system you want! 1-10, 1-3, whatever as long as the higher number is favorable\n",
        "import random\n",
        "from IPython.display import display, Image\n",
        "from IPython.display import clear_output\n",
        "from google.colab import output\n",
        "# wordlist = input(\"words:\").split(\" \")\n",
        "# wordlist = \"1girl solo highres long_hair commentary_request breasts looking_at_viewer blush smile short_hair open_mouth bangs blue_eyes multiple_girls blonde_hair skirt brown_hair large_breasts simple_background black_hair shirt hair_ornament red_eyes thighhighs absurdres hat gloves 1boy bad_id long_sleeves white_background dress original ribbon touhou bow navel 2girls bad_pixiv_id photoshop_(medium) holding animal-ears cleavage hair_between_eyes brown_eyes bare_shoulders twintails medium_breasts commentary jewelry sitting very_long_hair underwear closed_mouth nipples school_uniform green_eyes blue_hair standing purple_eyes collarbone panties monochrome tail jacket translated swimsuit full_body closed_eyes hair_ribbon kantai_collection yellow_eyes weapon ponytail purple_hair upper_body ass pink_hair comic white_shirt braid flower ahoge short_sleeves greyscale hair_bow hetero male_focus heart pantyhose bikini white_hair sidelocks nude thighs red_hair cowboy_shot pleated_skirt sweat translation_request hairband small_breasts earrings boots multicolored_hair lying censored frills parted_lips detached_sleeves one_eye_closed outdoors food japanese_clothes multiple_boys green_hair wings open_clothes sky necktie horns shoes penis fate_(series) grey_hair glasses barefoot shorts serafuku silver_hair pussy teeth day solo_focus sleeveless choker alternate_costume tongue pointy_ears socks black_gloves elbow_gloves hairclip fang striped midriff puffy_sleeves shiny looking_back belt sword official_art collared_shirt pants cloud artist_name black_thighhighs tears fate/grand_order cat_ears indoors white_gloves 3girls hair_flower signature virtual_youtuber dark_skin hand_up spread_legs cum 2boys idolmaster hood sex miniskirt wide_sleeves tongue_out fingerless_gloves on_back blunt_bangs black_skirt bowtie armpits pink_eyes sailor_collar black_legwear kimono english_commentary pokemon medium_hair water grey_background necklace chibi off_shoulder bag clothes_lift hair_bun scarf\"#@param {\"type\":\"string\"}\n",
        "# wordlist = wordlist.split(\" \")\n",
        "\n",
        "diffusers_model = \"doohickey/trinart-waifu-diffusion-50-50\" #@param {type:\"string\"}\n",
        "steps = 45 #@param\n",
        "scale = 8 #@param\n",
        "height = 640 #@param \n",
        "width = 448 #@param\n",
        "use_danbooru_tags = False #@param {type:\"boolean\"}\n",
        "use_interrogator_concepts = True #@param {type:\"boolean\"}\n",
        "# keep_top_fraction = 0.5 #@param\n",
        "mutation_rate = 0.1 #@param\n",
        "crossover_rate = 0.25 #@param\n",
        "wordlist = []\n",
        "if use_danbooru_tags == True:\n",
        "    wordlist += open(\"danbooru_tags.txt\", \"r\").read().split(\" \")\n",
        "    wordlist += open(\"danbooru_2.txt\", \"r\").read().split(\" \")\n",
        "    if use_interrogator_concepts:\n",
        "        wordlist = wordlist * 25 # the other ones are so long it gets under-represented\n",
        "if use_interrogator_concepts == True:\n",
        "    wordlist += [i.replace(\"\\n\",\"\") for i in open(\"artists.txt\", \"r\").readlines()]\n",
        "    wordlist += [i.replace(\"\\n\",\"\") for i in open(\"movements.txt\", \"r\").readlines()]\n",
        "    wordlist += [i.replace(\"\\n\",\"\") for i in open(\"flavors.txt\", \"r\").readlines()]\n",
        "    wordlist += [i.replace(\"\\n\",\"\") for i in open(\"mediums.txt\", \"r\").readlines()]\n",
        "\n",
        "base = \"\" #@param {type:\"string\"}\n",
        "prompt_length = 8 #@param\n",
        "population = 4 #@param\n",
        "clear_after_each_generation = True #@param {type:\"boolean\"}\n",
        "\n",
        "class Population:\n",
        "    def __init__(self, size):\n",
        "        self.size = size\n",
        "        self.members = []\n",
        "        for i in range(size):\n",
        "            self.members.append(Member())\n",
        "        self.fitness = []\n",
        "        self.best = None\n",
        "        self.best_fitness = 0\n",
        "        self.avg_fitness = 0\n",
        "        self.generation = 0\n",
        "        self.mutation_rate = mutation_rate\n",
        "        self.crossover_rate = crossover_rate\n",
        "        # self.top_frac = keep_top_fraction\n",
        "    def get_fitness(self):\n",
        "        if clear_after_each_generation:\n",
        "            clear_output(wait=False)\n",
        "        print(\"Current Generation\", self.generation)\n",
        "        # print string and use input to get fitness from user\n",
        "        for member in self.members:\n",
        "            current_prompt = ((base + \" \" + member.string) if len(base) > 1 else member.string)\n",
        "            print(\"=\"*20)\n",
        "\n",
        "            # im re-loading the model every time because i want NO chance of a memory leak, \n",
        "            # imagine you sit here for an hour and then BOOM cuda oom? that would SUCK\n",
        "            t_0 = time.time()\n",
        "            !python doohickey-slim.py \"$current_prompt\" $height $width \"$diffusers_model\" $steps $scale\n",
        "            t_1 = time.time()\n",
        "            output.eval_js('new Audio(\"https://freesound.org/data/previews/80/80921_1022651-lq.ogg\").play()')\n",
        "            print(f\"Took {t_1-t_0} seconds.\")\n",
        "            display(Image(\"progress.png\"))\n",
        "            print(\"-\"*20)\n",
        "            # member.fitness = float(input(\"fitness:\"))\n",
        "            a = %sx read -p \"Score:\"\n",
        "            a = a[0].replace(\"Score:\",\"\").replace(\" \", \"\")\n",
        "            if len(a)<1 or a==\" \":\n",
        "                a = self.avg_fitness\n",
        "            member.fitness = float(a) + random.random() * 0.01 # so that if you have all the same number it wont error\n",
        "            print(\"Your score:\", member.fitness)\n",
        "\n",
        "            if member.fitness > self.best_fitness:\n",
        "                self.best_fitness = member.fitness\n",
        "                self.best = member\n",
        "            self.fitness.append(member.fitness)\n",
        "        self.avg_fitness = sum(self.fitness)/len(self.fitness)\n",
        "    def next_generation(self):\n",
        "        # create new generation\n",
        "        self.generation += 1\n",
        "        new_members = []\n",
        "        print(sorted(self.fitness))\n",
        "        self.members = list(reversed([x for _, x in sorted(zip(self.fitness, self.members), key=lambda pair: pair[0])]))\n",
        "        self.members = self.members[:round(len(self.members)*0.5)] * 2\n",
        "        print(self.members)\n",
        "        for i in range(self.size):\n",
        "            # select parents\n",
        "            parent1 = self.members[random.randint(0,len(self.members)-1)]\n",
        "            parent2 = self.members[random.randint(0,len(self.members)-1)]\n",
        "            # crossover\n",
        "            if random.random() < self.crossover_rate:\n",
        "                child = parent1.crossover(parent2)\n",
        "            else:\n",
        "                child = parent1\n",
        "            # mutate\n",
        "            if random.random() < self.mutation_rate:\n",
        "                child.mutate()\n",
        "            new_members.append(child)\n",
        "        self.members = new_members\n",
        "        self.fitness = []\n",
        "\n",
        "class Member:\n",
        "    def __init__(self):\n",
        "        self.string = \"\"\n",
        "        for i in range(prompt_length):\n",
        "            self.string += wordlist[random.randint(0, len(wordlist)-1)] + \" \"\n",
        "        self.fitness = 0\n",
        "    def mutate(self):\n",
        "        # mutate string\n",
        "        self.string = \"\"\n",
        "        for i in range(prompt_length):\n",
        "            self.string += wordlist[random.randint(0, len(wordlist)-1)] + \" \"\n",
        "    def crossover(self, other):\n",
        "        # crossover strings\n",
        "        child = Member()\n",
        "        child.string = \"\"\n",
        "        for i in range(prompt_length):\n",
        "            if random.random() < 0.5:\n",
        "                child.string += self.string.split(\" \")[i] + \" \"\n",
        "            else:\n",
        "                child.string += other.string.split(\" \")[i] + \" \"\n",
        "        return child\n",
        "\n",
        "population = Population(population)\n",
        "\n",
        "while True:\n",
        "    population.get_fitness()\n",
        "    print(\"generation:\", population.generation)\n",
        "    print(\"best fitness:\", population.best_fitness)\n",
        "    print(\"best member:\", population.best.string)\n",
        "    print(\"average fitness:\", population.avg_fitness)\n",
        "    population.next_generation()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### tip:\n",
        "it's good to have a specific scoring criteria, for example:\n",
        "```\n",
        "0 - no waterbottle\n",
        "1 - waterbottle in picture\n",
        "2 - waterbottle is main focus\n",
        "3 - waterbottle is main focus + image is aesthetically pleasing\n",
        "```\n",
        "\n",
        "```\n",
        "-1 - dislike\n",
        "0 - neutral\n",
        "1 - like\n",
        "```\n",
        "\n",
        "another thing if you aren't sure of your scoring criteria, is to set the first image of each generation at the middle of your range (5 if 1-10) and base the rest from there."
      ],
      "metadata": {
        "id": "fZptSmg2-gV3"
      }
    }
  ]
}