008_text_to_world.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/karaage0703/stable-diffusion-colab-tools/blob/main/008_text_to_world.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5dcb2f23-7bcb-42d6-ae00-c336222dcd93",
      "metadata": {
        "id": "5dcb2f23-7bcb-42d6-ae00-c336222dcd93"
      },
      "source": [
        "# Text-to-World\n",
        "\n",
        "Text-to-World by using GA(Genetic Algorithm) like algorithm and Stable Diffusion\n",
        "\n",
        "Reference notebook:  \n",
        "https://github.com/fastai/diffusion-nbs/blob/master/Stable%20Diffusion%20Deep%20Dive.ipynb\n",
        "\n",
        "About license of this notebook refer to reference notebook. "
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4435401f-aa50-442b-91bd-bbdc17c81f5c",
      "metadata": {
        "id": "4435401f-aa50-442b-91bd-bbdc17c81f5c"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8a267a12-2bfa-4685-a5b2-ff1bc7771853",
      "metadata": {
        "id": "8a267a12-2bfa-4685-a5b2-ff1bc7771853"
      },
      "outputs": [],
      "source": [
        "!pip install -qq --upgrade transformers diffusers ftfy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "650e1224-d757-45e9-a7cc-67d608bd73bd",
      "metadata": {
        "id": "650e1224-d757-45e9-a7cc-67d608bd73bd"
      },
      "outputs": [],
      "source": [
        "from base64 import b64encode\n",
        "\n",
        "import numpy as np\n",
        "import random\n",
        "import torch\n",
        "from diffusers import AutoencoderKL, LMSDiscreteScheduler, UNet2DConditionModel\n",
        "from huggingface_hub import notebook_login\n",
        "\n",
        "# For video display:\n",
        "from IPython.display import HTML\n",
        "from matplotlib import pyplot as plt\n",
        "from PIL import Image\n",
        "from torch import autocast\n",
        "from torchvision import transforms as tfms\n",
        "from tqdm.auto import tqdm\n",
        "from transformers import CLIPTextModel, CLIPTokenizer, logging\n",
        "\n",
        "notebook_login()\n",
        "\n",
        "# Supress some unnecessary warnings when loading the CLIPTextModel\n",
        "logging.set_verbosity_error()\n",
        "\n",
        "# Set device\n",
        "torch_device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
      ]
    },
    {
      "cell_type": "markdown",
      "id": "c1e4e75b-21fc-496f-8b0c-371353b145f4",
      "metadata": {
        "id": "c1e4e75b-21fc-496f-8b0c-371353b145f4"
      },
      "source": [
        "## Loading the models\n",
        "\n",
        "This code (and that in the next section) comes from the [Huggingface example notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion.ipynb). \n",
        "\n",
        "This will download and set up the relevant models and components we'll be using. Let's just run this for now and move on to the next section to check that it all works before diving deeper.\n",
        "\n",
        "If you've loaded a pipeline, you can also access these components using `pipe.unet`, `pipe.vae` and so on.\n",
        "\n",
        "In this notebook we aren't doing any memory-saving tricks - if you find yourself running out of GPU RAM, look at the pipeline code for inspiration with things like attention slicing, switching to half precision (fp16), keeping the VAE on the CPU and other modifications."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "eb58f35f-dde0-428f-92d8-606fd3c0676d",
      "metadata": {
        "id": "eb58f35f-dde0-428f-92d8-606fd3c0676d"
      },
      "outputs": [],
      "source": [
        "# Load the autoencoder model which will be used to decode the latents into image space. \n",
        "vae = AutoencoderKL.from_pretrained(\"CompVis/stable-diffusion-v1-4\", subfolder=\"vae\")\n",
        "\n",
        "# Load the tokenizer and text encoder to tokenize and encode the text. \n",
        "tokenizer = CLIPTokenizer.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
        "text_encoder = CLIPTextModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
        "\n",
        "# The UNet model for generating the latents.\n",
        "unet = UNet2DConditionModel.from_pretrained(\"CompVis/stable-diffusion-v1-4\", subfolder=\"unet\")\n",
        "\n",
        "# The noise scheduler\n",
        "scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", num_train_timesteps=1000)\n",
        "\n",
        "# To the GPU we go!\n",
        "vae = vae.to(torch_device)\n",
        "text_encoder = text_encoder.to(torch_device)\n",
        "unet = unet.to(torch_device);"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4e6c78e1-f013-4926-9022-c0b3916b5bde",
      "metadata": {
        "id": "4e6c78e1-f013-4926-9022-c0b3916b5bde"
      },
      "source": [
        "## Text to Image"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "height = 512                        # default height of Stable Diffusion\n",
        "width = 512                         # default width of Stable Diffusion\n",
        "num_inference_steps = 20            # Number of denoising steps\n",
        "guidance_scale = 7.5                # Scale for classifier-free guidance\n",
        "batch_size = 1"
      ],
      "metadata": {
        "id": "GSjPEIasz01a"
      },
      "id": "GSjPEIasz01a",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def text_emb_to_image(text_embeddings, seed):\n",
        "    uncond_input = tokenizer(\n",
        "        [\"\"] * batch_size, padding=\"max_length\", max_length=77, return_tensors=\"pt\"\n",
        "    )\n",
        "    with torch.no_grad():\n",
        "        uncond_embeddings = text_encoder(uncond_input.input_ids.to(torch_device))[0]\n",
        "    text_embeddings = torch.cat([uncond_embeddings, text_embeddings])\n",
        "\n",
        "    # Prep Scheduler\n",
        "    scheduler.set_timesteps(num_inference_steps)\n",
        "\n",
        "    generator = torch.manual_seed(seed) \n",
        "    # Prep latents\n",
        "    latents = torch.randn(\n",
        "      (batch_size, unet.in_channels, height // 8, width // 8),\n",
        "      generator=generator,\n",
        "    )\n",
        "    latents = latents.to(torch_device)\n",
        "    latents = latents * scheduler.init_noise_sigma # Scaling (previous versions did latents = latents * self.scheduler.sigmas[0]\n",
        "\n",
        "    # Loop\n",
        "    with autocast(\"cuda\"):\n",
        "        for i, t in tqdm(enumerate(scheduler.timesteps)):\n",
        "            # expand the latents if we are doing classifier-free guidance to avoid doing two forward passes.\n",
        "            latent_model_input = torch.cat([latents] * 2)\n",
        "            sigma = scheduler.sigmas[i]\n",
        "            # Scale the latents (preconditioning):\n",
        "            # latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5) # Diffusers 0.3 and below\n",
        "            latent_model_input = scheduler.scale_model_input(latent_model_input, t)\n",
        "\n",
        "            # predict the noise residual\n",
        "            with torch.no_grad():\n",
        "                noise_pred = unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample\n",
        "\n",
        "            # perform guidance\n",
        "            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
        "            noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)\n",
        "\n",
        "            # compute the previous noisy sample x_t -> x_t-1\n",
        "            # latents = scheduler.step(noise_pred, i, latents)[\"prev_sample\"] # Diffusers 0.3 and below\n",
        "            latents = scheduler.step(noise_pred, t, latents).prev_sample\n",
        "\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\n",
        "\n",
        "    # Display\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[0]"
      ],
      "metadata": {
        "id": "Vsw2lNrAnLaq"
      },
      "id": "Vsw2lNrAnLaq",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def prompt_to_text_emb(prompt):\n",
        "    text_input = tokenizer(prompt, padding=\"max_length\", max_length=tokenizer.model_max_length, truncation=True, return_tensors=\"pt\")\n",
        "    with torch.no_grad():\n",
        "        text_embeddings = text_encoder(text_input.input_ids.to(torch_device))[0]\n",
        "    max_length = text_input.input_ids.shape[-1]\n",
        "    uncond_input = tokenizer(\n",
        "        [\"\"] * batch_size, padding=\"max_length\", max_length=max_length, return_tensors=\"pt\"\n",
        "    )\n",
        "\n",
        "    return text_embeddings"
      ],
      "metadata": {
        "id": "zhTo1PPOnusr"
      },
      "id": "zhTo1PPOnusr",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Gene"
      ],
      "metadata": {
        "id": "_zR8TmfhzOl5"
      },
      "id": "_zR8TmfhzOl5"
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setup\n",
        "Create Gene class"
      ],
      "metadata": {
        "id": "I6VcCjiLl2c7"
      },
      "id": "I6VcCjiLl2c7"
    },
    {
      "cell_type": "code",
      "source": [
        "class Gene:\n",
        "    def __init__(self, text_emb, seed):\n",
        "        self.text_embeddings = text_emb\n",
        "        self.seed = seed\n",
        "        self.image = None\n",
        "\n",
        "    def generate_image(self):\n",
        "        self.image = text_emb_to_image(self.text_embeddings.to(torch_device), self.seed)\n",
        "\n",
        "    def get_image(self):\n",
        "        if self.image == None:\n",
        "            self.generate_image()\n",
        "\n",
        "        return self.image"
      ],
      "metadata": {
        "id": "oYwAhH5UGR3O"
      },
      "id": "oYwAhH5UGR3O",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def cross_over(targets):\n",
        "    samples = random.sample(targets, 2)\n",
        "    cross_over_point = random.randint(22, 55)\n",
        "    sample_0 = samples[0].text_embeddings.cpu().detach().numpy().copy()\n",
        "    sample_1 = samples[1].text_embeddings.cpu().detach().numpy().copy()\n",
        "\n",
        "    sample_child = np.concatenate([sample_0[0][0:cross_over_point], sample_1[0][cross_over_point:77]])\n",
        "    sample_child = np.expand_dims(sample_child, axis = 0)\n",
        "    sample_child = torch.from_numpy(sample_child.astype(np.float32)).clone()\n",
        "\n",
        "    prob = random.random()\n",
        "    if prob < 0.2:\n",
        "        seed = random.randint(0, 1000)\n",
        "    else:\n",
        "        seed = samples[0].seed\n",
        "\n",
        "    return Gene(sample_child, seed)"
      ],
      "metadata": {
        "id": "sqsV4_1tzrB4"
      },
      "id": "sqsV4_1tzrB4",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def display_genes(targets):\n",
        "    plt.figure(figsize=(16, 9))\n",
        "    plt.subplots_adjust(hspace=0.5)\n",
        "\n",
        "    for i, target in enumerate(targets):\n",
        "        colum_numb = 5\n",
        "        plt.subplot(int(len(targets) / colum_numb) + 1, colum_numb, i + 1)\n",
        "\n",
        "        plt.imshow(target.get_image())\n",
        "        plt.title(i)\n",
        "        plt.axis('off')\n",
        "\n",
        "    _ = plt.suptitle('display gene images')"
      ],
      "metadata": {
        "id": "w8W3W_ne4f0t"
      },
      "id": "w8W3W_ne4f0t",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Initialize"
      ],
      "metadata": {
        "id": "636di_R7l--R"
      },
      "id": "636di_R7l--R"
    },
    {
      "cell_type": "markdown",
      "source": [
        "Initialize gene"
      ],
      "metadata": {
        "id": "BSxs7YrandEi"
      },
      "id": "BSxs7YrandEi"
    },
    {
      "cell_type": "code",
      "source": [
        "genes = []"
      ],
      "metadata": {
        "id": "6BWjR3WAIoKr"
      },
      "id": "6BWjR3WAIoKr",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "genes.append(Gene(prompt_to_text_emb('adam'), random.randint(0, 1000)))\n",
        "genes.append(Gene(prompt_to_text_emb('eve'), random.randint(0, 1000)))"
      ],
      "metadata": {
        "id": "TRTQ4Xt0zOHJ"
      },
      "id": "TRTQ4Xt0zOHJ",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "display_genes(genes)"
      ],
      "metadata": {
        "id": "crORAJgsn5V4"
      },
      "id": "crORAJgsn5V4",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## breed gene"
      ],
      "metadata": {
        "id": "0-2g0TWZn-W3"
      },
      "id": "0-2g0TWZn-W3"
    },
    {
      "cell_type": "code",
      "source": [
        "genes.append(cross_over(genes))"
      ],
      "metadata": {
        "id": "nHBJxXsGoIuN"
      },
      "id": "nHBJxXsGoIuN",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "display_genes(genes)"
      ],
      "metadata": {
        "id": "8P81FZEzoadV"
      },
      "id": "8P81FZEzoadV",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "add gene"
      ],
      "metadata": {
        "id": "hT4Snd60p6Uw"
      },
      "id": "hT4Snd60p6Uw"
    },
    {
      "cell_type": "code",
      "source": [
        "genes.append(Gene(prompt_to_text_emb('light'), random.randint(0, 1000)))\n",
        "genes.append(Gene(prompt_to_text_emb('fire'), random.randint(0, 1000)))\n",
        "genes.append(Gene(prompt_to_text_emb('rock'), random.randint(0, 1000)))\n",
        "genes.append(Gene(prompt_to_text_emb('apple'), random.randint(0, 1000)))\n",
        "genes.append(Gene(prompt_to_text_emb('water'), random.randint(0, 1000)))"
      ],
      "metadata": {
        "id": "88rvSV1GMS4c"
      },
      "id": "88rvSV1GMS4c",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "set cross over trial number"
      ],
      "metadata": {
        "id": "jM2ucBvr22Pj"
      },
      "id": "jM2ucBvr22Pj"
    },
    {
      "cell_type": "code",
      "source": [
        "trial_numb = 5"
      ],
      "metadata": {
        "id": "vxZ0S9Vq21SX"
      },
      "id": "vxZ0S9Vq21SX",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "cross over"
      ],
      "metadata": {
        "id": "8lfj2Li1pGoV"
      },
      "id": "8lfj2Li1pGoV"
    },
    {
      "cell_type": "code",
      "source": [
        "for i in range(trial_numb):\n",
        "    genes.append(cross_over(genes))"
      ],
      "metadata": {
        "id": "BwrPzlj12nxC"
      },
      "id": "BwrPzlj12nxC",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "display_genes(genes)"
      ],
      "metadata": {
        "id": "fPxwvVFU2s62"
      },
      "id": "fPxwvVFU2s62",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "delete gene"
      ],
      "metadata": {
        "id": "q_o3msjH3A4r"
      },
      "id": "q_o3msjH3A4r"
    },
    {
      "cell_type": "code",
      "source": [
        "del genes[1:2]"
      ],
      "metadata": {
        "id": "Tdrj4FgEqCGc"
      },
      "id": "Tdrj4FgEqCGc",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "display_genes(genes)"
      ],
      "metadata": {
        "id": "sYwNJXv-pi8Y"
      },
      "id": "sYwNJXv-pi8Y",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "check gene"
      ],
      "metadata": {
        "id": "ax2F2V-X3E1X"
      },
      "id": "ax2F2V-X3E1X"
    },
    {
      "cell_type": "code",
      "source": [
        "gene_numb = 17"
      ],
      "metadata": {
        "id": "xLJIohPd3uUI"
      },
      "id": "xLJIohPd3uUI",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(genes[gene_numb].seed)\n",
        "print(genes[gene_numb].text_embeddings)"
      ],
      "metadata": {
        "id": "-vuTvD1_3dF_"
      },
      "id": "-vuTvD1_3dF_",
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "plt.imshow(genes[gene_numb].get_image())"
      ],
      "metadata": {
        "id": "k-b1YPe33FlF"
      },
      "id": "k-b1YPe33FlF",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.4"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "nbformat": 4,
  "nbformat_minor": 5
}