imaginación-mecanizada-ej3.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3e2O4uXnhTya"
      },
      "source": [
        "# El arte de la imaginación mecanizada - Ejercicio 3\n",
        "\n",
        "#### Prácticas y pensamiento artísticos en torno al giro tecnológico y a Inteligencias Artificiales, Bilbo, 11 de julio de 2023\n",
        "\n",
        "Basado muy fuertemente en [Grokking Stable Diffusion](https://colab.research.google.com/drive/1dlgggNa5Mz8sEAGU0wFCHhGLFooW_pf1?usp=sharing) de [Jonathan Whitaker](https://github.com/johnowhitaker) y en [Stable Diffusion with 🧨 diffusers](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion.ipynb) de [🤗 Hugging Face](https://github.com/huggingface/diffusers)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "NLBlDaLyUKDe"
      },
      "outputs": [],
      "source": [
        "#@markdown #Comprobar GPU\n",
        "!nvidia-smi -L"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9vzfOZ27UsTm"
      },
      "outputs": [],
      "source": [
        "#@markdown #Librerías + Modelos + Funciones de ayuda\n",
        "\n",
        "from IPython.display import clear_output\n",
        "\n",
        "!pip install diffusers==0.3.0\n",
        "!pip install transformers scipy ftfy\n",
        "\n",
        "from google.colab import files\n",
        "from transformers import CLIPTextModel, CLIPTokenizer\n",
        "from tqdm.auto import tqdm\n",
        "from torch import autocast\n",
        "from PIL import Image\n",
        "from huggingface_hub import notebook_login\n",
        "from diffusers import AutoencoderKL, UNet2DConditionModel, LMSDiscreteScheduler\n",
        "import torch\n",
        "import re\n",
        "clear_output()\n",
        "\n",
        "\n",
        "vae = AutoencoderKL.from_pretrained(\n",
        "    \"runwayml/stable-diffusion-v1-5\", subfolder=\"vae\")\n",
        "tokenizer = CLIPTokenizer.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
        "text_encoder = CLIPTextModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
        "unet = UNet2DConditionModel.from_pretrained(\n",
        "    \"runwayml/stable-diffusion-v1-5\", subfolder=\"unet\")\n",
        "scheduler = LMSDiscreteScheduler(\n",
        "    beta_start=0.00085, beta_end=0.012, beta_schedule=\"scaled_linear\", num_train_timesteps=1000)\n",
        "clear_output()\n",
        "\n",
        "torch_device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "vae = vae.to(torch_device)\n",
        "text_encoder = text_encoder.to(torch_device)\n",
        "unet = unet.to(torch_device)\n",
        "\n",
        "\n",
        "def latents_to_pil(latents):\n",
        "    latents = (1 / 0.18215) * latents\n",
        "    with torch.no_grad():\n",
        "        images = vae.decode(latents).sample\n",
        "    images = (images / 2 + 0.5).clamp(0, 1)\n",
        "    images = images.detach().cpu().permute(0, 2, 3, 1).numpy()\n",
        "    images = (images * 255).round().astype(\"uint8\")\n",
        "    pil_images = [Image.fromarray(image) for image in images]\n",
        "\n",
        "    return pil_images\n",
        "\n",
        "\n",
        "def image_grid(imgs, cols):\n",
        "    grid_w = min([cols, len(imgs)])\n",
        "    grid_h = len(imgs)//cols + 1\n",
        "    w, h = imgs[0].size\n",
        "    grid = Image.new('RGB', size=(grid_w*w, grid_h*h))\n",
        "    #grid_w, grid_h = grid.size\n",
        "    for i, img in enumerate(imgs):\n",
        "        grid.paste(img, box=(i % cols*w, i//cols*h))\n",
        "\n",
        "    return grid\n",
        "\n",
        "\n",
        "def get_seed_gen(seed):\n",
        "    if seed == 0:\n",
        "      seed = torch.randint(2**32, (1, 1))[0, 0].item()\n",
        "    generator = torch.Generator(device=torch_device)\n",
        "    generator.manual_seed(seed)\n",
        "    return seed, generator\n",
        "\n",
        "\n",
        "def render(steps, scale, generator, embeddings, batch_size):\n",
        "\n",
        "    width = 512\n",
        "    height = 512\n",
        "\n",
        "    scheduler.set_timesteps(steps)\n",
        "\n",
        "\n",
        "    latents = torch.randn(\n",
        "        (batch_size, unet.in_channels, height // 8, width // 8),\n",
        "        generator=generator,\n",
        "        device=torch_device\n",
        "    )\n",
        "    # latents = latents.to(torch_device)  # [batch_size, 4, 64, 64]\n",
        "    latents = latents * scheduler.sigmas[0]\n",
        "\n",
        "    with autocast(\"cuda\"):\n",
        "\n",
        "        for i, t in tqdm(enumerate(scheduler.timesteps), total=steps):\n",
        "\n",
        "            sigma = scheduler.sigmas[i]\n",
        "            latent_model_input = torch.cat(\n",
        "                [latents] * 2)  # [batch_size*2, 4, 64, 64]\n",
        "            latent_model_input = latent_model_input / ((sigma**2 + 1) ** 0.5)\n",
        "\n",
        "            with torch.no_grad():\n",
        "                # [2, 4, 64, 64]\n",
        "                noise_pred = unet(latent_model_input, t,\n",
        "                                  encoder_hidden_states=embeddings).sample\n",
        "\n",
        "            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)\n",
        "            noise_pred = noise_pred_uncond + scale * \\\n",
        "                (noise_pred_text - noise_pred_uncond)  # [1, 4, 64, 64]\n",
        "\n",
        "            latents = scheduler.step(\n",
        "                noise_pred, i, latents).prev_sample  # [1, 4, 64, 64]\n",
        "            # print(f\"i={i} t={t}, sigma={sigma}\")\n",
        "\n",
        "    return latents_to_pil(latents)\n",
        "\n",
        "\n",
        "token_EOS_value = 49407\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "rI0Kp9QgW6WW"
      },
      "outputs": [],
      "source": [
        "#@markdown #1. Hiperscala\n",
        "#@markdown Un único prompt que multiplicamos para hacerlo 'hiper'.\n",
        "\n",
        "num_images = 1\n",
        "\n",
        "prompt = \"dog natgeo\"  # @param {type:\"string\"}\n",
        "semilla = 2100112712  # @param {type:\"number\"}\n",
        "guia = 7.5  # -@param {type:\"slider\", min:4.5, max:21, step:1.5}\n",
        "hiperescala = 1  # @param {type:\"slider\", min:-2, max:2, step:0.05}\n",
        "pasos = 60  # @param {type:\"slider\", min:25, max:150, step:5}\n",
        "\n",
        "\n",
        "prompts = [\"\"] * num_images  # Unconditional\n",
        "prompts.extend([prompt] * num_images)  # Prompts\n",
        "\n",
        "\n",
        "inputs = tokenizer(prompts, padding=\"max_length\", max_length=tokenizer.model_max_length,\n",
        "                   truncation=True, return_tensors=\"pt\")  # [num_images * 2, 77]\n",
        "\n",
        "with torch.no_grad():\n",
        "  text_embeddings = text_encoder(inputs.input_ids.to(torch_device))[\n",
        "      0]  # [num_images * 2, 77, 768]\n",
        "\n",
        "token_list = inputs.input_ids[1].tolist()\n",
        "\n",
        "# Reescalamos los embeddings del prompt\n",
        "text_embeddings[num_images:, 1:token_list.index(\n",
        "    token_EOS_value)] *= hiperescala\n",
        "\n",
        "\n",
        "semilla, generator = get_seed_gen(semilla)\n",
        "images = render(pasos, guia, generator, text_embeddings, num_images)\n",
        "filename = f\"1_{prompt}_{semilla}_{hiperescala:0.02f}_{pasos}.jpg\"\n",
        "\n",
        "print(f\"Semilla: {semilla}\")\n",
        "print(f\"Fichero: {filename}\\n\")\n",
        "\n",
        "grid = image_grid(images, cols=3)\n",
        "grid.save(filename, quality=100, subsampling=0)\n",
        "files.download(filename)\n",
        "grid\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "fB3o7-WZGo5l"
      },
      "outputs": [],
      "source": [
        "#@markdown #2. Híbridos\n",
        "#@markdown Combinación de dos prompts.\n",
        "\n",
        "num_images = 1\n",
        "\n",
        "pasos = 50\n",
        "guia = 7.5\n",
        "\n",
        "prompt1 = \"elephant natgeo\"  # @param {type:\"string\"}\n",
        "prompt2 = \"macaw natgeo\"  # @param {type:\"string\"}\n",
        "semilla = 3032598914  # @param {type:\"number\"}\n",
        "mezcla = 52  # @param {type:\"slider\", min:0, max:100, step:1}\n",
        "mezcla /= 100\n",
        "\n",
        "\n",
        "inputs = tokenizer([\"\", prompt1, prompt2], padding=\"max_length\",\n",
        "                   max_length=tokenizer.model_max_length, truncation=True, return_tensors=\"pt\")  # [3,77]\n",
        "\n",
        "with torch.no_grad():\n",
        "  embeddings = text_encoder(inputs.input_ids.to(torch_device))[0]  # [3,77,768]\n",
        "\n",
        "text_embeddings = torch.stack(\n",
        "    [embeddings[0], embeddings[1]*(1-mezcla) + embeddings[2]*mezcla])  # [2,77,768]\n",
        "\n",
        "\n",
        "semilla, generator = get_seed_gen(semilla)\n",
        "images = render(pasos, guia, generator, text_embeddings, num_images)\n",
        "filename = f\"21_{prompt1}_{prompt2}_{semilla}_{mezcla:0.02f}.jpg\"\n",
        "\n",
        "print(f\"Semilla: {semilla}\")\n",
        "print(f\"Fichero: {filename}\\n\")\n",
        "\n",
        "grid = image_grid(images, cols=3)\n",
        "grid.save(filename, quality=100, subsampling=0)\n",
        "files.download(filename)\n",
        "grid"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ctlMLKRQMMbS"
      },
      "outputs": [],
      "source": [
        "#@markdown #¡Ayuda! - Cuentatokens\n",
        "prompt1 = \"teenager face natgeo golden hour\"  # @param {type:\"string\"}\n",
        "prompt2 = \"greyhound, natgeo\"  # @param {type:\"string\"}\n",
        "\n",
        "\n",
        "def tokens_to_list(tokens):\n",
        "  tokens = list(filter(lambda id: id != 49406 and id != 49407, tokens))\n",
        "  token_list = [tokenizer.decoder.get(t) for t in tokens]\n",
        "  return token_list\n",
        "\n",
        "\n",
        "tokens = tokenizer([prompt1, prompt2], return_tensors=\"np\",\n",
        "                   padding=True).input_ids\n",
        "tokens = [tokens_to_list(tokens[0]), tokens_to_list(tokens[1])]\n",
        "print(f\"{tokens[0]}: {len(tokens[0])}\")\n",
        "print(f\"{tokens[1]}: {len(tokens[1])}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5Nc34EEYkC7y"
      },
      "outputs": [],
      "source": [
        "#@markdown #1. Hiperescala - AVANZADO\n",
        "#@markdown Un único prompt que multiplicamos para hacerlo hiper. Podemos elegir el rango a multiplicar y añadir ruido.\n",
        "\n",
        "num_images = 1\n",
        "\n",
        "prompt = \"dog natgeo\"  # @param {type:\"string\"}\n",
        "semilla = 2100112712  # @param {type:\"number\"}\n",
        "guia = 7.5  #-@param {type:\"slider\", min:4.5, max:21, step:1.5}\n",
        "hiperescala = 0.05  # @param {type:\"slider\", min:-1.95, max:2, step:0.05}\n",
        "pasos = 85  # @param {type:\"slider\", min:25, max:250, step:5}\n",
        "\n",
        "desde = 3  # @param {type:\"slider\", min:0, max:76, step:1}\n",
        "hasta = 76  # @param {type:\"slider\", min:0, max:76, step:1}\n",
        "\n",
        "ruido = 0.5 # @param {type:\"slider\", min:0, max:2, step:0.1}\n",
        "\n",
        "\n",
        "prompts = [\"\"] * num_images  # Unconditional\n",
        "prompts.extend([prompt] * num_images)  # Prompts\n",
        "\n",
        "\n",
        "inputs = tokenizer(prompts, padding=\"max_length\", max_length=tokenizer.model_max_length,\n",
        "                   truncation=True, return_tensors=\"pt\")  # [num_images * 2, 77]\n",
        "\n",
        "with torch.no_grad():\n",
        "  text_embeddings = text_encoder(inputs.input_ids.to(torch_device))[\n",
        "      0]  # [num_images * 2, 77, 768]\n",
        "\n",
        "# Reescalamos los embeddings del prompt\n",
        "text_embeddings[num_images:, desde:(hasta+1)] *= hiperescala\n",
        "\n",
        "# Y añadimos el ruido\n",
        "semilla, generator = get_seed_gen(semilla)\n",
        "text_embeddings[num_images:] += (torch.rand(text_embeddings[num_images:].size(),\n",
        "                                            generator=generator,\n",
        "                                            device=torch_device,\n",
        "                                            )-0.5)*ruido\n",
        "\n",
        "#Esto es una chapu, pero así nos resetea el generador y podemos replicar los\n",
        "#resultados de la versión no avanzada\n",
        "semilla, generator = get_seed_gen(semilla)\n",
        "images = render(pasos, guia, generator, text_embeddings, num_images)\n",
        "filename = f\"1_{prompt}_{semilla}_{hiperescala:0.02f}_{pasos}_{desde}_{hasta}_{ruido:0.01f}.jpg\"\n",
        "\n",
        "print(f\"Semilla: {semilla}\")\n",
        "print(f\"Fichero: {filename}\\n\")\n",
        "\n",
        "grid = image_grid(images, cols=3)\n",
        "grid.save(filename, quality=100, subsampling=0)\n",
        "files.download(filename)\n",
        "grid\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OVLAPjOjpEeN"
      },
      "outputs": [],
      "source": [
        "#@markdown #2. Híbridos - AVANZADO\n",
        "#@markdown Combinación de varios prompts con pesos individuales. Ajuste de pasos y guía.\n",
        "\n",
        "num_images = 1\n",
        "\n",
        "prompt_body = \"a portrait of a @, natgeo\" # @param {type:\"string\"}\n",
        "prompt_subjects = \"seal:1 parrot:3 skull:4\"  # @param {type:\"string\"}\n",
        "prompt_remove = \"hair\" # @param {type:\"string\"}\n",
        "semilla = 1823707604  # @param {type:\"number\"}\n",
        "pasos = 75  # @param {type:\"slider\", min:50, max:250, step:25}\n",
        "guia = 7.5  # @param {type:\"slider\", min:4.5, max:21, step:1.5}\n",
        "\n",
        "regex = re.compile(r'[a-záéíóúñ][a-záéíóúñ ]+:-?\\d+')\n",
        "\n",
        "weights = []\n",
        "prompts = [prompt_remove]\n",
        "weight_sum = 0\n",
        "for subject in regex.findall(prompt_subjects):\n",
        "  values = subject.split(\":\")\n",
        "  values[1] = int(values[1])\n",
        "  weights.append(values[1])\n",
        "  prompts.append(prompt_body.replace('@', values[0].strip()))\n",
        "  weight_sum += abs(values[1])\n",
        "#print(weight_sum, weights, prompts)\n",
        "\n",
        "inputs = tokenizer(prompts, padding=\"max_length\", max_length=tokenizer.model_max_length,\n",
        "                   truncation=True, return_tensors=\"pt\")  # [1+len(subject_list),77]\n",
        "\n",
        "with torch.no_grad():\n",
        "  embeddings = text_encoder(inputs.input_ids.to(torch_device))[\n",
        "      0]  # [1+len(subject_list),77,768]\n",
        "\n",
        "for i in range(embeddings.shape[0]-1):\n",
        "  # print(100*weights[i]/weight_sum)\n",
        "  embeddings[i+1] *= weights[i]/weight_sum\n",
        "\n",
        "text_embeddings = torch.stack(\n",
        "    [embeddings[0], torch.sum(embeddings[1:], axis=0)])  # [2,77,768]\n",
        "\n",
        "\n",
        "semilla, generator = get_seed_gen(semilla)\n",
        "images = render(pasos, guia, generator, text_embeddings, num_images)\n",
        "filename = f\"2_{prompt_body}_{prompt_subjects}_{semilla}_{pasos}_{guia}.jpg\"\n",
        "\n",
        "print(f\"Semilla: {semilla}\")\n",
        "print(filename)\n",
        "\n",
        "grid = image_grid(images, cols=3)\n",
        "grid.save(filename, quality=100, subsampling=0)\n",
        "files.download(filename)\n",
        "grid"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "gpuClass": "premium"
    },
    "gpuClass": "premium",
    "kernelspec": {
      "display_name": "Python 3.8.11 ('base')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.8.11"
    },
    "vscode": {
      "interpreter": {
        "hash": "19c52f9406ed8ee88add4472b642b79c35ade567538d579e8a14fbefe9c2ac9c"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}