trial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Originally made by Katherine Crowson (https://github.com/crowsonkb, https://twitter.com/RiversHaveWings)\n",
    "# The original BigGAN+CLIP method was by https://twitter.com/advadnoun\n",
    "\n",
    "import argparse\n",
    "import math\n",
    "import random\n",
    "# from email.policy import default\n",
    "from urllib.request import urlopen\n",
    "from tqdm import tqdm\n",
    "import sys\n",
    "import os\n",
    "\n",
    "# pip install taming-transformers doesn't work with Gumbel, but does not yet work with coco etc\n",
    "# appending the path does work with Gumbel, but gives ModuleNotFoundError: No module named 'transformers' for coco etc\n",
    "sys.path.append('taming-transformers')\n",
    "\n",
    "from omegaconf import OmegaConf\n",
    "from taming.models import cond_transformer, vqgan\n",
    "#import taming.modules \n",
    "\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.nn import functional as F\n",
    "from torchvision import transforms\n",
    "from torchvision.transforms import functional as TF\n",
    "from torch.cuda import get_device_properties\n",
    "torch.backends.cudnn.benchmark = False\t\t# NR: True is a bit faster, but can lead to OOM. False is more deterministic.\n",
    "#torch.use_deterministic_algorithms(True)\t# NR: grid_sampler_2d_backward_cuda does not have a deterministic implementation\n",
    "\n",
    "from torch_optimizer import DiffGrad, AdamP, RAdam\n",
    "\n",
    "from CLIP import clip\n",
    "import kornia.augmentation as K\n",
    "import numpy as np\n",
    "import imageio\n",
    "\n",
    "from PIL import ImageFile, Image, PngImagePlugin, ImageChops\n",
    "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
    "\n",
    "from subprocess import Popen, PIPE\n",
    "import re\n",
    "\n",
    "# Supress warnings\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# VQGAN"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check for GPU and reduce the default image size if low VRAM\n",
    "default_image_size = 512  # >8GB VRAM\n",
    "if not torch.cuda.is_available():\n",
    "    default_image_size = 256  # no GPU found\n",
    "elif get_device_properties(0).total_memory <= 2 ** 33:  # 2 ** 33 = 8,589,934,592 bytes = 8 GB\n",
    "    default_image_size = 318  # <8GB VRAM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts = None\n",
    "max_iterations = 500\n",
    "save_every = 50\n",
    "size = [256, 256]\n",
    "clip_model = 'ViT-B/32'\n",
    "vqgan_config = f'checkpoints/vqgan_imagenet_f16_16384.yaml'\n",
    "vqgan_checkpoint = f'checkpoints/vqgan_imagenet_f16_16384.ckpt'\n",
    "learning_rate = 0.1\n",
    "step_size = 0.1\n",
    "cut_method = 'latest'\n",
    "cutn = 32\n",
    "cut_pow = 1.0\n",
    "output = \"output.png\"\n",
    "augments = [['Af', 'Pe', 'Ji', 'Er']]\n",
    "cuda_device = \"cuda:0\"\n",
    "seed = None\n",
    "display_freq = 50"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "if prompts == None:\n",
    "    prompts = \"A cute, smiling, Nerdy Rodent\"\n",
    "\n",
    "# Split text prompts using the pipe character (weights are split later)\n",
    "if prompts != None:\n",
    "    # For stories, there will be many phrases\n",
    "    story_phrases = [phrase.strip() for phrase in prompts.split(\"^\")]\n",
    "    \n",
    "    # Make a list of all phrases\n",
    "    all_phrases = []\n",
    "    for phrase in story_phrases:\n",
    "        all_phrases.append(phrase.split(\"|\"))\n",
    "    \n",
    "    # First phrase\n",
    "    prompts = all_phrases[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not cuda_device == 'cpu' and not torch.cuda.is_available():\n",
    "    cuda_device = 'cpu'\n",
    "    print(\"Warning: No GPU found! Using the CPU instead. The iterations will be slow.\")\n",
    "    print(\"Perhaps CUDA/ROCm or the right pytorch version is not properly installed?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Various functions and classes\n",
    "def sinc(x):\n",
    "    return torch.where(x != 0, torch.sin(math.pi * x) / (math.pi * x), x.new_ones([]))\n",
    "\n",
    "\n",
    "def lanczos(x, a):\n",
    "    cond = torch.logical_and(-a < x, x < a)\n",
    "    out = torch.where(cond, sinc(x) * sinc(x/a), x.new_zeros([]))\n",
    "    return out / out.sum()\n",
    "\n",
    "\n",
    "def ramp(ratio, width):\n",
    "    n = math.ceil(width / ratio + 1)\n",
    "    out = torch.empty([n])\n",
    "    cur = 0\n",
    "    for i in range(out.shape[0]):\n",
    "        out[i] = cur\n",
    "        cur += ratio\n",
    "    return torch.cat([-out[1:].flip([0]), out])[1:-1]\n",
    "\n",
    "# Used in older MakeCutouts\n",
    "def resample(input, size, align_corners=True):\n",
    "    n, c, h, w = input.shape\n",
    "    dh, dw = size\n",
    "\n",
    "    input = input.view([n * c, 1, h, w])\n",
    "\n",
    "    if dh < h:\n",
    "        kernel_h = lanczos(ramp(dh / h, 2), 2).to(input.device, input.dtype)\n",
    "        pad_h = (kernel_h.shape[0] - 1) // 2\n",
    "        input = F.pad(input, (0, 0, pad_h, pad_h), 'reflect')\n",
    "        input = F.conv2d(input, kernel_h[None, None, :, None])\n",
    "\n",
    "    if dw < w:\n",
    "        kernel_w = lanczos(ramp(dw / w, 2), 2).to(input.device, input.dtype)\n",
    "        pad_w = (kernel_w.shape[0] - 1) // 2\n",
    "        input = F.pad(input, (pad_w, pad_w, 0, 0), 'reflect')\n",
    "        input = F.conv2d(input, kernel_w[None, None, None, :])\n",
    "\n",
    "    input = input.view([n, c, h, w])\n",
    "    return F.interpolate(input, size, mode='bicubic', align_corners=align_corners)\n",
    "\n",
    "\n",
    "class ReplaceGrad(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, x_forward, x_backward):\n",
    "        ctx.shape = x_backward.shape\n",
    "        return x_forward\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_in):\n",
    "        return None, grad_in.sum_to_size(ctx.shape)\n",
    "\n",
    "replace_grad = ReplaceGrad.apply\n",
    "\n",
    "\n",
    "class ClampWithGrad(torch.autograd.Function):\n",
    "    @staticmethod\n",
    "    def forward(ctx, input, min, max):\n",
    "        ctx.min = min\n",
    "        ctx.max = max\n",
    "        ctx.save_for_backward(input)\n",
    "        return input.clamp(min, max)\n",
    "\n",
    "    @staticmethod\n",
    "    def backward(ctx, grad_in):\n",
    "        input, = ctx.saved_tensors\n",
    "        return grad_in * (grad_in * (input - input.clamp(ctx.min, ctx.max)) >= 0), None, None\n",
    "\n",
    "clamp_with_grad = ClampWithGrad.apply\n",
    "\n",
    "\n",
    "def vector_quantize(x, codebook):\n",
    "    d = x.pow(2).sum(dim=-1, keepdim=True) + codebook.pow(2).sum(dim=1) - 2 * x @ codebook.T\n",
    "    indices = d.argmin(-1)\n",
    "    x_q = F.one_hot(indices, codebook.shape[0]).to(d.dtype) @ codebook\n",
    "    return replace_grad(x_q, x)\n",
    "\n",
    "\n",
    "class Prompt(nn.Module):\n",
    "    def __init__(self, embed, weight=1., stop=float('-inf')):\n",
    "        super().__init__()\n",
    "        self.register_buffer('embed', embed)\n",
    "        self.register_buffer('weight', torch.as_tensor(weight))\n",
    "        self.register_buffer('stop', torch.as_tensor(stop))\n",
    "\n",
    "    def forward(self, input):\n",
    "        input_normed = F.normalize(input.unsqueeze(1), dim=2)\n",
    "        embed_normed = F.normalize(self.embed.unsqueeze(0), dim=2)\n",
    "        dists = input_normed.sub(embed_normed).norm(dim=2).div(2).arcsin().pow(2).mul(2)\n",
    "        dists = dists * self.weight.sign()\n",
    "        return self.weight.abs() * replace_grad(dists, torch.maximum(dists, self.stop)).mean()\n",
    "\n",
    "\n",
    "#NR: Split prompts and weights\n",
    "def split_prompt(prompt):\n",
    "    vals = prompt.rsplit(':', 2)\n",
    "    vals = vals + ['', '1', '-inf'][len(vals):]\n",
    "    return vals[0], float(vals[1]), float(vals[2])\n",
    "\n",
    "\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 # not used with pooling\n",
    "        \n",
    "        # Pick your own augments & their order\n",
    "        augment_list = []\n",
    "        for item in augments[0]:\n",
    "            if item == 'Ji':\n",
    "                augment_list.append(K.ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1, p=0.7))\n",
    "            elif item == 'Sh':\n",
    "                augment_list.append(K.RandomSharpness(sharpness=0.3, p=0.5))\n",
    "            elif item == 'Gn':\n",
    "                augment_list.append(K.RandomGaussianNoise(mean=0.0, std=1., p=0.5))\n",
    "            elif item == 'Pe':\n",
    "                augment_list.append(K.RandomPerspective(distortion_scale=0.7, p=0.7))\n",
    "            elif item == 'Ro':\n",
    "                augment_list.append(K.RandomRotation(degrees=15, p=0.7))\n",
    "            elif item == 'Af':\n",
    "                augment_list.append(K.RandomAffine(degrees=15, translate=0.1, shear=5, p=0.7, padding_mode='zeros', keepdim=True)) # border, reflection, zeros\n",
    "            elif item == 'Et':\n",
    "                augment_list.append(K.RandomElasticTransform(p=0.7))\n",
    "            elif item == 'Ts':\n",
    "                augment_list.append(K.RandomThinPlateSpline(scale=0.8, same_on_batch=True, p=0.7))\n",
    "            elif item == 'Cr':\n",
    "                augment_list.append(K.RandomCrop(size=(self.cut_size,self.cut_size), pad_if_needed=True, padding_mode='reflect', p=0.5))\n",
    "            elif item == 'Er':\n",
    "                augment_list.append(K.RandomErasing(scale=(.1, .4), ratio=(.3, 1/.3), same_on_batch=True, p=0.7))\n",
    "            elif item == 'Re':\n",
    "                augment_list.append(K.RandomResizedCrop(size=(self.cut_size,self.cut_size), scale=(0.1,1),  ratio=(0.75,1.333), cropping_mode='resample', p=0.5))\n",
    "                \n",
    "        self.augs = nn.Sequential(*augment_list)\n",
    "        self.noise_fac = 0.1\n",
    "        # self.noise_fac = False\n",
    "\n",
    "        # Uncomment if you like seeing the list ;)\n",
    "        # print(augment_list)\n",
    "        \n",
    "        # Pooling\n",
    "        self.av_pool = nn.AdaptiveAvgPool2d((self.cut_size, self.cut_size))\n",
    "        self.max_pool = nn.AdaptiveMaxPool2d((self.cut_size, self.cut_size))\n",
    "\n",
    "    def forward(self, input):\n",
    "        cutouts = []\n",
    "        \n",
    "        for _ in range(self.cutn):            \n",
    "            # Use Pooling\n",
    "            cutout = (self.av_pool(input) + self.max_pool(input))/2\n",
    "            cutouts.append(cutout)\n",
    "            \n",
    "        batch = self.augs(torch.cat(cutouts, dim=0))\n",
    "        \n",
    "        if self.noise_fac:\n",
    "            facs = batch.new_empty([self.cutn, 1, 1, 1]).uniform_(0, self.noise_fac)\n",
    "            batch = batch + facs * torch.randn_like(batch)\n",
    "        return batch\n",
    "\n",
    "\n",
    "\n",
    "def load_vqgan_model(config_path, checkpoint_path):\n",
    "    global gumbel\n",
    "    gumbel = False\n",
    "    config = OmegaConf.load(config_path)\n",
    "    if config.model.target == 'taming.models.vqgan.VQModel':\n",
    "        model = vqgan.VQModel(**config.model.params)\n",
    "        model.eval().requires_grad_(False)\n",
    "        model.init_from_ckpt(checkpoint_path)\n",
    "    elif config.model.target == 'taming.models.vqgan.GumbelVQ':\n",
    "        model = vqgan.GumbelVQ(**config.model.params)\n",
    "        model.eval().requires_grad_(False)\n",
    "        model.init_from_ckpt(checkpoint_path)\n",
    "        gumbel = True\n",
    "    elif config.model.target == 'taming.models.cond_transformer.Net2NetTransformer':\n",
    "        parent_model = cond_transformer.Net2NetTransformer(**config.model.params)\n",
    "        parent_model.eval().requires_grad_(False)\n",
    "        parent_model.init_from_ckpt(checkpoint_path)\n",
    "        model = parent_model.first_stage_model\n",
    "    else:\n",
    "        raise ValueError(f'unknown model type: {config.model.target}')\n",
    "    del model.loss\n",
    "    return model\n",
    "\n",
    "\n",
    "def resize_image(image, out_size):\n",
    "    ratio = image.size[0] / image.size[1]\n",
    "    area = min(image.size[0] * image.size[1], out_size[0] * out_size[1])\n",
    "    size = round((area * ratio)**0.5), round((area / ratio)**0.5)\n",
    "    return image.resize(size, Image.LANCZOS)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Working with z of shape (1, 256, 16, 16) = 65536 dimensions.\n",
      "loaded pretrained LPIPS loss from taming/modules/autoencoder/lpips/vgg.pth\n",
      "VQLPIPSWithDiscriminator running with hinge loss.\n",
      "Restored from checkpoints/vqgan_imagenet_f16_16384.ckpt\n"
     ]
    }
   ],
   "source": [
    "# Do it\n",
    "device = torch.device(cuda_device)\n",
    "model = load_vqgan_model(vqgan_config, vqgan_checkpoint).to(device)\n",
    "jit = True if float(torch.__version__[:3]) < 1.8 else False\n",
    "perceptor = clip.load(clip_model, jit=jit)[0].eval().requires_grad_(False).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "cut_size = perceptor.visual.input_resolution\n",
    "f = 2**(model.decoder.num_resolutions - 1)\n",
    "\n",
    "make_cutouts = MakeCutouts(cut_size, cutn, cut_pow=cut_pow) \n",
    "\n",
    "toksX, toksY = size[0] // f, size[1] // f\n",
    "sideX, sideY = toksX * f, toksY * f\n",
    "\n",
    "# Gumbel or not?\n",
    "e_dim = model.quantize.e_dim\n",
    "n_toks = model.quantize.n_e\n",
    "z_min = model.quantize.embedding.weight.min(dim=0).values[None, :, None, None]\n",
    "z_max = model.quantize.embedding.weight.max(dim=0).values[None, :, None, None]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[[-5.0996e-05, -2.9291e-05,  3.8511e-05,  ...,  2.7785e-06,\n",
       "            9.0637e-01, -3.5288e-05],\n",
       "          [-1.7766e-05,  8.3750e-06,  4.0015e-05,  ..., -1.2230e-05,\n",
       "            1.9764e-05, -2.3202e-05],\n",
       "          [-5.8987e-05, -5.3469e-05,  3.0875e-05,  ..., -1.4589e-05,\n",
       "            2.1192e-05,  3.8347e-05],\n",
       "          ...,\n",
       "          [ 3.1803e-05,  4.0603e-05,  5.4634e-06,  ...,  8.0669e-05,\n",
       "           -1.6351e-05,  3.0171e-05],\n",
       "          [-3.7514e-05,  4.3857e-05,  2.1694e-05,  ..., -5.6673e-05,\n",
       "           -5.6358e-05, -3.4189e-06],\n",
       "          [ 9.3955e-06, -2.3645e-01,  4.3247e-05,  ..., -4.8752e-05,\n",
       "            3.2483e-05,  2.9966e-05]],\n",
       "\n",
       "         [[ 1.1908e-05,  5.4353e-05, -8.9054e-06,  ...,  5.3405e-05,\n",
       "           -1.4026e+00, -8.2203e-06],\n",
       "          [-4.1242e-05,  4.8444e-05,  4.1693e-05,  ...,  3.6219e-06,\n",
       "           -2.8750e-05, -1.8468e-05],\n",
       "          [ 3.7632e-05,  5.3459e-05, -5.9309e-05,  ..., -5.2000e-06,\n",
       "            8.9573e-06, -4.7822e-05],\n",
       "          ...,\n",
       "          [-2.6927e-05, -3.5945e-05, -5.5189e-05,  ..., -1.5114e-04,\n",
       "            4.5025e-05,  2.2841e-05],\n",
       "          [ 4.5073e-05,  3.2473e-05,  4.5534e-05,  ..., -3.6037e-05,\n",
       "            1.0783e-05,  5.7692e-05],\n",
       "          [ 1.9665e-05,  1.2950e+00, -3.8374e-05,  ..., -4.6952e-05,\n",
       "           -6.0295e-05, -4.9865e-05]],\n",
       "\n",
       "         [[ 2.1975e-05, -5.7940e-05,  4.9627e-05,  ...,  1.7473e-05,\n",
       "           -1.1939e+00,  1.5664e-06],\n",
       "          [-2.3618e-05,  2.7310e-05, -5.3913e-05,  ..., -4.0666e-05,\n",
       "           -4.4916e-05,  5.8760e-06],\n",
       "          [ 3.8953e-05,  5.8530e-06, -2.6140e-05,  ..., -3.5582e-05,\n",
       "           -1.1016e-05,  3.0788e-05],\n",
       "          ...,\n",
       "          [-5.4711e-05, -2.7799e-05, -3.3730e-05,  ...,  7.8487e-05,\n",
       "            3.3357e-05, -1.9324e-05],\n",
       "          [-5.4624e-05, -3.6238e-05, -5.6496e-05,  ..., -1.0613e-06,\n",
       "            4.6534e-06,  4.9167e-05],\n",
       "          [-5.3570e-05,  1.1726e+00,  4.1411e-05,  ..., -1.2569e-05,\n",
       "           -7.2079e-06,  1.1028e-06]],\n",
       "\n",
       "         ...,\n",
       "\n",
       "         [[ 5.0567e-05, -5.7720e-05,  3.7970e-05,  ..., -3.7104e-05,\n",
       "           -4.5995e-01, -3.7684e-05],\n",
       "          [ 5.9341e-05, -9.6673e-06,  1.6685e-06,  ...,  5.1958e-06,\n",
       "           -2.5065e-06,  3.1783e-05],\n",
       "          [ 6.7426e-06,  5.7238e-05, -5.5164e-05,  ...,  4.2481e-05,\n",
       "            4.1842e-05, -3.0508e-05],\n",
       "          ...,\n",
       "          [ 2.4290e-05,  3.5704e-05,  4.3226e-05,  ...,  1.6763e-04,\n",
       "            8.3660e-06,  4.5990e-05],\n",
       "          [-1.3562e-05,  5.9056e-05,  2.5376e-05,  ..., -3.0215e-05,\n",
       "           -3.1404e-05, -1.4916e-05],\n",
       "          [-2.8042e-05,  7.6410e-01,  6.9015e-06,  ..., -3.8785e-05,\n",
       "            5.7003e-05, -1.0879e-06]],\n",
       "\n",
       "         [[-2.8585e-05,  5.5114e-05, -5.1230e-05,  ...,  3.2287e-05,\n",
       "           -1.0514e+00,  2.5540e-05],\n",
       "          [-7.6186e-06, -4.7043e-05, -4.7699e-05,  ..., -5.0602e-05,\n",
       "           -1.8998e-05, -2.9058e-05],\n",
       "          [ 3.2962e-05, -2.8809e-05,  2.2126e-05,  ...,  3.1326e-06,\n",
       "            3.6045e-05, -5.4259e-05],\n",
       "          ...,\n",
       "          [ 5.0576e-05,  8.4278e-06, -3.0068e-05,  ...,  1.1636e-04,\n",
       "           -5.7496e-05,  5.3331e-07],\n",
       "          [-2.5560e-05, -5.8547e-05, -4.2752e-05,  ..., -8.6749e-06,\n",
       "           -3.0010e-05, -3.6472e-06],\n",
       "          [-3.6751e-06,  9.8818e-01,  4.1716e-05,  ..., -4.0955e-05,\n",
       "           -5.0036e-05, -5.2589e-05]],\n",
       "\n",
       "         [[ 1.6364e-05, -1.8956e-05, -3.0594e-05,  ..., -2.8950e-05,\n",
       "            3.1147e+00, -4.7086e-06],\n",
       "          [ 4.3059e-05,  1.0571e-06, -3.7838e-05,  ..., -6.6329e-06,\n",
       "           -2.7313e-06, -2.2532e-05],\n",
       "          [ 4.7620e-05,  3.5929e-05, -5.2193e-05,  ..., -2.6612e-05,\n",
       "            5.9918e-05, -1.3491e-05],\n",
       "          ...,\n",
       "          [-2.3156e-05, -4.1573e-05, -5.2237e-05,  ..., -4.1410e-05,\n",
       "            5.1623e-05,  4.9292e-05],\n",
       "          [ 2.3111e-05,  3.7747e-05,  3.0495e-05,  ..., -1.7259e-05,\n",
       "           -9.1463e-06, -4.2736e-05],\n",
       "          [ 1.4226e-05, -2.0345e+00,  5.0168e-06,  ..., -3.6918e-05,\n",
       "           -2.6196e-05,  5.3498e-06]]]], device='cuda:0', requires_grad=True)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "one_hot = F.one_hot(torch.randint(n_toks, [toksY * toksX], device=device), n_toks).float()\n",
    "z = one_hot @ model.quantize.embedding.weight\n",
    "\n",
    "z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2) \n",
    "    #z = torch.rand_like(z)*2\t\t\t\t\t\t# NR: check\n",
    "\n",
    "z_orig = z.clone()\n",
    "z.requires_grad_(True)#Init Z\n",
    "one_hot = F.one_hot(torch.randint(n_toks, [toksY * toksX], device=device), n_toks).float()\n",
    "\n",
    "z = one_hot @ model.quantize.embedding.weight\n",
    "z = z.view([-1, toksY, toksX, e_dim]).permute(0, 3, 1, 2) \n",
    "\n",
    "z_orig = z.clone()\n",
    "z.requires_grad_(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "pMs = []\n",
    "normalize = transforms.Normalize(mean=[0.48145466, 0.4578275, 0.40821073],\n",
    "                                  std=[0.26862954, 0.26130258, 0.27577711])\n",
    "\n",
    "# From imagenet - Which is better?\n",
    "#normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "#                                  std=[0.229, 0.224, 0.225])\n",
    "\n",
    "# CLIP tokenize/encode   \n",
    "if prompts != None:\n",
    "    for prompt in prompts:\n",
    "        txt, weight, stop = split_prompt(prompt)\n",
    "        embed = perceptor.encode_text(clip.tokenize(txt).to(device)).float()\n",
    "        pMs.append(Prompt(embed, weight, stop).to(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt = optim.Adam([z], lr= step_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using device: cuda:0\n",
      "Optimising using: Adam (\n",
      "Parameter Group 0\n",
      "    amsgrad: False\n",
      "    betas: (0.9, 0.999)\n",
      "    eps: 1e-08\n",
      "    lr: 0.1\n",
      "    weight_decay: 0\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "# Output for the user\n",
    "print('Using device:', device)\n",
    "print('Optimising using:', opt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using seed: 11014765562215412717\n"
     ]
    }
   ],
   "source": [
    "if seed is None:\n",
    "    seed = torch.seed()\n",
    "else:\n",
    "    seed = seed  \n",
    "torch.manual_seed(seed)\n",
    "print('Using seed:', seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Vector quantize\n",
    "def synth(z):\n",
    "    if gumbel:\n",
    "        z_q = vector_quantize(z.movedim(1, 3), model.quantize.embed.weight).movedim(3, 1)\n",
    "    else:\n",
    "        z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)\n",
    "    return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)\n",
    "\n",
    "\n",
    "#@torch.no_grad()\n",
    "@torch.inference_mode()\n",
    "def checkin(i, losses):\n",
    "    losses_str = ', '.join(f'{loss.item():g}' for loss in losses)\n",
    "    tqdm.write(f'i: {i}, loss: {sum(losses).item():g}, losses: {losses_str}')\n",
    "    out = synth(z)\n",
    "    info = PngImagePlugin.PngInfo()\n",
    "    TF.to_pil_image(out[0].cpu()).save(output, pnginfo=info) \n",
    "\n",
    "\n",
    "def ascend_txt():\n",
    "    global i\n",
    "    out = synth(z)\n",
    "    iii = perceptor.encode_image(normalize(make_cutouts(out))).float()\n",
    "    \n",
    "    result = []\n",
    "    \n",
    "    for prompt in pMs:\n",
    "        result.append(prompt(iii))\n",
    "    return result # return loss\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(i):\n",
    "    opt.zero_grad(set_to_none=True)\n",
    "    lossAll = ascend_txt()\n",
    "    \n",
    "    if i % display_freq == 0:\n",
    "        checkin(i, lossAll)\n",
    "       \n",
    "    loss = sum(lossAll)\n",
    "    loss.backward()\n",
    "    opt.step()\n",
    "    \n",
    "    #with torch.no_grad():\n",
    "    with torch.inference_mode():\n",
    "        z.copy_(z.maximum(z_min).minimum(z_max))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "0it [00:00, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 0, loss: 0.709187, losses: 0.709187\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "51it [00:07,  6.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 50, loss: 0.723484, losses: 0.723484\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "101it [00:14,  6.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 100, loss: 0.705107, losses: 0.705107\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "151it [00:21,  6.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 150, loss: 0.782248, losses: 0.782248\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "201it [00:27,  6.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 200, loss: 0.778452, losses: 0.778452\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "251it [00:34,  6.17it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 250, loss: 0.704561, losses: 0.704561\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "301it [00:41,  6.20it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 300, loss: 0.733989, losses: 0.733989\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "351it [00:48,  6.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 350, loss: 0.791059, losses: 0.791059\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "401it [00:55,  6.13it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 400, loss: 0.686357, losses: 0.686357\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "451it [01:02,  6.04it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 450, loss: 0.761763, losses: 0.761763\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "500it [01:09,  7.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "i: 500, loss: 0.732906, losses: 0.732906\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "i = 0 # Iteration counter\n",
    "\n",
    "# Do it\n",
    "try:\n",
    "    with tqdm() as pbar:\n",
    "        while True:            \n",
    "           \n",
    "            train(i)\n",
    "            \n",
    "            # Ready to stop yet?\n",
    "            if i == max_iterations:\n",
    "                break\n",
    "                \n",
    "            i += 1\n",
    "            pbar.update()\n",
    "except KeyboardInterrupt:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[0.4809, 0.4662, 0.4732,  ..., 0.4739, 0.4643, 0.4305],\n",
      "          [0.4786, 0.4794, 0.4599,  ..., 0.4594, 0.4569, 0.4655],\n",
      "          [0.4772, 0.4740, 0.4666,  ..., 0.4532, 0.4464, 0.4363],\n",
      "          ...,\n",
      "          [0.4734, 0.4593, 0.4507,  ..., 0.4788, 0.4685, 0.4824],\n",
      "          [0.5044, 0.4770, 0.4625,  ..., 0.4776, 0.4695, 0.4834],\n",
      "          [0.4667, 0.4809, 0.4625,  ..., 0.4916, 0.4818, 0.4698]],\n",
      "\n",
      "         [[0.4151, 0.3737, 0.3839,  ..., 0.3466, 0.3557, 0.3049],\n",
      "          [0.4018, 0.3868, 0.3937,  ..., 0.3576, 0.3685, 0.3652],\n",
      "          [0.4033, 0.3896, 0.3964,  ..., 0.3545, 0.3631, 0.3520],\n",
      "          ...,\n",
      "          [0.4008, 0.4060, 0.4050,  ..., 0.4001, 0.3997, 0.3875],\n",
      "          [0.3797, 0.3889, 0.3900,  ..., 0.3941, 0.3914, 0.3968],\n",
      "          [0.3749, 0.3811, 0.3819,  ..., 0.3816, 0.3885, 0.3757]],\n",
      "\n",
      "         [[0.4368, 0.4145, 0.4331,  ..., 0.3667, 0.3751, 0.3416],\n",
      "          [0.4187, 0.4316, 0.4390,  ..., 0.3703, 0.3825, 0.3876],\n",
      "          [0.4301, 0.4235, 0.4302,  ..., 0.3790, 0.3916, 0.3761],\n",
      "          ...,\n",
      "          [0.4507, 0.4315, 0.4297,  ..., 0.4327, 0.4297, 0.4569],\n",
      "          [0.4773, 0.4498, 0.4582,  ..., 0.4383, 0.4417, 0.4633],\n",
      "          [0.4329, 0.4422, 0.4407,  ..., 0.4617, 0.4541, 0.4449]]]],\n",
      "       device='cuda:0', grad_fn=<ClampWithGradBackward>)\n"
     ]
    }
   ],
   "source": [
    "out = synth(z)\n",
    "iii = perceptor.encode_image(normalize(make_cutouts(out))).float()\n",
    "\n",
    "print(out)\n",
    "\n",
    "result = []\n",
    "\n",
    "for prompt in pMs:\n",
    "    result.append(prompt(iii))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[-0.0074, -0.2058,  0.4287,  ...,  0.3250,  0.1367, -0.2632],\n",
      "        [ 0.1456, -0.2710,  0.3555,  ...,  0.2939, -0.1127,  0.0892],\n",
      "        [ 0.1191,  0.1833,  0.4578,  ...,  0.1143,  0.1541, -0.1753],\n",
      "        ...,\n",
      "        [ 0.0941, -0.1650,  0.1428,  ...,  0.1637, -0.2142, -0.1042],\n",
      "        [ 0.0948, -0.1512,  0.3711,  ...,  0.5181, -0.0595, -0.0899],\n",
      "        [ 0.0200, -0.0715,  0.3022,  ...,  0.5234,  0.1709, -0.2812]],\n",
      "       device='cuda:0', grad_fn=<CopyBackwards>)\n"
     ]
    }
   ],
   "source": [
    "print(iii)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "def synth(z):\n",
    "    if gumbel:\n",
    "        z_q = vector_quantize(z.movedim(1, 3), model.quantize.embed.weight).movedim(3, 1)\n",
    "    else:\n",
    "        z_q = vector_quantize(z.movedim(1, 3), model.quantize.embedding.weight).movedim(3, 1)\n",
    "    print(z_q)\n",
    "    return clamp_with_grad(model.decode(z_q).add(1).div(2), 0, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[-5.0996e-05, -2.9291e-05,  3.8511e-05,  ...,  2.7785e-06,\n",
      "            9.0637e-01, -3.5288e-05],\n",
      "          [-1.7766e-05,  8.3750e-06,  4.0015e-05,  ..., -1.2230e-05,\n",
      "            1.9764e-05, -2.3202e-05],\n",
      "          [-5.8987e-05, -5.3469e-05,  3.0875e-05,  ..., -1.4589e-05,\n",
      "            2.1192e-05,  3.8347e-05],\n",
      "          ...,\n",
      "          [ 3.1803e-05,  4.0603e-05,  5.4634e-06,  ...,  8.0669e-05,\n",
      "           -1.6351e-05,  3.0171e-05],\n",
      "          [-3.7514e-05,  4.3857e-05,  2.1694e-05,  ..., -5.6673e-05,\n",
      "           -5.6358e-05, -3.4189e-06],\n",
      "          [ 9.3955e-06, -2.3645e-01,  4.3247e-05,  ..., -4.8752e-05,\n",
      "            3.2483e-05,  2.9966e-05]],\n",
      "\n",
      "         [[ 1.1908e-05,  5.4353e-05, -8.9054e-06,  ...,  5.3405e-05,\n",
      "           -1.4026e+00, -8.2203e-06],\n",
      "          [-4.1242e-05,  4.8444e-05,  4.1693e-05,  ...,  3.6219e-06,\n",
      "           -2.8750e-05, -1.8468e-05],\n",
      "          [ 3.7632e-05,  5.3459e-05, -5.9309e-05,  ..., -5.2000e-06,\n",
      "            8.9573e-06, -4.7822e-05],\n",
      "          ...,\n",
      "          [-2.6927e-05, -3.5945e-05, -5.5189e-05,  ..., -1.5114e-04,\n",
      "            4.5025e-05,  2.2841e-05],\n",
      "          [ 4.5073e-05,  3.2473e-05,  4.5534e-05,  ..., -3.6037e-05,\n",
      "            1.0783e-05,  5.7692e-05],\n",
      "          [ 1.9665e-05,  1.2950e+00, -3.8374e-05,  ..., -4.6952e-05,\n",
      "           -6.0295e-05, -4.9865e-05]],\n",
      "\n",
      "         [[ 2.1975e-05, -5.7940e-05,  4.9627e-05,  ...,  1.7473e-05,\n",
      "           -1.1939e+00,  1.5664e-06],\n",
      "          [-2.3618e-05,  2.7310e-05, -5.3913e-05,  ..., -4.0666e-05,\n",
      "           -4.4916e-05,  5.8760e-06],\n",
      "          [ 3.8953e-05,  5.8530e-06, -2.6140e-05,  ..., -3.5582e-05,\n",
      "           -1.1016e-05,  3.0788e-05],\n",
      "          ...,\n",
      "          [-5.4711e-05, -2.7799e-05, -3.3730e-05,  ...,  7.8487e-05,\n",
      "            3.3357e-05, -1.9324e-05],\n",
      "          [-5.4624e-05, -3.6238e-05, -5.6496e-05,  ..., -1.0613e-06,\n",
      "            4.6534e-06,  4.9167e-05],\n",
      "          [-5.3570e-05,  1.1726e+00,  4.1411e-05,  ..., -1.2569e-05,\n",
      "           -7.2079e-06,  1.1028e-06]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 5.0567e-05, -5.7720e-05,  3.7970e-05,  ..., -3.7104e-05,\n",
      "           -4.5995e-01, -3.7684e-05],\n",
      "          [ 5.9341e-05, -9.6673e-06,  1.6685e-06,  ...,  5.1958e-06,\n",
      "           -2.5065e-06,  3.1783e-05],\n",
      "          [ 6.7426e-06,  5.7238e-05, -5.5164e-05,  ...,  4.2481e-05,\n",
      "            4.1842e-05, -3.0508e-05],\n",
      "          ...,\n",
      "          [ 2.4290e-05,  3.5704e-05,  4.3226e-05,  ...,  1.6763e-04,\n",
      "            8.3660e-06,  4.5990e-05],\n",
      "          [-1.3562e-05,  5.9056e-05,  2.5376e-05,  ..., -3.0215e-05,\n",
      "           -3.1404e-05, -1.4916e-05],\n",
      "          [-2.8042e-05,  7.6410e-01,  6.9015e-06,  ..., -3.8785e-05,\n",
      "            5.7003e-05, -1.0879e-06]],\n",
      "\n",
      "         [[-2.8585e-05,  5.5114e-05, -5.1230e-05,  ...,  3.2287e-05,\n",
      "           -1.0514e+00,  2.5540e-05],\n",
      "          [-7.6186e-06, -4.7043e-05, -4.7699e-05,  ..., -5.0602e-05,\n",
      "           -1.8998e-05, -2.9058e-05],\n",
      "          [ 3.2962e-05, -2.8809e-05,  2.2126e-05,  ...,  3.1326e-06,\n",
      "            3.6045e-05, -5.4259e-05],\n",
      "          ...,\n",
      "          [ 5.0576e-05,  8.4278e-06, -3.0068e-05,  ...,  1.1636e-04,\n",
      "           -5.7496e-05,  5.3331e-07],\n",
      "          [-2.5560e-05, -5.8547e-05, -4.2752e-05,  ..., -8.6749e-06,\n",
      "           -3.0010e-05, -3.6472e-06],\n",
      "          [-3.6751e-06,  9.8818e-01,  4.1716e-05,  ..., -4.0955e-05,\n",
      "           -5.0036e-05, -5.2589e-05]],\n",
      "\n",
      "         [[ 1.6364e-05, -1.8956e-05, -3.0594e-05,  ..., -2.8950e-05,\n",
      "            3.1147e+00, -4.7086e-06],\n",
      "          [ 4.3059e-05,  1.0571e-06, -3.7838e-05,  ..., -6.6329e-06,\n",
      "           -2.7313e-06, -2.2532e-05],\n",
      "          [ 4.7620e-05,  3.5929e-05, -5.2193e-05,  ..., -2.6612e-05,\n",
      "            5.9918e-05, -1.3491e-05],\n",
      "          ...,\n",
      "          [-2.3156e-05, -4.1573e-05, -5.2237e-05,  ..., -4.1410e-05,\n",
      "            5.1623e-05,  4.9292e-05],\n",
      "          [ 2.3111e-05,  3.7747e-05,  3.0495e-05,  ..., -1.7259e-05,\n",
      "           -9.1463e-06, -4.2736e-05],\n",
      "          [ 1.4226e-05, -2.0345e+00,  5.0168e-06,  ..., -3.6918e-05,\n",
      "           -2.6196e-05,  5.3498e-06]]]], device='cuda:0',\n",
      "       grad_fn=<PermuteBackward>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[[0.4809, 0.4662, 0.4732,  ..., 0.4739, 0.4643, 0.4305],\n",
       "          [0.4786, 0.4794, 0.4599,  ..., 0.4594, 0.4569, 0.4655],\n",
       "          [0.4772, 0.4740, 0.4666,  ..., 0.4532, 0.4464, 0.4363],\n",
       "          ...,\n",
       "          [0.4734, 0.4593, 0.4507,  ..., 0.4788, 0.4685, 0.4824],\n",
       "          [0.5044, 0.4770, 0.4625,  ..., 0.4776, 0.4695, 0.4834],\n",
       "          [0.4667, 0.4809, 0.4625,  ..., 0.4916, 0.4818, 0.4698]],\n",
       "\n",
       "         [[0.4151, 0.3737, 0.3839,  ..., 0.3466, 0.3557, 0.3049],\n",
       "          [0.4018, 0.3868, 0.3937,  ..., 0.3576, 0.3685, 0.3652],\n",
       "          [0.4033, 0.3896, 0.3964,  ..., 0.3545, 0.3631, 0.3520],\n",
       "          ...,\n",
       "          [0.4008, 0.4060, 0.4050,  ..., 0.4001, 0.3997, 0.3875],\n",
       "          [0.3797, 0.3889, 0.3900,  ..., 0.3941, 0.3914, 0.3968],\n",
       "          [0.3749, 0.3811, 0.3819,  ..., 0.3816, 0.3885, 0.3757]],\n",
       "\n",
       "         [[0.4368, 0.4145, 0.4331,  ..., 0.3667, 0.3751, 0.3416],\n",
       "          [0.4187, 0.4316, 0.4390,  ..., 0.3703, 0.3825, 0.3876],\n",
       "          [0.4301, 0.4235, 0.4302,  ..., 0.3790, 0.3916, 0.3761],\n",
       "          ...,\n",
       "          [0.4507, 0.4315, 0.4297,  ..., 0.4327, 0.4297, 0.4569],\n",
       "          [0.4773, 0.4498, 0.4582,  ..., 0.4383, 0.4417, 0.4633],\n",
       "          [0.4329, 0.4422, 0.4407,  ..., 0.4617, 0.4541, 0.4449]]]],\n",
       "       device='cuda:0', grad_fn=<ClampWithGradBackward>)"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "synth(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[-5.0996e-05, -2.9291e-05,  3.8511e-05,  ...,  2.7785e-06,\n",
      "            9.0637e-01, -3.5288e-05],\n",
      "          [-1.7766e-05,  8.3750e-06,  4.0015e-05,  ..., -1.2230e-05,\n",
      "            1.9764e-05, -2.3202e-05],\n",
      "          [-5.8987e-05, -5.3469e-05,  3.0875e-05,  ..., -1.4589e-05,\n",
      "            2.1192e-05,  3.8347e-05],\n",
      "          ...,\n",
      "          [ 3.1803e-05,  4.0603e-05,  5.4634e-06,  ...,  8.0669e-05,\n",
      "           -1.6351e-05,  3.0171e-05],\n",
      "          [-3.7514e-05,  4.3857e-05,  2.1694e-05,  ..., -5.6673e-05,\n",
      "           -5.6358e-05, -3.4189e-06],\n",
      "          [ 9.3955e-06, -2.3645e-01,  4.3247e-05,  ..., -4.8752e-05,\n",
      "            3.2483e-05,  2.9966e-05]],\n",
      "\n",
      "         [[ 1.1908e-05,  5.4353e-05, -8.9054e-06,  ...,  5.3405e-05,\n",
      "           -1.4026e+00, -8.2203e-06],\n",
      "          [-4.1242e-05,  4.8444e-05,  4.1693e-05,  ...,  3.6219e-06,\n",
      "           -2.8750e-05, -1.8468e-05],\n",
      "          [ 3.7632e-05,  5.3459e-05, -5.9309e-05,  ..., -5.2000e-06,\n",
      "            8.9573e-06, -4.7822e-05],\n",
      "          ...,\n",
      "          [-2.6927e-05, -3.5945e-05, -5.5189e-05,  ..., -1.5114e-04,\n",
      "            4.5025e-05,  2.2841e-05],\n",
      "          [ 4.5073e-05,  3.2473e-05,  4.5534e-05,  ..., -3.6037e-05,\n",
      "            1.0783e-05,  5.7692e-05],\n",
      "          [ 1.9665e-05,  1.2950e+00, -3.8374e-05,  ..., -4.6952e-05,\n",
      "           -6.0295e-05, -4.9865e-05]],\n",
      "\n",
      "         [[ 2.1975e-05, -5.7940e-05,  4.9627e-05,  ...,  1.7473e-05,\n",
      "           -1.1939e+00,  1.5664e-06],\n",
      "          [-2.3618e-05,  2.7310e-05, -5.3913e-05,  ..., -4.0666e-05,\n",
      "           -4.4916e-05,  5.8760e-06],\n",
      "          [ 3.8953e-05,  5.8530e-06, -2.6140e-05,  ..., -3.5582e-05,\n",
      "           -1.1016e-05,  3.0788e-05],\n",
      "          ...,\n",
      "          [-5.4711e-05, -2.7799e-05, -3.3730e-05,  ...,  7.8487e-05,\n",
      "            3.3357e-05, -1.9324e-05],\n",
      "          [-5.4624e-05, -3.6238e-05, -5.6496e-05,  ..., -1.0613e-06,\n",
      "            4.6534e-06,  4.9167e-05],\n",
      "          [-5.3570e-05,  1.1726e+00,  4.1411e-05,  ..., -1.2569e-05,\n",
      "           -7.2079e-06,  1.1028e-06]],\n",
      "\n",
      "         ...,\n",
      "\n",
      "         [[ 5.0567e-05, -5.7720e-05,  3.7970e-05,  ..., -3.7104e-05,\n",
      "           -4.5995e-01, -3.7684e-05],\n",
      "          [ 5.9341e-05, -9.6673e-06,  1.6685e-06,  ...,  5.1958e-06,\n",
      "           -2.5065e-06,  3.1783e-05],\n",
      "          [ 6.7426e-06,  5.7238e-05, -5.5164e-05,  ...,  4.2481e-05,\n",
      "            4.1842e-05, -3.0508e-05],\n",
      "          ...,\n",
      "          [ 2.4290e-05,  3.5704e-05,  4.3226e-05,  ...,  1.6763e-04,\n",
      "            8.3660e-06,  4.5990e-05],\n",
      "          [-1.3562e-05,  5.9056e-05,  2.5376e-05,  ..., -3.0215e-05,\n",
      "           -3.1404e-05, -1.4916e-05],\n",
      "          [-2.8042e-05,  7.6410e-01,  6.9015e-06,  ..., -3.8785e-05,\n",
      "            5.7003e-05, -1.0879e-06]],\n",
      "\n",
      "         [[-2.8585e-05,  5.5114e-05, -5.1230e-05,  ...,  3.2287e-05,\n",
      "           -1.0514e+00,  2.5540e-05],\n",
      "          [-7.6186e-06, -4.7043e-05, -4.7699e-05,  ..., -5.0602e-05,\n",
      "           -1.8998e-05, -2.9058e-05],\n",
      "          [ 3.2962e-05, -2.8809e-05,  2.2126e-05,  ...,  3.1326e-06,\n",
      "            3.6045e-05, -5.4259e-05],\n",
      "          ...,\n",
      "          [ 5.0576e-05,  8.4278e-06, -3.0068e-05,  ...,  1.1636e-04,\n",
      "           -5.7496e-05,  5.3331e-07],\n",
      "          [-2.5560e-05, -5.8547e-05, -4.2752e-05,  ..., -8.6749e-06,\n",
      "           -3.0010e-05, -3.6472e-06],\n",
      "          [-3.6751e-06,  9.8818e-01,  4.1716e-05,  ..., -4.0955e-05,\n",
      "           -5.0036e-05, -5.2589e-05]],\n",
      "\n",
      "         [[ 1.6364e-05, -1.8956e-05, -3.0594e-05,  ..., -2.8950e-05,\n",
      "            3.1147e+00, -4.7086e-06],\n",
      "          [ 4.3059e-05,  1.0571e-06, -3.7838e-05,  ..., -6.6329e-06,\n",
      "           -2.7313e-06, -2.2532e-05],\n",
      "          [ 4.7620e-05,  3.5929e-05, -5.2193e-05,  ..., -2.6612e-05,\n",
      "            5.9918e-05, -1.3491e-05],\n",
      "          ...,\n",
      "          [-2.3156e-05, -4.1573e-05, -5.2237e-05,  ..., -4.1410e-05,\n",
      "            5.1623e-05,  4.9292e-05],\n",
      "          [ 2.3111e-05,  3.7747e-05,  3.0495e-05,  ..., -1.7259e-05,\n",
      "           -9.1463e-06, -4.2736e-05],\n",
      "          [ 1.4226e-05, -2.0345e+00,  5.0168e-06,  ..., -3.6918e-05,\n",
      "           -2.6196e-05,  5.3498e-06]]]], device='cuda:0', requires_grad=True)\n"
     ]
    }
   ],
   "source": [
    "print(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "audio-clip",
   "language": "python",
   "name": "audio-clip"
  },
  "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}