JAX_CLIP_Guided_diffusion_nshepperd.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "JAX-CLIP-Guided-diffusion_nshepperd.ipynb",
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/olaviinha/NeuralTextToImage/blob/main/JAX_CLIP_Guided_diffusion_nshepperd.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QZz-fFxNBCug"
      },
      "source": [
        "#<font face=\"Trebuchet MS\" size=\"6\">Big Sleep: JAX CLIP Guided Diffusion <font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><font color=\"#999\" size=\"4\">Neural text-to-image</font><font color=\"#999\" size=\"4\">&nbsp;&nbsp;&nbsp;|&nbsp;&nbsp;&nbsp;</font><a href=\"https://github.com/olaviinha/NeuralImageGeneration\" target=\"_blank\"><font color=\"#999\" size=\"4\">Github</font></a>\n",
        "\n",
        "Big Sleep generates images from text input. It's originally a combination of [CLIP](https://github.com/openai/CLIP) by OpenAI and [BigGAN](https://arxiv.org/abs/1809.11096) by Andrew Brock et al., a concept introduced by [Ryan Murdock](https://github.com/rynmurdock) in his [original notebook](https://colab.research.google.com/drive/1NCceX2mbiKOSlAd_o7IU7nA9UskKN5WR?usp=sharing). The concept has since evolved to multiple directions. This notebook is based on  nshepperd's [JAX CLIP Guided Diffusion v2.3](https://colab.research.google.com/drive/12Bod44YVIXYRh39WRqp0kNz8OUBNFk9Z), which in turn is based on [Katherine Crowson](https://github.com/crowsonkb)'s work.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1zKX4uWFBks2"
      },
      "source": [
        "#@title Licensed under the MIT License { display-mode: \"form\" }\n",
        "\n",
        "# Copyright (c) 2021 Katherine Crowson; nshepperd\n",
        "\n",
        "# Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "# of this software and associated documentation files (the \"Software\"), to deal\n",
        "# in the Software without restriction, including without limitation the rights\n",
        "# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "# copies of the Software, and to permit persons to whom the Software is\n",
        "# furnished to do so, subject to the following conditions:\n",
        "\n",
        "# The above copyright notice and this permission notice shall be included in\n",
        "# all copies or substantial portions of the Software.\n",
        "\n",
        "# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN\n",
        "# THE SOFTWARE."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "r_gJy8nAJfL2"
      },
      "source": [
        "#@markdown ## Check GPU\n",
        "def whatGPU():\n",
        "  x = !nvidia-smi\n",
        "  g = ''.join(x);\n",
        "  restart = False\n",
        "  gpu = 'None'\n",
        "  if 'A100' in g:\n",
        "    # !pip install torch==1.9.0+cu111 torchvision==0.10.0+cu111 torchaudio==0.9.0 -f https://download.pytorch.org/whl/torch_stable.html\n",
        "    gpu = 'A100 40 GB'\n",
        "    restart = True\n",
        "  if 'V100' in g:\n",
        "    gpu = 'V100 32 GB'\n",
        "  if 'P100' in g:\n",
        "    gpu = 'P100 16 GB'\n",
        "  if gpu is 'None':\n",
        "    gpu = x\n",
        "  if restart is True:\n",
        "    gpu = gpu + ' - RESTART RUNTIME!'\n",
        "  return gpu\n",
        "\n",
        "print('GPU:', whatGPU())"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "id": "FmheNGEYE5oV"
      },
      "source": [
        "#@title #Setup\n",
        "#@markdown This cell needs to be run only once. It will mount your Google Drive and setup prerequisites.\n",
        "\n",
        "force_setup = False\n",
        "pip_packages = ''\n",
        "main_repository = ''\n",
        "\n",
        "import os\n",
        "from google.colab import output\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "%cd /content/\n",
        "\n",
        "# inhagcutils\n",
        "if not os.path.isfile('/content/inhagcutils.ipynb') or force_setup == True:\n",
        "  !pip -q install import-ipynb {pip_packages}\n",
        "  !apt -q install exempi\n",
        "  !curl -s -O https://raw.githubusercontent.com/olaviinha/inhagcutils/master/inhagcutils.ipynb\n",
        "import import_ipynb\n",
        "from inhagcutils import *\n",
        "\n",
        "# Mount Drive\n",
        "if not os.path.isdir('/content/drive') and force_setup == False:\n",
        "  from google.colab import drive\n",
        "  drive.mount('/content/drive')\n",
        "\n",
        "# Drive symlink\n",
        "if not os.path.isdir('/content/mydrive') and force_setup == False:\n",
        "  os.symlink('/content/drive/My Drive', '/content/mydrive')\n",
        "  drive_root_set = True\n",
        "drive_root = '/content/mydrive/'\n",
        "\n",
        "\n",
        "\n",
        "# Install dependencies\n",
        "!pip install dm-haiku cbor2 ftfy einops\n",
        "!git clone https://github.com/kingoflolz/CLIP_JAX\n",
        "!git clone https://github.com/nshepperd/jax-guided-diffusion -b v2\n",
        "\n",
        "import sys\n",
        "sys.path.append('./CLIP_JAX')\n",
        "sys.path.append('./jax-guided-diffusion')\n",
        "\n",
        "import math\n",
        "import io\n",
        "import time\n",
        "import os\n",
        "import functools\n",
        "from functools import partial\n",
        "from dataclasses import dataclass\n",
        "\n",
        "from PIL import Image\n",
        "import requests\n",
        "\n",
        "import numpy as np\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import jaxtorch\n",
        "from jaxtorch import PRNG, Context, Module, nn, init\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "import clip_jax\n",
        "\n",
        "from lib.script_util import create_model_and_diffusion, model_and_diffusion_defaults\n",
        "from lib import util\n",
        "from lib.util import pil_from_tensor, pil_to_tensor\n",
        "\n",
        "from IPython import display\n",
        "from torchvision import datasets, transforms, utils\n",
        "from torchvision.transforms import functional as TF\n",
        "import torch.utils.data\n",
        "import torch\n",
        "\n",
        "\n",
        "\n",
        "devices = jax.devices()\n",
        "n_devices = len(devices)\n",
        "print('Using device:', devices)\n",
        "\n",
        "\n",
        "\n",
        "# Define necessary functions\n",
        "\n",
        "def fetch(url_or_path):\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 fd\n",
        "    return open(url_or_path, 'rb')\n",
        "\n",
        "def fetch_model(url_or_path):\n",
        "    basename = os.path.basename(url_or_path)\n",
        "    if os.path.exists(basename):\n",
        "        return basename\n",
        "    else:\n",
        "        !curl -OL '{url_or_path}'\n",
        "        return basename\n",
        "\n",
        "def grey(image):\n",
        "    [*_, c, h, w] = image.shape\n",
        "    return jnp.broadcast_to(image.mean(axis=-3, keepdims=True), image.shape)\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class MakeCutouts(object):\n",
        "    def __init__(self, cut_size, cutn, cut_pow=1., p_grey=0.2):\n",
        "        self.cut_size = cut_size\n",
        "        self.cutn = cutn\n",
        "        self.cut_pow = cut_pow\n",
        "        self.p_grey = p_grey\n",
        "\n",
        "    def __call__(self, input, key):\n",
        "        [b, c, h, w] = input.shape\n",
        "        rng = PRNG(key)\n",
        "        max_size = min(h, w)\n",
        "        min_size = min(h, w, self.cut_size)\n",
        "        cut_us = jax.random.uniform(rng.split(), shape=[self.cutn//2])**self.cut_pow\n",
        "        sizes = (min_size + cut_us * (max_size - min_size + 1)).astype(jnp.int32).clamp(min_size, max_size)\n",
        "        offsets_x = jax.random.uniform(rng.split(), [self.cutn//2], minval=0, maxval=w - sizes)\n",
        "        offsets_y = jax.random.uniform(rng.split(), [self.cutn//2], minval=0, maxval=h - sizes)\n",
        "        cutouts = util.cutouts_images(input, offsets_x, offsets_y, sizes)\n",
        "\n",
        "        lcut_us = jax.random.uniform(rng.split(), shape=[self.cutn//2])\n",
        "        lsizes = (max(h,w) + 10 + lcut_us * 10).astype(jnp.int32)\n",
        "        loffsets_x = jax.random.uniform(rng.split(), [self.cutn//2], minval=-20, maxval=0)\n",
        "        loffsets_y = jax.random.uniform(rng.split(), [self.cutn//2], minval=-20, maxval=0)\n",
        "        lcutouts = util.cutouts_images(input, loffsets_x, loffsets_y, lsizes)\n",
        "\n",
        "        cutouts = jnp.concatenate([cutouts, lcutouts], axis=1)\n",
        "\n",
        "        grey_us = jax.random.uniform(rng.split(), shape=[b, self.cutn, 1, 1, 1])\n",
        "        cutouts = jnp.where(grey_us < self.p_grey, grey(cutouts), cutouts)\n",
        "        cutouts = cutouts.rearrange('b n c h w -> (n b) c h w')\n",
        "        return cutouts\n",
        "\n",
        "    def tree_flatten(self):\n",
        "        return ([self.p_grey, self.cut_pow], (self.cut_size, self.cutn))\n",
        "\n",
        "    @staticmethod\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        (cut_size, cutn) = static\n",
        "        (p_grey, cut_pow) = dynamic\n",
        "        return MakeCutouts(cut_size, cutn, cut_pow, p_grey)\n",
        "\n",
        "def Normalize(mean, std):\n",
        "    mean = jnp.array(mean).reshape(3,1,1)\n",
        "    std = jnp.array(std).reshape(3,1,1)\n",
        "    def forward(image):\n",
        "        return (image - mean) / std\n",
        "    return forward\n",
        "\n",
        "def norm1(x):\n",
        "    \"\"\"Normalize to the unit sphere.\"\"\"\n",
        "    return x / x.square().sum(axis=-1, keepdims=True).sqrt()\n",
        "\n",
        "def spherical_dist_loss(x, y):\n",
        "    x = norm1(x)\n",
        "    y = norm1(y)\n",
        "    return (x - y).square().sum(axis=-1).sqrt().div(2).arcsin().square().mul(2)\n",
        "\n",
        "def tv_loss(input):\n",
        "    \"\"\"L2 total variation loss, as in Mahendran et al.\"\"\"\n",
        "    # input = jnp.pad(input, ((0,0), (0,0), (0,1), (0,1)), mode='edge')\n",
        "    # x_diff = input[..., :-1, 1:] - input[..., :-1, :-1]\n",
        "    # y_diff = input[..., 1:, :-1] - input[..., :-1, :-1]\n",
        "    # return (x_diff**2 + y_diff**2).mean([1, 2, 3])\n",
        "    x_diff = input[..., :, 1:] - input[..., :, :-1]\n",
        "    y_diff = input[..., 1:, :] - input[..., :-1, :]\n",
        "    return x_diff.square().mean([1,2,3]) + y_diff.square().mean([1,2,3])\n",
        "\n",
        "def downscale2d(image, f):\n",
        "  [c, n, h, w] = image.shape\n",
        "  return jax.image.resize(image, [c, n, h//f, w//f], method='cubic')\n",
        "\n",
        "def upscale2d(image, f):\n",
        "  [c, n, h, w] = image.shape\n",
        "  return jax.image.resize(image, [c, n, h*f, w*f], method='cubic')\n",
        "\n",
        "def rms(x):\n",
        "  return x.square().mean().sqrt()\n",
        "\n",
        "@dataclass\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class DiffusionOutput:\n",
        "    v: torch.Tensor\n",
        "    pred: torch.Tensor\n",
        "    eps: torch.Tensor\n",
        "\n",
        "    def tree_flatten(self):\n",
        "        return [self.v, self.pred, self.eps], []\n",
        "\n",
        "    @classmethod\n",
        "    def tree_unflatten(cls, static, dynamic):\n",
        "        return cls(*dynamic)\n",
        "  \n",
        "def alpha_sigma_to_t(alpha, sigma):\n",
        "    return jnp.arctan2(sigma, alpha) * 2 / math.pi\n",
        "\n",
        "def get_ddpm_alphas_sigmas(t):\n",
        "    log_snrs = -jnp.expm1(1e-4 + 10 * t**2).log()\n",
        "    alphas_squared = jax.nn.sigmoid(log_snrs)\n",
        "    sigmas_squared = jax.nn.sigmoid(-log_snrs)\n",
        "    return alphas_squared.sqrt(), sigmas_squared.sqrt()\n",
        "\n",
        "def get_cosine_alphas_sigmas(t):\n",
        "    return jnp.cos(t * math.pi/2), jnp.sin(t * math.pi/2)\n",
        "\n",
        "\n",
        "\n",
        "# Common nn modules.\n",
        "class SkipBlock(nn.Module):\n",
        "    def __init__(self, main, skip=None):\n",
        "        super().__init__()\n",
        "        self.main = nn.Sequential(*main)\n",
        "        self.skip = skip if skip else nn.Identity()\n",
        "\n",
        "    def forward(self, cx, input):\n",
        "        return jnp.concatenate([self.main(cx, input), self.skip(cx, input)], axis=1)\n",
        "\n",
        "\n",
        "class FourierFeatures(nn.Module):\n",
        "    def __init__(self, in_features, out_features, std=1.):\n",
        "        super().__init__()\n",
        "        assert out_features % 2 == 0\n",
        "        self.weight = init.normal(out_features // 2, in_features, stddev=std)\n",
        "\n",
        "    def forward(self, cx, input):\n",
        "        f = 2 * math.pi * input @ cx[self.weight].transpose()\n",
        "        return jnp.concatenate([f.cos(), f.sin()], axis=-1)\n",
        "\n",
        "\n",
        "class AvgPool2d(nn.Module):\n",
        "    def forward(self, cx, x):\n",
        "        [n, c, h, w] = x.shape\n",
        "        x = x.reshape([n, c, h//2, 2, w//2, 2])\n",
        "        x = x.mean((3,5))\n",
        "        return x\n",
        "\n",
        "\n",
        "def expand_to_planes(input, shape):\n",
        "    return input[..., None, None].broadcast_to(list(input.shape) + [shape[2], shape[3]])\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# Secondary Model \n",
        "class ConvBlock(nn.Sequential):\n",
        "    def __init__(self, c_in, c_out):\n",
        "        super().__init__(\n",
        "            nn.Conv2d(c_in, c_out, 3, padding=1),\n",
        "            nn.ReLU(),\n",
        "        )\n",
        "\n",
        "class SecondaryDiffusionImageNet(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        c = 64  # The base channel count\n",
        "\n",
        "        self.timestep_embed = FourierFeatures(1, 16)\n",
        "\n",
        "        self.net = nn.Sequential(\n",
        "            ConvBlock(3 + 16, c),\n",
        "            ConvBlock(c, c),\n",
        "            SkipBlock([\n",
        "                AvgPool2d(),\n",
        "                # nn.image.Downsample2d('linear'),\n",
        "                ConvBlock(c, c * 2),\n",
        "                ConvBlock(c * 2, c * 2),\n",
        "                SkipBlock([\n",
        "                    AvgPool2d(),\n",
        "                    # nn.image.Downsample2d('linear'),\n",
        "                    ConvBlock(c * 2, c * 4),\n",
        "                    ConvBlock(c * 4, c * 4),\n",
        "                    SkipBlock([\n",
        "                        AvgPool2d(),\n",
        "                        # nn.image.Downsample2d('linear'),\n",
        "                        ConvBlock(c * 4, c * 8),\n",
        "                        ConvBlock(c * 8, c * 4),\n",
        "                        nn.image.Upsample2d('linear'),\n",
        "                    ]),\n",
        "                    ConvBlock(c * 8, c * 4),\n",
        "                    ConvBlock(c * 4, c * 2),\n",
        "                    nn.image.Upsample2d('linear'),\n",
        "                ]),\n",
        "                ConvBlock(c * 4, c * 2),\n",
        "                ConvBlock(c * 2, c),\n",
        "                nn.image.Upsample2d('linear'),\n",
        "            ]),\n",
        "            ConvBlock(c * 2, c),\n",
        "            nn.Conv2d(c, 3, 3, padding=1),\n",
        "        )\n",
        "\n",
        "    def forward(self, cx, input, t):\n",
        "        timestep_embed = expand_to_planes(self.timestep_embed(cx, t[:, None]), input.shape)\n",
        "        v = self.net(cx, jnp.concatenate([input, timestep_embed], axis=1))\n",
        "        alphas, sigmas = get_cosine_alphas_sigmas(t)\n",
        "        alphas = alphas[:, None, None, None]\n",
        "        sigmas = sigmas[:, None, None, None]\n",
        "        pred = input * alphas - v * sigmas\n",
        "        eps = input * sigmas + v * alphas\n",
        "        return DiffusionOutput(v, pred, eps)\n",
        "\n",
        "class SecondaryDiffusionImageNet2(nn.Module):\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        c = 64  # The base channel count\n",
        "        cs = [c, c * 2, c * 2, c * 4, c * 4, c * 8]\n",
        "\n",
        "        self.timestep_embed = FourierFeatures(1, 16)\n",
        "        self.down = AvgPool2d()\n",
        "        self.up = nn.image.Upsample2d('linear')\n",
        "\n",
        "        self.net = nn.Sequential(\n",
        "            ConvBlock(3 + 16, cs[0]),\n",
        "            ConvBlock(cs[0], cs[0]),\n",
        "            SkipBlock([\n",
        "                self.down,\n",
        "                ConvBlock(cs[0], cs[1]),\n",
        "                ConvBlock(cs[1], cs[1]),\n",
        "                SkipBlock([\n",
        "                    self.down,\n",
        "                    ConvBlock(cs[1], cs[2]),\n",
        "                    ConvBlock(cs[2], cs[2]),\n",
        "                    SkipBlock([\n",
        "                        self.down,\n",
        "                        ConvBlock(cs[2], cs[3]),\n",
        "                        ConvBlock(cs[3], cs[3]),\n",
        "                        SkipBlock([\n",
        "                            self.down,\n",
        "                            ConvBlock(cs[3], cs[4]),\n",
        "                            ConvBlock(cs[4], cs[4]),\n",
        "                            SkipBlock([\n",
        "                                self.down,\n",
        "                                ConvBlock(cs[4], cs[5]),\n",
        "                                ConvBlock(cs[5], cs[5]),\n",
        "                                ConvBlock(cs[5], cs[5]),\n",
        "                                ConvBlock(cs[5], cs[4]),\n",
        "                                self.up,\n",
        "                            ]),\n",
        "                            ConvBlock(cs[4] * 2, cs[4]),\n",
        "                            ConvBlock(cs[4], cs[3]),\n",
        "                            self.up,\n",
        "                        ]),\n",
        "                        ConvBlock(cs[3] * 2, cs[3]),\n",
        "                        ConvBlock(cs[3], cs[2]),\n",
        "                        self.up,\n",
        "                    ]),\n",
        "                    ConvBlock(cs[2] * 2, cs[2]),\n",
        "                    ConvBlock(cs[2], cs[1]),\n",
        "                    self.up,\n",
        "                ]),\n",
        "                ConvBlock(cs[1] * 2, cs[1]),\n",
        "                ConvBlock(cs[1], cs[0]),\n",
        "                self.up,\n",
        "            ]),\n",
        "            ConvBlock(cs[0] * 2, cs[0]),\n",
        "            nn.Conv2d(cs[0], 3, 3, padding=1),\n",
        "        )\n",
        "\n",
        "    def forward(self, cx, input, t):\n",
        "        timestep_embed = expand_to_planes(self.timestep_embed(cx, t[:, None]), input.shape)\n",
        "        v = self.net(cx, jnp.concatenate([input, timestep_embed], axis=1))\n",
        "        alphas, sigmas = get_cosine_alphas_sigmas(t)\n",
        "        alphas = alphas[:, None, None, None]\n",
        "        sigmas = sigmas[:, None, None, None]\n",
        "        pred = input * alphas - v * sigmas\n",
        "        eps = input * sigmas + v * alphas\n",
        "        return DiffusionOutput(v, pred, eps)\n",
        "\n",
        "secondary1_model = SecondaryDiffusionImageNet()\n",
        "secondary1_params = secondary1_model.init_weights(jax.random.PRNGKey(0))\n",
        "secondary1_params = jaxtorch.pt.load(fetch_model('https://v-diffusion.s3.us-west-2.amazonaws.com/secondary_model_imagenet.pth'))\n",
        "\n",
        "secondary2_model = SecondaryDiffusionImageNet2()\n",
        "secondary2_params = secondary2_model.init_weights(jax.random.PRNGKey(0))\n",
        "secondary2_params = jaxtorch.pt.load(fetch_model('https://v-diffusion.s3.us-west-2.amazonaws.com/secondary_model_imagenet_2.pth'))\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# Anti-JPEG model\n",
        "class ResidualBlock(nn.Module):\n",
        "    def __init__(self, main, skip=None):\n",
        "        super().__init__()\n",
        "        self.main = nn.Sequential(*main)\n",
        "        self.skip = skip if skip else nn.Identity()\n",
        "\n",
        "    def forward(self, cx, input):\n",
        "        return self.main(cx, input) + self.skip(cx, input)\n",
        "\n",
        "\n",
        "class ResConvBlock(ResidualBlock):\n",
        "    def __init__(self, c_in, c_mid, c_out, dropout=True):\n",
        "        skip = None if c_in == c_out else nn.Conv2d(c_in, c_out, 1, bias=False)\n",
        "        super().__init__([\n",
        "            nn.LeakyReLU(),\n",
        "            nn.Conv2d(c_in, c_mid, 3, padding=1),\n",
        "            nn.LeakyReLU(),\n",
        "            nn.Conv2d(c_mid, c_out, 3, padding=1),\n",
        "        ], skip)\n",
        "        \n",
        "\n",
        "CHANNELS=64\n",
        "class JPEGModel(nn.Module):\n",
        "    def __init__(self, c=CHANNELS):\n",
        "        super().__init__()\n",
        "\n",
        "        self.timestep_embed = FourierFeatures(1, 16, std=1.0)\n",
        "        self.class_embed = nn.Embedding(3, 16)\n",
        "\n",
        "        self.arch = '11(22(22(2)22)22)11'\n",
        "\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Conv2d(3 + 16 + 16, c, 1),\n",
        "            ResConvBlock(c, c, c),\n",
        "            ResConvBlock(c, c, c),\n",
        "            SkipBlock([\n",
        "                nn.image.Downsample2d(),\n",
        "                ResConvBlock(c,     c * 2, c * 2),\n",
        "                ResConvBlock(c * 2, c * 2, c * 2),\n",
        "                SkipBlock([\n",
        "                    nn.image.Downsample2d(),\n",
        "                    ResConvBlock(c * 2, c * 2, c * 2),\n",
        "                    ResConvBlock(c * 2, 2 * 2, c * 2),\n",
        "                    SkipBlock([\n",
        "                        nn.image.Downsample2d(),\n",
        "                        ResConvBlock(c * 2, c * 2, c * 2),\n",
        "                        nn.image.Upsample2d(),\n",
        "                    ]),\n",
        "                    ResConvBlock(c * 4, c * 2, c * 2),\n",
        "                    ResConvBlock(c * 2, c * 2, c * 2),\n",
        "                    nn.image.Upsample2d(),\n",
        "                ]),\n",
        "                ResConvBlock(c * 4, c * 2, c * 2),\n",
        "                ResConvBlock(c * 2, c * 2, c),\n",
        "                nn.image.Upsample2d(),\n",
        "            ]),\n",
        "            ResConvBlock(c * 2, c, c),\n",
        "            ResConvBlock(c, c, 3, dropout=False),\n",
        "        )\n",
        "\n",
        "    def forward(self, cx, input, ts, cond):\n",
        "        [n, c, h, w] = input.shape\n",
        "        timestep_embed = expand_to_planes(self.timestep_embed(cx, ts[:, None]), input.shape)\n",
        "        class_embed = expand_to_planes(self.class_embed(cx, cond), input.shape)\n",
        "        v = self.net(cx, jnp.concatenate([input, timestep_embed, class_embed], axis=1))\n",
        "        alphas, sigmas = get_cosine_alphas_sigmas(ts)\n",
        "        alphas = alphas[:, None, None, None]\n",
        "        sigmas = sigmas[:, None, None, None]\n",
        "        pred = input * alphas - v * sigmas\n",
        "        eps = input * sigmas + v * alphas\n",
        "        return DiffusionOutput(v, pred, eps)\n",
        "\n",
        "jpeg_model = JPEGModel()\n",
        "jpeg_params = jpeg_model.init_weights(jax.random.PRNGKey(0))\n",
        "jpeg_params = jaxtorch.pt.load(fetch_model('https://set.zlkj.in/models/diffusion/jpeg-db-oi-614.pt'))['params_ema']\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "# Secondary Anti-JPEG Classifier\n",
        "\n",
        "CHANNELS=64\n",
        "class Classifier(nn.Module):\n",
        "    def __init__(self, c=CHANNELS):\n",
        "        super().__init__()\n",
        "\n",
        "        self.timestep_embed = FourierFeatures(1, 16, std=1.0)\n",
        "\n",
        "        self.arch = '11-22-22-22'\n",
        "\n",
        "        self.net = nn.Sequential(\n",
        "            nn.Conv2d(3 + 16, c, 1),\n",
        "            ResConvBlock(c, c, c),\n",
        "            ResConvBlock(c, c, c),\n",
        "            nn.image.Downsample2d(),\n",
        "            ResConvBlock(c,     c * 2, c * 2),\n",
        "            ResConvBlock(c * 2, c * 2, c * 2),\n",
        "            nn.image.Downsample2d(),\n",
        "            ResConvBlock(c * 2, c * 2, c * 2),\n",
        "            ResConvBlock(c * 2, 2 * 2, c * 2),\n",
        "            nn.image.Downsample2d(),\n",
        "            ResConvBlock(c * 2, c * 2, c * 2),\n",
        "            ResConvBlock(c * 2, c * 2, c * 2),\n",
        "            ResConvBlock(c * 2, c * 2, 1, dropout=False),\n",
        "        )\n",
        "\n",
        "    def forward(self, cx, input, ts):\n",
        "        [n, c, h, w] = input.shape\n",
        "        timestep_embed = expand_to_planes(self.timestep_embed(cx, ts[:, None]), input.shape)\n",
        "        return self.net(cx, jnp.concatenate([input, timestep_embed], axis=1))\n",
        "\n",
        "    def score(self, cx, reals, ts, cond, flood_level):\n",
        "        cond = cond[:, None, None, None]\n",
        "        logits = self.forward(cx, reals, ts)\n",
        "        loss = -jax.nn.log_sigmoid(jnp.where(cond==0, logits, -logits))\n",
        "        loss = loss.clamp(minval=flood_level, maxval=None)\n",
        "        return loss.mean()\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def classifier_probs(classifier_params, x, ts):\n",
        "  n = x.shape[0]\n",
        "  cx = Context(classifier_params, jax.random.PRNGKey(0)).eval_mode_()\n",
        "  probs = jax.nn.sigmoid(classifier_model(cx, x, ts.broadcast_to([n])))\n",
        "  return probs\n",
        "\n",
        "classifier_model = Classifier()\n",
        "classifier_params = classifier_model.init_weights(jax.random.PRNGKey(0))\n",
        "classifier_params = jaxtorch.pt.load(fetch_model('https://set.zlkj.in/models/diffusion/jpeg-classifier-72.pt'))['params_ema']\n",
        "\n",
        "\n",
        "\n",
        "model_config = model_and_diffusion_defaults()\n",
        "model_config.update({\n",
        "    'attention_resolutions': '32, 16, 8',\n",
        "    'class_cond': False,\n",
        "    'diffusion_steps': 1000,\n",
        "    'rescale_timesteps': True,\n",
        "    'timestep_respacing': '1000',\n",
        "    'image_size': 512,     # Change to either 256 or 512 to select the openai model\n",
        "    'learn_sigma': True,\n",
        "    'noise_schedule': 'linear',\n",
        "    'num_channels': 256,\n",
        "    'num_head_channels': 64,\n",
        "    'num_res_blocks': 2,\n",
        "    'resblock_updown': True,\n",
        "    'use_scale_shift_norm': True,\n",
        "    'use_checkpoint': False # Set to True to save memory\n",
        "})\n",
        "\n",
        "# Load models\n",
        "\n",
        "model, diffusion = create_model_and_diffusion(**model_config)\n",
        "model_params = model.init_weights(jax.random.PRNGKey(0))\n",
        "\n",
        "print('Loading state dict...')\n",
        "model_urls = {\n",
        "    512: 'https://set.zlkj.in/models/diffusion/512x512_diffusion_uncond_finetune_008100.pt',\n",
        "    256: 'https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt'\n",
        "}\n",
        "with torch.no_grad():\n",
        "  model_params = model.load_state_dict(model_params, jaxtorch.pt.load(fetch_model(model_urls[model_config['image_size']])))\n",
        "\n",
        "\n",
        "\n",
        "# Define combinators.\n",
        "\n",
        "# These (ab)use the jax pytree registration system to define parameterised\n",
        "# objects for doing various things, which are compatible with jax.jit.\n",
        "\n",
        "# For jit compatibility an object needs to act as a pytree, which means implementing two methods:\n",
        "#  - tree_flatten(self): returns two lists of the object's fields: \n",
        "#       1. 'dynamic' parameters: things which can be jax tensors, or other pytrees\n",
        "#       2. 'static' parameters: arbitrary python objects, will trigger recompilation when changed\n",
        "#  - tree_unflatten(static, dynamic): reconstitutes the object from its parts\n",
        "\n",
        "# With these tricks, you can simply define your cond_fn as an object, as is done\n",
        "# below, and pass it into the jitted sample step as a regular argument. JAX will\n",
        "# handle recompiling the jitted code whenever a control-flow affecting parameter\n",
        "# is changed (such as cut_batches).\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class CosineModel(object):\n",
        "    def __init__(self, model, params, **kwargs):\n",
        "      self.model = model\n",
        "      self.params = params\n",
        "      self.kwargs = kwargs\n",
        "\n",
        "    @jax.jit\n",
        "    def __call__(self, x, t, key):\n",
        "        n = x.shape[0]\n",
        "        alpha, sigma = get_ddpm_alphas_sigmas(t)\n",
        "        cosine_t = alpha_sigma_to_t(alpha, sigma)\n",
        "        cx = Context(self.params, key).eval_mode_()\n",
        "        return self.model(cx, x, cosine_t.broadcast_to([n]), **self.kwargs)\n",
        "\n",
        "    def tree_flatten(self):\n",
        "        return [self.params, self.kwargs], [self.model]\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        [params, kwargs] = dynamic\n",
        "        [model] = static\n",
        "        return CosineModel(model, params, **kwargs)\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class OpenaiModel(object):\n",
        "    def __init__(self, model, params):\n",
        "      self.model = model\n",
        "      self.params = params\n",
        "\n",
        "    @jax.jit\n",
        "    def __call__(self, x, t, key):\n",
        "        n = x.shape[0]\n",
        "        alpha, sigma = get_ddpm_alphas_sigmas(t)\n",
        "        cx = Context(self.params, key).eval_mode_()\n",
        "        openai_t = (t * 1000).broadcast_to([n])\n",
        "        eps = self.model(cx, x, openai_t)[:, :3, :, :]\n",
        "        pred = (x - eps * sigma) / alpha\n",
        "        v    = (eps - x * sigma) / alpha\n",
        "        return DiffusionOutput(v, pred, eps)\n",
        "\n",
        "    def tree_flatten(self):\n",
        "        return [self.params], [self.model]\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        [params] = dynamic\n",
        "        [model] = static\n",
        "        return OpenaiModel(model, params)\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class Perceptor(object):\n",
        "    # Wraps a CLIP instance and its parameters.\n",
        "    def __init__(self, image_fn, text_fn, clip_params):\n",
        "        self.image_fn = image_fn\n",
        "        self.text_fn = text_fn\n",
        "        self.clip_params = clip_params\n",
        "        \n",
        "    @jax.jit\n",
        "    def embed_cutouts(self, cutouts):\n",
        "        return norm1(self.image_fn(self.clip_params, cutouts))\n",
        "    \n",
        "    def embed_text(self, text):\n",
        "        tokens = clip_jax.tokenize([text])\n",
        "        text_embed = self.text_fn(self.clip_params, tokens)\n",
        "        return norm1(text_embed.reshape(512))\n",
        "\n",
        "    def tree_flatten(self):\n",
        "        return [self.clip_params], [self.image_fn, self.text_fn]\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        [clip_params] = dynamic\n",
        "        [image_fn, text_fn] = static\n",
        "        return Perceptor(image_fn, text_fn, clip_params)\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class LerpModels(object):\n",
        "    \"\"\"Linear combination of diffusion models.\"\"\"\n",
        "    def __init__(self, models):\n",
        "        self.models = models\n",
        "\n",
        "    def __call__(self, x, t, key):\n",
        "        outputs = [m(x,t,key) for (m,w) in self.models]\n",
        "        v = sum(out.v * w for (out, (m,w)) in zip(outputs, self.models))\n",
        "        pred = sum(out.pred * w for (out, (m,w)) in zip(outputs, self.models))\n",
        "        eps = sum(out.eps * w for (out, (m,w)) in zip(outputs, self.models))\n",
        "        return DiffusionOutput(v, pred, eps)\n",
        "\n",
        "    def tree_flatten(self):\n",
        "        return [self.models], []\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        return LerpModels(*dynamic)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class CondCLIP(object):\n",
        "    # CLIP guidance loss. Pushes the image toward a text prompt.\n",
        "    def __init__(self, text_embed, clip_guidance_scale, perceptor, make_cutouts, cut_batches):\n",
        "        self.text_embed = text_embed\n",
        "        self.clip_guidance_scale = clip_guidance_scale\n",
        "        self.perceptor = perceptor\n",
        "        self.make_cutouts = make_cutouts\n",
        "        self.cut_batches = cut_batches\n",
        "    def __call__(self, x_in, key):\n",
        "        n = x_in.shape[0]\n",
        "        def main_clip_loss(x_in, key):\n",
        "            cutouts = normalize(self.make_cutouts(x_in.add(1).div(2), key))\n",
        "            image_embeds = self.perceptor.embed_cutouts(cutouts).reshape([self.make_cutouts.cutn, n, 512])\n",
        "            losses = spherical_dist_loss(image_embeds, self.text_embed).mean(0)\n",
        "            return losses.sum() * self.clip_guidance_scale\n",
        "        num_cuts = self.cut_batches\n",
        "        keys = jnp.stack(jax.random.split(key, num_cuts))\n",
        "        main_clip_grad = jax.lax.scan(lambda total, key: (total + jax.grad(main_clip_loss)(x_in, key), key),\n",
        "                                        jnp.zeros_like(x_in),\n",
        "                                        keys)[0] / num_cuts\n",
        "        return main_clip_grad\n",
        "    def tree_flatten(self):\n",
        "        return [self.text_embed, self.clip_guidance_scale, self.perceptor, self.make_cutouts], [self.cut_batches]\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        [text_embed, clip_guidance_scale, perceptor, make_cutouts] = dynamic\n",
        "        [cut_batches] = static\n",
        "        return CondCLIP(text_embed, clip_guidance_scale, perceptor, make_cutouts, cut_batches)\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class CondTV(object):\n",
        "    # Multiscale Total Variation loss. Tries to smooth out the image.\n",
        "    def __init__(self, tv_scale):\n",
        "        self.tv_scale = tv_scale\n",
        "    def __call__(self, x_in, key):\n",
        "        def sum_tv_loss(x_in, f=None):\n",
        "            if f is not None:\n",
        "                x_in = downscale2d(x_in, f)\n",
        "            return tv_loss(x_in).sum() * self.tv_scale\n",
        "        tv_grad_512 = jax.grad(sum_tv_loss)(x_in)\n",
        "        tv_grad_256 = jax.grad(partial(sum_tv_loss,f=2))(x_in)\n",
        "        tv_grad_128 = jax.grad(partial(sum_tv_loss,f=4))(x_in)\n",
        "        return tv_grad_512 + tv_grad_256 + tv_grad_128        \n",
        "    def tree_flatten(self):\n",
        "        return [self.tv_scale], []\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        return CondTV(*dynamic)\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class CondSat(object):\n",
        "    # Saturation loss. Tries to prevent the image from going out of range.\n",
        "    def __init__(self, sat_scale):\n",
        "        self.sat_scale = sat_scale\n",
        "    def __call__(self, x_in, key):\n",
        "        def saturation_loss(x_in):\n",
        "            return jnp.abs(x_in - x_in.clamp(minval=-1,maxval=1)).mean()\n",
        "        return self.sat_scale * jax.grad(saturation_loss)(x_in)\n",
        "    def tree_flatten(self):\n",
        "        return [self.sat_scale], []\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        return CondSat(*dynamic)\n",
        "\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class CondMSE(object):\n",
        "    # MSE loss. Targets the output towards an image.\n",
        "    def __init__(self, target, mse_scale):\n",
        "        self.target = target\n",
        "        self.mse_scale = mse_scale\n",
        "    def __call__(self, x_in, key):\n",
        "        def mse_loss(x_in):\n",
        "            return (x_in - self.target).square().mean()\n",
        "        return self.mse_scale * jax.grad(mse_loss)(x_in)\n",
        "    def tree_flatten(self):\n",
        "        return [self.target, self.mse_scale], []\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        return CondMSE(*dynamic)\n",
        "\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class MainCondFn(object):\n",
        "    # Used to construct the main cond_fn. Accepts a diffusion model which will \n",
        "    # be used for denoising, plus a list of 'conditions' which will\n",
        "    # generate gradient of a loss wrt the denoised, to be summed together.\n",
        "    def __init__(self, diffusion, conditions, use='pred'):\n",
        "        self.diffusion = diffusion\n",
        "        self.conditions = conditions\n",
        "        self.use = use\n",
        "\n",
        "    @jax.jit\n",
        "    def __call__(self, key, x, t):\n",
        "        rng = PRNG(key)\n",
        "        n = x.shape[0]\n",
        "\n",
        "        alphas, sigmas = get_ddpm_alphas_sigmas(t)\n",
        "\n",
        "        def denoise(key, x):\n",
        "            pred = self.diffusion(x, t, key).pred\n",
        "            if self.use == 'pred':\n",
        "                return pred\n",
        "            elif self.use == 'x_in':\n",
        "                return pred * sigmas + x * alphas\n",
        "        (x_in, backward) = jax.vjp(partial(denoise, rng.split()), x)\n",
        "\n",
        "        total = jnp.zeros_like(x_in)\n",
        "        for cond in self.conditions:\n",
        "            total += cond(x_in, rng.split())\n",
        "        final_grad = -backward(total)[0]\n",
        "\n",
        "        # clamp gradients to a max of 0.2\n",
        "        magnitude = final_grad.square().mean(axis=(1,2,3), keepdims=True).sqrt()\n",
        "        final_grad = final_grad * jnp.where(magnitude > 0.2, 0.2 / magnitude, 1.0)\n",
        "        return final_grad  \n",
        "    def tree_flatten(self):\n",
        "        return [self.diffusion, self.conditions], [self.use]\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        return MainCondFn(*dynamic, *static)\n",
        "\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class ClassifierFn(object):\n",
        "    def __init__(self, model, params, guidance_scale, **kwargs):\n",
        "       self.model = model\n",
        "       self.params = params\n",
        "       self.guidance_scale = guidance_scale\n",
        "       self.kwargs = kwargs\n",
        "\n",
        "    @jax.jit\n",
        "    def __call__(self, key, x, t):\n",
        "        n = x.shape[0]\n",
        "        alpha, sigma = get_ddpm_alphas_sigmas(t)\n",
        "        cosine_t = alpha_sigma_to_t(alpha, sigma).broadcast_to([n])\n",
        "        def fwd(x):\n",
        "          cx = Context(self.params, key).eval_mode_()\n",
        "          return self.guidance_scale * self.model.score(cx, x, cosine_t, **self.kwargs)\n",
        "        return -jax.grad(fwd)(x)\n",
        "    def tree_flatten(self):\n",
        "        return [self.params, self.guidance_scale, self.kwargs], [self.model]\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        [params, guidance_scale, kwargs] = dynamic\n",
        "        [model] = static\n",
        "        return ClassifierFn(model, params, guidance_scale, **kwargs)\n",
        "\n",
        "\n",
        "@jax.tree_util.register_pytree_node_class\n",
        "class CondFns(object):\n",
        "    def __init__(self, *conditions):\n",
        "        self.conditions = conditions\n",
        "    def __call__(self, key, x, t):\n",
        "        rng = PRNG(key)\n",
        "        total = jnp.zeros_like(x)\n",
        "        for cond in self.conditions:\n",
        "          total += cond(rng.split(), x, t)\n",
        "        return total\n",
        "    def tree_flatten(self):\n",
        "        return [self.conditions], []\n",
        "    def tree_unflatten(static, dynamic):\n",
        "        [conditions] = dynamic\n",
        "        return MixCondFn(*conditions)\n",
        "\n",
        "\n",
        "\n",
        "def sample_step(key, x, t1, t2, diffusion, cond_fn, eta):\n",
        "    rng = PRNG(key)\n",
        "\n",
        "    n = x.shape[0]\n",
        "    alpha1, sigma1 = get_ddpm_alphas_sigmas(t1)\n",
        "    alpha2, sigma2 = get_ddpm_alphas_sigmas(t2)\n",
        "\n",
        "    # Run the model\n",
        "    out = diffusion(x, t1, rng.split())\n",
        "    eps = out.eps\n",
        "    pred0 = out.pred\n",
        "\n",
        "    # # Predict the denoised image\n",
        "    # pred0 = (x - eps * sigma1) / alpha1\n",
        "\n",
        "    # Adjust eps with conditioning gradient\n",
        "    cond_score = cond_fn(rng.split(), x, t1)\n",
        "    eps = eps - sigma1 * cond_score\n",
        "\n",
        "    # Predict the denoised image with conditioning\n",
        "    pred = (x - eps * sigma1) / alpha1\n",
        "\n",
        "    # Negative eta allows more extreme levels of noise.\n",
        "    ddpm_sigma = (sigma2**2 / sigma1**2).sqrt() * (1 - alpha1**2 / alpha2**2).sqrt()\n",
        "    ddim_sigma = jnp.where(eta >= 0.0, \n",
        "                           eta * ddpm_sigma, # Normal: eta interpolates between ddim and ddpm\n",
        "                           -eta * sigma2)    # Extreme: eta interpolates between ddim and q_sample(pred)\n",
        "    adjusted_sigma = (sigma2**2 - ddim_sigma**2).sqrt()\n",
        "\n",
        "    # Recombine the predicted noise and predicted denoised image in the\n",
        "    # correct proportions for the next step\n",
        "    x = pred * alpha2 + eps * adjusted_sigma\n",
        "\n",
        "    # Add the correct amount of fresh noise\n",
        "    x += jax.random.normal(rng.split(), x.shape) * ddim_sigma\n",
        "    return x, pred0\n",
        "\n",
        "\n",
        "\n",
        "clip_size = 224\n",
        "normalize = Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
        "                      std=[0.26862954, 0.26130258, 0.27577711])\n",
        "\n",
        "image_fn, text_fn, clip_params, _ = clip_jax.load('ViT-B/32')\n",
        "vit32 = Perceptor(image_fn, text_fn, clip_params)\n",
        "image_fn, text_fn, clip_params, _ = clip_jax.load('ViT-B/16')\n",
        "vit16 = Perceptor(image_fn, text_fn, clip_params)\n",
        "\n",
        "output.clear()\n",
        "op(c.ok, 'Done.', time=True)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OoIL7ayzq7kC",
        "cellView": "form"
      },
      "source": [
        "generate_image_of = \"\" #@param {type:\"string\"}\n",
        "output_dir = \"\" #@param {type:\"string\"}\n",
        "iterations = \"250\" #@param {type:\"string\"}\n",
        "#@markdown <small>Input step numbers, in case you want to save some earlier iterations to `output_dir`, **in addition** to the finished images. Separate by comma, e.g. _200, 225_\n",
        "save_steps = \"\" #@param {type:\"string\"}\n",
        "\n",
        "save_location = drive_root+output_dir\n",
        "image_size = (640, 640)\n",
        "\n",
        "n_batches = 1\n",
        "batch_size = 2\n",
        "repeats = 1\n",
        "#title = \"Unreal Engine style CGI rendering of 3D video game scene\"\n",
        "\n",
        "texts = [generate_image_of]\n",
        "\n",
        "\n",
        "for txt_index, title in enumerate(texts):\n",
        "\n",
        "  clip_guidance_scale = 2000 # Note: with two perceptors, combined guidance scale is 2x because they are added together.\n",
        "  tv_scale = 150  # Smooths out the image.\n",
        "  sat_scale = 600 # Prevents image from over/under-saturating.\n",
        "  steps = int(iterations)    # Number of steps for sampling. Generally, more = better.\n",
        "  eta = 1.0 # 0.0: DDIM | 1.0: DDPM | -1.0: Extreme noise (q_sample)\n",
        "\n",
        "  cutn = 16 # effective cutn is cut_batches * this\n",
        "  cut_pow = 1.0\n",
        "  cut_batches = 4\n",
        "  make_cutouts = MakeCutouts(clip_size, cutn, cut_pow=cut_pow)\n",
        "\n",
        "  # n_batches = 4\n",
        "  init_image = None\n",
        "  skip_timesteps = 0\n",
        "  seed = None # if None, uses the current time in seconds.\n",
        "\n",
        "  # OpenAI used T=1000 to 0. We've just rescaled to between 1 and 0.\n",
        "  schedule = jnp.linspace(1, 0, steps)[skip_timesteps:]\n",
        "\n",
        "\n",
        "  openai = OpenaiModel(model, model_params)\n",
        "  secondary1 = CosineModel(secondary1_model, secondary1_params)\n",
        "  secondary2 = CosineModel(secondary2_model, secondary2_params)\n",
        "  jpeg_0 = CosineModel(jpeg_model, jpeg_params, cond=jnp.array([0]*batch_size)) # Clean class\n",
        "  jpeg_1 = CosineModel(jpeg_model, jpeg_params, cond=jnp.array([2]*batch_size)) # Noisy class\n",
        "\n",
        "  jpeg_classifier_fn = ClassifierFn(classifier_model, classifier_params, \n",
        "                                    guidance_scale=10000.0, # will generally depend on image size\n",
        "                                    cond=jnp.array([0]*batch_size), # Clean class\n",
        "                                    flood_level=0.7) # Prevent over-optimization\n",
        "\n",
        "  diffusion = LerpModels([(openai, 1.0),\n",
        "                          (jpeg_0, 1.0),\n",
        "                          (jpeg_1, -1.0)])\n",
        "  cond_model = secondary2\n",
        "\n",
        "  # target = Image.open(fetch('https://pbs.twimg.com/media/FBbTU8hVEAM10dL?format=png&name=small')).convert('RGB')\n",
        "  # target = target.resize(image_size, Image.LANCZOS)\n",
        "  # target = jnp.array(TF.to_tensor(target)) * 2 - 1\n",
        "\n",
        "  cond_fn = CondFns(MainCondFn(cond_model, [\n",
        "                      CondCLIP(vit32.embed_text(title), clip_guidance_scale, vit32, make_cutouts, cut_batches),\n",
        "                      CondCLIP(vit16.embed_text(title), clip_guidance_scale, vit16, make_cutouts, cut_batches),\n",
        "                      # CondTV(tv_scale),\n",
        "                      # CondMSE(target, 256*256*3),\n",
        "                      CondSat(sat_scale),\n",
        "                      ], use='pred'),\n",
        "                    jpeg_classifier_fn)\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "  def proc_init_image(init_image):\n",
        "      if init_image.endswith(':parts512'):\n",
        "          url = init_image.rsplit(':', 1)[0]\n",
        "          init = Image.open(fetch(url)).convert('RGB')\n",
        "          init = pil_to_tensor(init).mul(2).sub(1)\n",
        "          [c, h, w] = init.shape\n",
        "          indices = [(x, y)\n",
        "                    for y in range(0, h, 512)\n",
        "                    for x in range(0, w, 512)]\n",
        "          indices = (indices * batch_size)[:batch_size]\n",
        "          parts = [init[:, y:y+512, x:x+512] for (x, y) in indices]\n",
        "          init = jnp.stack(parts)\n",
        "          init = jax.image.resize(init, [batch_size, c, image_size[1], image_size[0]], method='lanczos3')\n",
        "          return init\n",
        "\n",
        "      init = Image.open(fetch(init_image)).convert('RGB')\n",
        "      init = init.resize(image_size, Image.LANCZOS)\n",
        "      init = pil_to_tensor(init).unsqueeze(0).mul(2).sub(1)\n",
        "      return init\n",
        "\n",
        "  @torch.no_grad()\n",
        "  def run():\n",
        "      if seed is None:\n",
        "          local_seed = int(time.time())\n",
        "      else:\n",
        "          local_seed = seed\n",
        "      #print(f'Starting run with seed {local_seed}...')\n",
        "      op(c.title, title)\n",
        "      rng = PRNG(jax.random.PRNGKey(local_seed))\n",
        "\n",
        "      init = None\n",
        "      if init_image is not None:\n",
        "          if type(init_image) is list:\n",
        "              init = jnp.concatenate([proc_init_image(url) for url in init_image], axis=0)\n",
        "          else:\n",
        "              init = proc_init_image(init_image)\n",
        "\n",
        "\n",
        "      for i in range(n_batches):\n",
        "          timestring = time.strftime('%Y%m%d%H%M%S')\n",
        "\n",
        "          ts = schedule\n",
        "          alphas, sigmas = get_ddpm_alphas_sigmas(ts)\n",
        "          cosine_ts = alpha_sigma_to_t(alphas, sigmas)\n",
        "\n",
        "          x = sigmas[0] * jax.random.normal(rng.split(), [batch_size, 3, image_size[1], image_size[0]])\n",
        "\n",
        "          if init is not None:\n",
        "              x = x + alphas[0] * init\n",
        "\n",
        "          # Main loop\n",
        "          local_steps = schedule.shape[0] - 1\n",
        "          for j in tqdm(range(local_steps)):\n",
        "              if ts[j] != ts[j+1]:\n",
        "                  # Skip steps where the ts are the same, to make it easier to\n",
        "                  # make complicated schedules out of cat'ing linspaces.\n",
        "                  x, pred = sample_step(rng.split(), x, ts[j], ts[j+1], diffusion, cond_fn, eta)\n",
        "              if j % 50 == 0 or j == local_steps - 1:\n",
        "                  images = pred.add(1).div(2).clamp(0, 1)\n",
        "                  # probs = classifier_probs(classifier_params, x, cosine_ts[j+1])\n",
        "                  # probs = jax.image.resize(probs * jnp.ones([batch_size, 3, 1, 1]), pred.shape, 'cubic')\n",
        "                  # images = jnp.concatenate([images, probs],axis=0)\n",
        "                  images = torch.tensor(np.array(images))\n",
        "                  display.display(TF.to_pil_image(utils.make_grid(images, 4).cpu()))\n",
        "                  if j in [int(x.strip()) for x in save_steps.split(',')]:\n",
        "                    for idx, temp_image in enumerate(images):\n",
        "                      this_title = title[:100]\n",
        "                      tmp_img = TF.to_pil_image(temp_image)\n",
        "                      tmp_img.save(f'{save_location}/{timestring}_{idx}_{j}_{this_title}.png')\n",
        "\n",
        "          # Save samples\n",
        "          os.makedirs('samples', exist_ok=True)\n",
        "          os.makedirs(save_location, exist_ok=True)\n",
        "          for k in range(batch_size):\n",
        "              this_title = title[:100]\n",
        "              dname = f'samples/{timestring}_{k}_{this_title}.png'\n",
        "              pil_image = TF.to_pil_image(images[k])\n",
        "              pil_image.save(dname)\n",
        "              pil_image.save(f'{save_location}/{timestring}_{k}_{this_title}.png')\n",
        "\n",
        "  try:\n",
        "    run()\n",
        "    success = True\n",
        "  except:\n",
        "    import traceback\n",
        "    traceback.print_exc()\n",
        "    success = False\n",
        "  assert success\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}